import os
import json
import pandas as pd
from PyPDF2 import PdfReader
from json_repair import repair_json
from typing import List, Dict, Any
from typing import List, Dict, Any, Optional
from crewai import Agent, Task, Crew, Process
from crewai_tools import SerperDevTool
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma


# ---------- MEAT Validator (MODIFIED) ----------
class MEATValidatorAgent:
    def __init__(self, model: str = "gpt-4o"):
        self.llm = ChatOpenAI(model=model, temperature=0)

        self.agent = Agent(
            role="MEAT Criteria Validator for Active Conditions",
            goal="Validate MEAT criteria for diagnoses that are confirmed to be ACTIVE.",
            backstory="You are a compliance and risk adjustment specialist. Your task is to scrutinize "
                      "chart documentation for ACTIVE conditions to ensure they meet MEAT criteria.",
            verbose=True,
            memory=False,
            llm=self.llm,
        )

    def validate_one(self, diagnosis_entry: dict) -> dict:
        """Validate MEAT criteria for a diagnosis that is confirmed as ACTIVE."""
        diagnosis = diagnosis_entry["diagnosis"]
        icd10 = diagnosis_entry["icd10"]
        context = diagnosis_entry["context"]

        task = Task(
            description=(
                f"The following diagnosis has been classified as ACTIVE: {diagnosis} (ICD-10: {icd10})\n\n"
                f"Patient chart excerpts:\n{context}\n\n"
                "Your task is to validate the MEAT Criteria (Y/N for each) based *only* on the provided text:\n"
                "- MONITOR: Are signs, symptoms, disease progression, or regression being tracked?\n"
                "- EVALUATE: Are test results, medication effectiveness, or patient response to treatment being reviewed?\n"
                "- ASSESS: Did the provider order tests, refer to specialists, or document their clinical assessment/plan?\n"
                "- TREAT: Are medications, therapies, or other therapeutic interventions prescribed or documented?\n\n"
                "Output must be strict JSON:\n"
                "{'meat': {'monitor': true/false, 'evaluate': true/false, 'assess': true/false, 'treat': true/false}, "
                "'rationale': 'short justification for the MEAT findings'}"
            ),
            expected_output="JSON with keys meat and rationale",
            agent=self.agent,
            json_mode=True,
        )

        crew = Crew(agents=[self.agent], tasks=[task], process=Process.sequential, verbose=True)
        result = crew.kickoff()
        result = json.loads(repair_json(result))

        return {
            **diagnosis_entry,
            "meat": result.get("meat", {}),
            "meat_rationale": result.get("rationale", "")
        }

    def run(self, active_diagnoses: list[dict]) -> list[dict]:
        """Loop through all ACTIVE diagnoses and validate MEAT criteria."""
        enriched_results = []
        for entry in active_diagnoses:
            # This function now expects only pre-filtered 'ACTIVE' diagnoses
            print(f"\n[INFO] Validating MEAT for ACTIVE diagnosis: {entry['diagnosis']} ({entry['icd10']})")
            enriched = self.validate_one(entry)
            enriched_results.append(enriched)
            print(f"[MEAT RESULT] {enriched}")
        return enriched_results