from app.ai.gemini_client import GeminiClient
import asyncio
import json

from app.ai.openai_client import OpenAIClient
from app.ai.prompts import build_system_prompt
from app.ai.validator import validate_ai_response
from app.config import OPENAI_API_KEY, GEMINI_API_KEY, GEMINI_MODEL, ANTHROPIC_API_KEY, ANTHROPIC_MODEL
from app.specialties.registry import get_specialty_config


SPECIALTY_RULES = [
    {
        "code": "psychiatry",
        "name": "Psychiatry & Mental Health Care",
        "keywords": [
            "depression", "depressed", "mental", "anxiety", "panic", "stress", "sadness", "suicidal", "suicide", "sleep", "insomnia", "emotional", "mood", "lonely", "hopeless",
            "மன அழுத்தம்", "மனச்சோர்வு", "தூக்கமின்மை", "கவலை", "உணர்ச்சி", "தற்கொலை", "டிப்ரெஷன்",
            "डिप्रेशन", "तनाव", "चिंता", "उदासी", "नींद न आना", "आत्महत्या",
            "మానసిక ఒత్తిడి", "డిప్రెషన్", "నిద్రలేమి",
            "ಮನೋಒತ್ತಡ", "ಖಿನ್ನತೆ", "ನಿದ್ರಾಹೀನತೆ"
        ]
    },
    {
        "code": "cardiology",
        "name": "Cardiology (Heart Care)",
        "keywords": [
            "chest pain", "heart", "palpitation", "breathless", "cardiac", "chest pressure",
            "நெஞ்சு வலி", "இருதய வலி", "மூச்சு திணறல்", "छाती में दर्द", "दिल की धड़कन", "గుండె నొప్పి"
        ]
    },
    {
        "code": "dermatology",
        "name": "Dermatology (Skin & Hair Care)",
        "keywords": [
            "skin", "rash", "itching", "eczema", "acne", "pimples", "fungal", "allergy",
            "தோல் அரிப்பு", "தழும்பு", "படை", "கொப்புளம்", "त्वचा में खुजली", "दाद", "पिंपल्स"
        ]
    },
    {
        "code": "neurology",
        "name": "Neurology (Brain & Nervous System)",
        "keywords": [
            "headache", "dizziness", "seizure", "stroke", "numbness", "paralysis", "migraine",
            "தலைவலி", "மயக்கம்", "கை கால் மரத்துப்போதல்", "सिरदर्द", "चक्कर आना", "दौरा"
        ]
    },
    {
        "code": "orthopedics",
        "name": "Orthopedics (Bone & Joint Care)",
        "keywords": [
            "knee pain", "joint pain", "fracture", "bone", "back pain", "arthritis", "ligament",
            "மூட்டு வலி", "இடுப்பு வலி", "எலும்பு முறிவு", "जोड़ों का दर्द", "घुटने में दर्द", "पीठ दर्द"
        ]
    },
    {
        "code": "gastroenterology",
        "name": "Gastroenterology (Digestive Health)",
        "keywords": [
            "stomach pain", "stomach ache", "acidity", "vomiting", "diarrhea", "indigestion", "gastric",
            "வயிற்று வலி", "வாந்தி", "பேதி", "நெஞ்செரிச்சல்", "पेट दर्द", "उल्टी", "दस्त"
        ]
    },
    {
        "code": "pulmonology",
        "name": "Pulmonology (Respiratory Care)",
        "keywords": [
            "cough", "asthma", "wheezing", "phlegm", "cold", "flu", "respiratory",
            "இருமல்", "சளி", "ஆஸ்துமா", "மூச்சு வாங்குதல்", "खांसी", "बलगम", "दमा"
        ]
    },
    {
        "code": "ent",
        "name": "ENT (Ear, Nose & Throat)",
        "keywords": [
            "ear pain", "throat pain", "sore throat", "tonsils", "nasal block", "hearing",
            "காது வலி", "தொண்டை வலி", "மூக்கடைப்பு", "कान का दर्द", "गले में खराश"
        ]
    },
    {
        "code": "ophthalmology",
        "name": "Ophthalmology (Eye Care)",
        "keywords": [
            "eye pain", "blurred vision", "red eye", "cataract", "glaucoma",
            "கண் வலி", "பார்வை மங்கல்", "சிவப்பு கண்", "आँख में दर्द", "धुंधला दिखना"
        ]
    },
    {
        "code": "pediatrics",
        "name": "Pediatrics (Child Health)",
        "keywords": [
            "baby", "child", "infant", "toddler", "pediatric",
            "குழந்தைக்கு", "குழந்தை", "बच्चे को", "शिशु"
        ]
    }
]


def classify_specialty_from_symptoms(text: str) -> dict:
    if not text:
        return {
            "type": "general",
            "code": "general_medicine",
            "name": "General Medicine / Primary Care",
            "confidence": 0.85,
            "reason": "Default primary care specialty assignment."
        }

    text_lower = text.lower()
    for rule in SPECIALTY_RULES:
        for kw in rule["keywords"]:
            if kw in text_lower:
                return {
                    "type": "speciality",
                    "code": rule["code"],
                    "name": rule["name"],
                    "confidence": 0.95,
                    "reason": f"Automatically classified based on clinical keyword: '{kw}'"
                }

    return {
        "type": "general",
        "code": "general_medicine",
        "name": "General Medicine / Primary Care",
        "confidence": 0.85,
        "reason": "General primary care consultation."
    }


class AIOrchestrator:

    def __init__(self):
        self.openai = OpenAIClient()
        self.gemini = GeminiClient()

    async def process(
        self,
        session,
        state,
        conversation_history,
        user_message,
    ):
        specialty = get_specialty_config(
            session.category_code
        )

        current_state = json.loads(
            state.state_json
        )

        # Dynamic Specialty Triage Classification
        suggested_specialty = classify_specialty_from_symptoms(user_message)

        context = {
            "patient": {
                "id": session.patient_id,
                "name": session.patient_name,
                "age": session.age,
                "gender": session.gender,
            },
            "category": {
                "type": session.category_type,
                "code": session.category_code,
                "name": session.category_name,
            },
            "specialty_config": specialty,
            "language": session.language,
            "current_stage": session.stage,
            "current_state": current_state,
            "conversation_history": conversation_history,
            "new_patient_message": user_message,
        }

        prompt = build_system_prompt(context)

        # Try real LLM dynamic generation (Gemini -> OpenAI -> Local Fallback)
        parsed = None

        # 1. Try Gemini LLM Generation first
        if self.gemini.api_key and self.gemini.api_key.strip() not in ["your_gemini_key", "your_gemini_api_key", ""]:
            try:
                raw_res = await self.gemini.generate_response(prompt, user_message)
                if raw_res:
                    cleaned_res = raw_res.replace("```json", "").replace("```", "").strip()
                    parsed = json.loads(cleaned_res)
                    print(f"[Gemini Dynamic Response SUCCESS]: {parsed.get('reply')}")
            except Exception as e:
                print(f"[Gemini Dynamic Intake Warning]: {e}")

        # 2. Try OpenAI LLM Generation if Gemini was not available or failed
        if not parsed and self.openai.api_key and self.openai.api_key.strip() not in ["your_openai_api_key", "your_openai_key", ""]:
            try:
                raw_res = await self.openai.generate_response(prompt, user_message)
                if raw_res:
                    parsed = json.loads(raw_res)
                    print(f"[OpenAI Dynamic Response SUCCESS]: {parsed.get('reply')}")
            except Exception as e:
                print(f"[OpenAI Dynamic Intake Warning]: {e}")

        # 3. Fallback to intelligent symptom-tailored intake engine
        if not parsed:
            parsed = self._simulate_specialty_response(session, specialty, user_message, len(conversation_history), suggested_specialty)

        parsed["suggested_specialty"] = suggested_specialty
        validated = validate_ai_response(parsed)
        res_dict = validated.model_dump()
        res_dict["suggested_specialty"] = suggested_specialty
        return res_dict

    def _simulate_specialty_response(self, session, specialty, user_message, turn_count=0, suggested_specialty=None):
        stage = session.stage
        lang = (session.language or "ta-IN").lower()
        msg_lower = (user_message or "").lower().strip()
        
        # Detect Mental Health / Depression / Psychiatry Complaints
        is_mental_health = False
        if suggested_specialty and suggested_specialty.get("code") == "psychiatry":
            is_mental_health = True
        elif session.category_code == "psychiatry":
            is_mental_health = True
        else:
            mental_kws = [
                "depress", "mental", "anxiet", "panic", "stress", "sadness", "suicid", "sleep", "insomn", "mood", "hopeless", "lonely",
                "மன அழுத்தம்", "மனச்சோர்வு", "தூக்கம்", "கவலை", "உணர்ச்சி", "டிப்ரெஷன்", "டிப்ரஷன்",
                "डिप्रेशन", "तनाव", "चिंता", "उदासी", "नींद"
            ]
            is_mental_health = any(k in msg_lower for k in mental_kws)

        # Check crisis/suicidal intent
        is_crisis = any(k in msg_lower for k in ["suicid", "end my life", "kill myself", "தற்கொலை", "आत्महत्या"])

        negative_keywords = [
            "no", "no more", "nothing", "nothing else", "none", "no illness", "no problem",
            "இல்லை", "இல்லைங்க", "இல்ல", "ஒன்றுமில்லை", "ஒன்னுமில்ல", "ஏதுமில்லை",
            "नहीं", "कुछ नहीं", "कोई नहीं", "नाही",
            "లేదు", "ఏమీ లేదు", "లేదండి",
            "ಇಲ್ಲ", "ಇಲ್ಲವೇ ಇಲ್ಲ", "ಏನೂ ಇಲ್ಲ",
            "ഇല്ല", "ഒന്നുമില്ല", "ഇല്ല കട്ടായം",
            "না", "কিছু না", "নাই",
            "नाही", "काही नाही",
            "નથી", "કાંઈ નથી",
            "ਨਹੀਂ", "ਕੁਝ ਨਹੀਂ"
        ]
        is_negative = any(k in msg_lower for k in negative_keywords)

        # Empathetic Mental Health Specific Questions
        if is_mental_health:
            if lang.startswith("ta"):
                conclusion_text = "சரிங்க, எல்லா விவரத்தையும் மனநல மருத்துவருகிட்ட (Psychiatrist) குடுத்துடுறேன். ரொம்ப நன்றிங்க!"
                q_onset_severity = "புரியுதுங்க. மன அழுத்தம் மற்றும் மனச்சோர்வை பற்றிப் பேசுவது ரொம்ப முக்கியமானது. இந்த கவலை எப்போலிருந்து இருக்குன்னு சொல்ல முடியுமா?"
                q_associated_symptoms = "சரிங்க. இதனால உங்களுக்கு தூக்கம் இல்ல அன்றாட வேலைகள்ல ஏதாவது சிரமம் இருக்கா?"
                q_past_history = "புரியுதுங்க. இதுக்காக ஏற்கனவே ஏதாவது கவுன்சிலிங் இல்ல மாத்திரை எடுக்குறீங்களா?"
                q_meds_allergies = "சரிங்க, உங்களுக்கு ஏதேனும் மாத்திரை அலர்ஜி இருக்கா?"
            elif lang.startswith("hi"):
                conclusion_text = "ठीक है जी, मैं यह सारी जानकारी मनोचिकित्सक (Psychiatrist) डॉक्टर साहब को दे देती हूँ। शुक्रिया!"
                q_onset_severity = "जी समझ गया। मानसिक तनाव और डिप्रेशन के बारे में बात करना बहुत जरूरी है। यह परेशानी कब से महसूस हो रही है?"
                q_associated_symptoms = "जी ठीक है, क्या इसकी वजह से आपको नींद या रोज़मर्रा के कामों में कोई परेशानी हो रही है?"
                q_past_history = "ठीक है जी, क्या आप इसके लिए कोई काउंसलिंग या मेडिसिन ले रहे हैं?"
                q_meds_allergies = "अच्छा जी, आपको किसी दवाई से एलर्जी तो नहीं है?"
            else:
                conclusion_text = "Got it. I have compiled all these details for the specialist psychiatrist to review. Thank you!"
                q_onset_severity = "I hear you, and it takes courage to talk about how you're feeling. How long have you been experiencing this stress or depression?"
                q_associated_symptoms = "Understood. Is this affecting your sleep, appetite, or daily activities?"
                q_past_history = "Got it. Are you currently seeing a therapist, receiving counseling, or taking any medications?"
                q_meds_allergies = "Understood. Do you have any known medication or drug allergies?"

            if is_crisis:
                if lang.startswith("ta"):
                    conclusion_text = "உடனடி உதவிக்கு மனநல உதவி எண் Tele-MANAS (14416) அல்லது 108 ஐ தொடர்பு கொள்ளவும். உங்கள் விவரங்களை மருத்துவரிடம் பகிர்ந்துள்ளேன்."
                elif lang.startswith("hi"):
                    conclusion_text = "अगर आपको मदद की ज़रूरत है तो तुरंत Tele-MANAS (14416) या 108 पर संपर्क करें। मैंने डॉक्टर साहब को सूचित कर दिया है।"
                else:
                    conclusion_text = "If you are in distress, please reach out to Tele-MANAS (14416) or emergency 108 immediately. I have shared your details with our doctor."

        else:
            # Physical General Medical Intake Prompts
            if lang.startswith("ta"):
                conclusion_text = "சரிங்க, எல்லா விவரத்தையும் டாக்டர்கிட்ட குடுத்துடுறேன். ரொம்ப நன்றிங்க!"
                q_onset_severity = "சரிங்க, இந்த பிரச்சனை எத்தனை நாளா இருக்குன்னு சொல்ல முடியுமா?"
                q_associated_symptoms = "புரியுதுங்க. உங்களுக்கு வேறு ஏதேனும் உடம்பு பிரச்சனை இருக்கா?"
                q_past_history = "சரிங்க. உங்களுக்கு பிபி, சுகர், ஆஸ்துமா போன்ற வேறு ஏதேனும் உடம்பு பிரச்சனை இருக்கா?"
                q_meds_allergies = "சரிங்க, இதுக்காக ஏதாவது மாத்திரை இல்ல மெடிசின் சாப்பிட்டீங்களா? உங்களுக்கு ஏதேனும் மாத்திரை அலர்ஜி இருக்கா?"
            elif lang.startswith("hi"):
                conclusion_text = "ठीक है जी, मैं यह सारी जानकारी डॉक्टर साहब को दे देती हूँ। शुक्रिया!"
                q_onset_severity = "समझ गया जी, यह तकलीफ कितने दिनों से हो रही है?"
                q_associated_symptoms = "ठीक है जी, क्या कोई और परेशानी भी महसूस हो रही है?"
                q_past_history = "जी ठीक है, क्या आपको बीपी, शुगर, या अस्थामा जैसी कोई पुरानी बीमारी है?"
                q_meds_allergies = "समझ गया, इसके लिए कोई दवाई ली है आपने? किसी दवा से एलर्जी तो नहीं है?"
            elif lang.startswith("te"):
                conclusion_text = "చాలా ధన్యవాదాలు! మీ వివరాలన్నీ డాక్టర్ గారికి అందిస్తాను."
                q_onset_severity = "అర్థమైందండి, ఈ ఇబ్బంది ఎన్ని రోజుల నుండి ఉంది?"
                q_associated_symptoms = "సరేనండి, వేరే ఏమైనా ఇబ్బంది ఉందా?"
                q_past_history = "అర్థమైంది, మీకు బీపీ, షుగర్ లేదా ఆస్తమా వంటి పాత ఆరోగ్య సమస్యలు ఉన్నాయా?"
                q_meds_allergies = "సరేనండి, దీనికి ఏమైనా మందులు వాడారా? ఏమైనా మందుల అలర్జీ ఉందా?"
            elif lang.startswith("kn"):
                conclusion_text = "ಸರಿ, ಈ ಎಲ್ಲಾ ವಿವರಗಳನ್ನು ಡಾಕ್ಟರ್‌ಗೆ ತಿಳಿಸ್ತೀನಿ. ಧನ್ಯವಾದಗಳು!"
                q_onset_severity = "ಅರ್ಥವಾಯಿತು, ಈ ತೊಂದರೆ ಎಷ್ಟು ದಿನದಿಂದ ಇದೆ?"
                q_associated_symptoms = "ಸರಿ, ಬೇರೆ ಏನಾದರೂ ತೊಂದರೆ ಇದೆಯಾ?"
                q_past_history = "ತಿಳಿಯಿತು, ನಿಮಗೆ ಬಿಪಿ, ಶುಗರ್ ಅಥವಾ ಅಸ್ತಮಾದಂತಹ ಹಳೆಯ ತೊಂದರೆ ಇದೆಯಾ?"
                q_meds_allergies = "ಸರಿ, ಇದಕ್ಕೆ ಏನಾದರೂ ಮಾತ್ರೆ ತಗೊಂಡಿದ್ದೀರಾ? ಏನಾದರೂ ಮಾತ್ರೆ ಅಲರ್ಜಿ ಇದೆಯಾ?"
            elif lang.startswith("ml"):
                conclusion_text = "ശരി, എല്ലാ വിവരങ്ങളും ഡോക്ടർക്ക് കൈമാറാം. വളരെ നന്ദി!"
                q_onset_severity = "ശരി, ഈ ബുദ്ധിമുട്ട് എത്ര ദിവസമായിട്ടുണ്ട്?"
                q_associated_symptoms = "മനസ്സിലായി, മറ്റ് ബുദ്ധിമുട്ടുകൾ ഉണ്ടോ?"
                q_past_history = "ശരി, പ്രഷർ (BP), ഷുഗർ, അല്ലെങ്കിൽ ആസ്ത്മ തുടങ്ങിയ രോഗങ്ങൾ മുൻപുണ്ടോ?"
                q_meds_allergies = "മനസ്സിലായി, ഇതിനായി മരുന്നുകൾ കഴിക്കുന്നുണ്ടോ? മരുന്ന് അലർജിയുണ്ടോ?"
            else:
                conclusion_text = "Got it. I have shared all these details directly with the doctor. Thank you!"
                q_onset_severity = "Got it. How long has this issue been bothering you?"
                q_associated_symptoms = "Understood. Are you experiencing any other symptoms?"
                q_past_history = "Got it. Do you have a history of high blood pressure, diabetes, or asthma?"
                q_meds_allergies = "Understood. Have you taken any medication for this yet, and do you have any allergies?"

        # Fast conclusion if patient says "No" / "Nothing" at any stage after initial chief complaint
        if is_negative and stage != "greeting":
            return {
                "reply": conclusion_text,
                "next_stage": "confirmation",
                "state_updates": {
                    "reason_for_visit": None,
                    "symptoms": [],
                    "medical_history": ["None reported"],
                    "medications": [],
                    "allergies": ["None"]
                },
                "completed": True
            }

        if stage == "greeting":
            return {
                "reply": q_onset_severity,
                "next_stage": "reason_for_visit",
                "state_updates": {
                    "reason_for_visit": user_message,
                    "symptoms": [{"name": user_message, "status": "reported"}],
                },
                "completed": False
            }
        elif stage == "reason_for_visit":
            return {
                "reply": q_associated_symptoms,
                "next_stage": "symptom_collection",
                "state_updates": {
                    "symptoms": [{"onset_and_severity": user_message}],
                },
                "completed": False
            }
        elif stage == "symptom_collection":
            return {
                "reply": q_past_history,
                "next_stage": "symptom_details",
                "state_updates": {
                    "symptoms": [{"associated_findings": user_message}],
                },
                "completed": False
            }
        elif stage == "symptom_details":
            return {
                "reply": q_meds_allergies,
                "next_stage": "medications",
                "state_updates": {
                    "medical_history": [user_message] if not is_negative else ["None"],
                },
                "completed": False
            }
        elif stage == "medications":
            return {
                "reply": conclusion_text,
                "next_stage": "confirmation",
                "state_updates": {
                    "medications": [user_message] if not is_negative else [],
                    "allergies": [user_message] if not is_negative else ["NKDA"],
                },
                "completed": True
            }
        else:
            return {
                "reply": conclusion_text,
                "next_stage": "confirmation",
                "state_updates": {},
                "completed": True
            }