from llm_factory import get_llm
from langchain_core.prompts import PromptTemplate
import os
from dotenv import load_dotenv

load_dotenv()

def get_triage_agent():
    llm = get_llm(model_type="text", temperature=0.1)
    
    template = """
    Tu es un Agent de Triage Médical expert pour les urgences de première ligne en Tunisie.
    
    TA MISSION :
    Analyser les symptômes du patient et déterminer le niveau d'urgence.
    
    NIVEAUX D'URGENCE :
    - 🔴 ROUGE (Urgence Vitale) : Intervention immédiate requise (ex: douleur thoracique, détresse respiratoire, AVC).
    - 🟠 ORANGE (Urgent) : Consultation nécessaire dans les 2-4 heures (ex: fièvre élevée mal tolérée, douleur intense, fracture).
    - 🟢 VERT (Non Urgent) : Consultation classique ou conseils (ex: rhume, petite plaie, renouvellement).
    
    SYMPTÔMES DU PATIENT :
    {symptoms}
    
    FORMAT DE RÉPONSE ATTENDU (Strictement en Français) :
    
    NIVEAU : [ROUGE/ORANGE/VERT]
    JUSTIFICATION : [Explication courte et précise en 1 phrase]
    ACTION RECOMMANDÉE : [Ex: "Appeler le 190", "Consultation immédiate", "Prendre RDV demain"]
    """
    
    prompt = PromptTemplate(
        input_variables=["symptoms"],
        template=template,
    )
    
    return prompt | llm

async def run_triage_agent(symptoms: str) -> str:
    agent = get_triage_agent()
    response = await agent.ainvoke({"symptoms": symptoms})
    return response.content