# TODO: Ask the original authors to clarify or include an explicit license.
# pending at https://github.com/anisha2102/molecular_facts/issues/1


from typing import Dict, override

from pydantic import BaseModel

from srs2std.evaluate.decontextualize_strategies.molecular_prompts import (
    MOLECULAR_STAGE1_PROMPT,
    MOLECULAR_STAGE2_PROMPT,
)
from srs2std.evaluate.decontextualize_strategies.strategy import DecontextualizeStrategy
from srs2std.llms import LLM, FinishReason, Prompt


class DisambiguitySchema(BaseModel):
    explanation: str | None
    subject: str | None
    disambiguation_criteria: str | None


class DecontextualizeSchema(BaseModel):
    explanation: str | None
    decontextualized_claim: str | None


class Molecular(DecontextualizeStrategy):
    def __init__(self, llm: LLM):
        self.llm = llm
        self.global_disambig_info = {}

    @override
    def run(self, sentence: str, context: str) -> str:
        disambig_info = self._identify_ambiguity(sentence)
        return self._decontextualize_ambiguity(sentence, context, disambig_info)

    def _identify_ambiguity(self, claim: str) -> Dict:
        prompt = Prompt(
            system="",
            user=MOLECULAR_STAGE1_PROMPT.replace("[claim]", claim),
        )
        result = self.llm.ask(prompt, DisambiguitySchema)
        if result.finish_reason != FinishReason.STOP or result.content is None:
            return {"subject": None, "disambiguation_criteria": None}

        subject = result.content.subject
        criteria = result.content.disambiguation_criteria

        if subject and subject != "None":
            if subject not in self.global_disambig_info:
                self.global_disambig_info[subject] = None
            if criteria and criteria != "None":
                self.global_disambig_info[subject] = criteria

        return {"subject": subject, "disambiguation_criteria": criteria}

    def _decontextualize_ambiguity(
        self, claim: str, context: str, disambig: Dict
    ) -> str:
        subject = disambig["subject"]
        if subject not in self.global_disambig_info:
            self.global_disambig_info[subject] = None

        disambig_new = {
            "subject": subject,
            "disambiguation_criteria": self.global_disambig_info[subject],
        }
        context = " ".join(context.split())

        prompt = Prompt(
            system="",
            user=MOLECULAR_STAGE2_PROMPT.replace("[claim]", claim)
            .replace("[disambiguation]", str(disambig_new))
            .replace("[context]", context),
        )

        result = self.llm.ask(prompt, DecontextualizeSchema)
        if result.finish_reason != FinishReason.STOP or result.content is None:
            return claim

        return result.content.decontextualized_claim or claim
