# summarizer.py
from typing import Callable
from crewai import Agent, Task, Crew, Process
from crewai_tools import tool
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI


class PatientChartSummarizer:
    """
    Builds a single manager-style agent that queries the patient chart
    (backed by Chroma) and produces a comprehensive markdown summary.
    """
    def __init__(self, vectordb: Chroma):
        self.vectordb = vectordb
        self.model = ChatOpenAI(model="gpt-4o", temperature=0)

        @tool("patient_chart_search")
        def patient_chart_search(query: str) -> str:
            """Search the patient chart embeddings for the given query."""
            results = self.vectordb.similarity_search(query, k=15)
            return "\n".join([res.page_content for res in results])

        self.summary_agent = Agent(
            role="Clinical Documentation Manager",
            goal=("Create a comprehensive, well-organized markdown summary of the entire patient chart."),
            verbose=True,
            memory=False,
            tools=[patient_chart_search],
            backstory=(
                "You coordinate information extraction for diagnoses, procedures, labs, vitals, and medications, "
                "then synthesize a clear, clinically relevant chart summary."
            ),
            llm=self.model
        )

        self.summary_task = Task(
            description=(
                "Search and compile a full patient chart summary in markdown.\n\n"
                "# Patient Chart Summary\n\n"
                "## Patient Demographics\n"
                "Include patient name, age, DOB, gender, MRN (if available)\n\n"
                "## Chief Complaint & History\n"
                "Extract reason for visit, chief complaint, relevant history\n\n"
                "## Diagnoses\n"
                "Organize as Primary/Secondary/Chronic/Past Medical History\n\n"
                "## Procedures & Interventions\n"
                "Chronological list with outcomes\n\n"
                "## Laboratory Results\n"
                "Organize by category; highlight abnormal findings\n\n"
                "## Vital Signs & Measurements\n"
                "Present trends and significant findings\n\n"
                "## Medications\n"
                "Current, discontinued, allergies & adverse reactions\n\n"
                "## Clinical Assessment & Plan\n"
                "Summarize assessment and plan sections\n\n"
                "## Key Clinical Findings\n"
                "Bullet key takeaways and recommendations\n"
            ),
            expected_output="Comprehensive patient chart summary in well-formatted markdown.",
            agent=self.summary_agent,
        )

        self.crew = Crew(
            agents=[self.summary_agent],
            tasks=[self.summary_task],
            process=Process.sequential,  # simple, deterministic flow
            verbose=True,
        )

    def summarize_chart(self) -> str:
        result = self.crew.kickoff()
        return str(result)