from typing import TypedDict, AnyStr

from .prompt import highlight_explain_chain


class State(TypedDict):
    domain: AnyStr
    highlight_terms: AnyStr
    before_highlight_paragraph: AnyStr
    after_highlight_paragraph: AnyStr
    question: AnyStr
    explanation: AnyStr
    language: AnyStr


async def highlight_explain(state: State):
    adjacent_paragraphs = (
        state["before_highlight_paragraph"]
        + "**"
        + state["highlight_terms"]
        + "**"
        + state["after_highlight_paragraph"]
    )
    response = await highlight_explain_chain.ainvoke(
        {
            "domain": state["domain"],
            "highlight_terms": state["highlight_terms"],
            "adjacent_paragraphs": adjacent_paragraphs,
            "question": state["question"],
            "language": state["language"],
        }
    )
    return {"explanation": response["explanation"]}