from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from src.agents.CallState import CallState
from src.agents.schemas import CallSummary

class SummarizationAgent:
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        
    def __call__(self, state: CallState) -> CallState:
        """Generates summaries and key points."""
        clean_text = state.get("clean_content", state.get("content", ""))
        if not clean_text:
            return state
            
        prompt = PromptTemplate.from_template(
            "Summarize the following call transcript and extract key points.\n"
            "Return a concise summary, a short list of key points, 2-8 action items, 3-8 topic tags, and 2-6 highlights.\n"
            "Action items must be concrete follow-ups; include the owner (Agent/Customer) when you can.\n\n"
            "Transcript:\n{text}"
        )
        
        structured_llm = self._structured_llm()
        chain = prompt | structured_llm
        result = chain.invoke({"text": clean_text})

        if isinstance(result, CallSummary):
            state["summary"] = result.summary
            state["key_points"] = result.key_points
            state["action_items"] = result.action_items
            state["tags"] = result.tags
            state["highlights"] = result.highlights
        else:
            state["summary"] = (result or {}).get("summary", "")
            state["key_points"] = (result or {}).get("key_points", [])
            state["action_items"] = (result or {}).get("action_items", [])
            state["tags"] = (result or {}).get("tags", [])
            state["highlights"] = (result or {}).get("highlights", [])
            
        return state

    def _structured_llm(self):
        # Some LangChain versions support different output enforcement methods.
        for kwargs in ({"method": "function_calling"}, {"method": "json_mode"}, {}):
            try:
                return self.llm.with_structured_output(CallSummary, **kwargs)
            except TypeError:
                continue
        return self.llm.with_structured_output(CallSummary)

    @staticmethod
    def fallback(state: CallState) -> CallState:
        text = state.get("clean_content") or state.get("content") or ""
        if not text:
            return state

        # Best-effort retry using structured output with alternate enforcement methods.
        llm = ChatOpenAI(model="gpt-4o", temperature=0)
        prompt = PromptTemplate.from_template(
            "Summarize the following call transcript.\n"
            "Return: summary, key_points, action_items, tags, highlights.\n\n"
            "Transcript:\n{text}"
        )

        for kwargs in ({"method": "function_calling"}, {"method": "json_mode"}, {}):
            try:
                structured_llm = llm.with_structured_output(CallSummary, **kwargs)
                result = (prompt | structured_llm).invoke({"text": text})
                if isinstance(result, CallSummary):
                    state["summary"] = result.summary
                    state["key_points"] = result.key_points
                    state["action_items"] = result.action_items
                    state["tags"] = result.tags
                    state["highlights"] = result.highlights
                    return state
                state["summary"] = (result or {}).get("summary", state.get("summary", ""))
                state["key_points"] = (result or {}).get("key_points", state.get("key_points") or [])
                state["action_items"] = (result or {}).get("action_items", state.get("action_items") or [])
                state["tags"] = (result or {}).get("tags", state.get("tags") or [])
                state["highlights"] = (result or {}).get("highlights", state.get("highlights") or [])
                return state
            except Exception:
                continue

        # Last resort heuristic fallback that avoids breaking the UI.
        summary = (text.strip()[:800] + ("…" if len(text.strip()) > 800 else "")).strip()
        state["summary"] = summary or state.get("summary", "")
        state["key_points"] = state.get("key_points") or []
        state["action_items"] = state.get("action_items") or []
        state["tags"] = state.get("tags") or []
        state["highlights"] = state.get("highlights") or []
        return state