from __future__ import annotations

from typing import Generator, List, Tuple

from pydantic import BaseModel

from srs2std.evaluate.decontextualize_strategies import DecontextualizeStrategy
from srs2std.srs.instance import SRS, Section


class Decontextualization(BaseModel):
    before: str
    after: str


class DecontextualizeContext:
    def __init__(self, srs_title: str, srs: SRS):
        self.srs_title = srs_title
        self.srs = srs

    def run(
        self,
        decontextualize_strategy: DecontextualizeStrategy,
    ) -> List[Decontextualization]:
        decontextualized_list = []
        for sentence, context in self.generate_examples():
            decontextualized = decontextualize_strategy.run(sentence, context)
            decontextualized = Decontextualization(
                before=sentence,
                after=decontextualized,
            )
            decontextualized_list.append(decontextualized)
        return decontextualized_list

    def generate_examples(self) -> Generator[Tuple[str, str], None, None]:
        for section in self.srs.sections:
            context = construct_context(self.srs_title, section)
            if not context:
                continue
            for sentence in section.to_sentences():
                yield (sentence, context)


def construct_context(srs_title: str, section: Section) -> str:
    section_content = section.to_str(ignore_title=False, ignore_headers=False)
    return "\n".join([srs_title, section_content]).strip()
