import yaml, json, os
from crewai import Agent, Task, Crew, Process
from langchain_openai import ChatOpenAI

class PatientChartCrewEngine:
    def __init__(self, folder_path: str, measure_name: str):
        if not measure_name.endswith(".yml"):
            measure_name += ".yml"
        
        file_path = os.path.join(folder_path, measure_name)
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Measure file not found at: {file_path}")

        with open(file_path, 'r') as f:
            self.cfg = yaml.safe_load(f)
        
        # Using a slightly lower temperature for the Architect to ensure syntax stability
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0.5)

    def run(self):
        # AGENT 1: Scribe - Focuses on clinical density
        scribe = Agent(
            role="Senior Clinical Quality Scribe",
            goal="Generate high-density clinical narratives for patient charts.",
            backstory="""You are a senior scribe. You provide extremely detailed HPIs 
            and expanded Social Histories. You avoid symbols like & or %; 
            instead, you write 'and' or 'percent' to ensure data compatibility.""",
            llm=self.llm,
            allow_delegation=False
        )

        # AGENT 2: Architect - Focuses on LaTeX Syntax and escaping
        architect = Agent(
            role="LaTeX Document Architect",
            goal="Convert clinical text into error-free, top-aligned LaTeX source code.",
            backstory="""You are a LaTeX expert. You know that characters like %, &, $, #, 
            and _ must be escaped (e.g., \%) or the code will crash. You prioritize 
            structural integrity and use packages like 'geometry' and 'microtype' 
            to ensure professional formatting.""",
            llm=self.llm,
            allow_delegation=False
        )

        # TASK 1: Narrative Volume
        content_task = Task(
            description=f"""
            RULES: {json.dumps(self.cfg.get('target_rules', {}), indent=2)}
            1. IDENTITY: Unique Name and MRN (e.g., Emily Carter, 92837465).
            2. DENSE HPI: Write a massive clinical narrative (minimum 800 words).
               Focus on Transition to Hospice (Z51.5) and Hysterectomy (58150).
            3. DATES: Use MM/DD/YYYY format for all entries.
            """,
            agent=scribe,
            expected_output="A JSON-like structure containing the clinical narrative sections."
        )

        # TASK 2: Rigid LaTeX Construction
        latex_task = Task(
            description="""
            Assemble the clinical text into a single RAW LaTeX document.
            STRICT REQUIREMENTS:
            - Start with \\documentclass{article}.
            - Use \\usepackage[utf8]{inputenc} and \\usepackage[margin=0.8in]{geometry}.
            - Include \\raggedbottom to keep all text top-aligned.
            - ESCAPE special characters: Change '%' to '\%', '&' to '\&', and '_' to '\_'.
            - Ensure there is a \\begin{document} and \\end{document}.
            - OUTPUT ONLY THE RAW LATEX. DO NOT USE MARKDOWN CODE BLOCKS (```).
            """,
            agent=architect,
            context=[content_task],
            expected_output="Valid LaTeX source code starting with \\documentclass."
        )

        crew = Crew(
            agents=[scribe, architect], 
            tasks=[content_task, latex_task], 
            process=Process.sequential
        )
        
        result = crew.kickoff()
        return result.raw if hasattr(result, 'raw') else str(result)