import json
from app.ai.openai_client import OpenAIClient
from app.config import OPENAI_API_KEY
from app.schemas.onboarding import PatientOnboardingResult


class ResultExtractionService:

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

    async def extract_result(
        self,
        session,
        state,
        conversation_history,
    ) -> dict:
        current_state = json.loads(
            state.state_json
        )

        prompt = f"""
You are a hospital patient onboarding data extraction assistant.

Your task is to convert the completed onboarding information
into structured patient intake data.

You are NOT a doctor.
Do not diagnose the patient.
Do not add information that was not provided.

PATIENT:
{json.dumps({
    "id": session.patient_id,
    "name": session.patient_name,
    "age": session.age,
    "gender": session.gender,
}, indent=2)}

CATEGORY:
{json.dumps({
    "type": session.category_type,
    "code": session.category_code,
    "name": session.category_name,
}, indent=2)}

CURRENT STATE:
{json.dumps(current_state, indent=2)}

CONVERSATION:
{json.dumps(conversation_history, indent=2)}

Return ONLY JSON conforming to this exact structure:
{{
    "patient_id": "{session.patient_id}",
    "category": {{
        "type": "{session.category_type}",
        "code": "{session.category_code}",
        "name": "{session.category_name}"
    }},
    "chief_complaint": "string or null",
    "symptoms": [
        {{
            "name": "string",
            "duration": "string or null",
            "severity": "string or null",
            "description": "string or null"
        }}
    ],
    "medical_history": [],
    "medications": [],
    "allergies": [],
    "additional_information": "string or null",
    "language": "{session.language}",
    "onboarding_completed": true
}}
"""

        if OPENAI_API_KEY and OPENAI_API_KEY.strip() not in ["your_openai_api_key", "your_openai_key", ""]:
            try:
                raw_response = await self.openai.generate_response(
                    system_prompt=prompt,
                    user_message="Generate the final structured patient onboarding result.",
                )
                cleaned = raw_response.strip()
                if cleaned.startswith("```json"):
                    cleaned = cleaned[7:]
                if cleaned.startswith("```"):
                    cleaned = cleaned[3:]
                if cleaned.endswith("```"):
                    cleaned = cleaned[:-3]
                cleaned = cleaned.strip()

                parsed = json.loads(cleaned)
                validated = PatientOnboardingResult.model_validate(parsed)
                return validated.model_dump()
            except Exception as e:
                print(f"[Result Extraction Error]: {e}, falling back to structured state")


        # Fallback simulator from state
        chief_complaint = current_state.get("reason_for_visit")
        if isinstance(chief_complaint, dict):
            chief_complaint = chief_complaint.get("value")

        symptoms_raw = current_state.get("symptoms")
        symptoms_list = []
        if isinstance(symptoms_raw, dict):
            symptoms_list = symptoms_raw.get("items", [])
        elif isinstance(symptoms_raw, list):
            symptoms_list = symptoms_raw

        meds_raw = current_state.get("medications")
        meds_list = meds_raw.get("items", []) if isinstance(meds_raw, dict) else (meds_raw or [])

        allergies_raw = current_state.get("allergies")
        allergies_list = allergies_raw.get("items", []) if isinstance(allergies_raw, dict) else (allergies_raw or [])

        history_raw = current_state.get("medical_history")
        history_list = history_raw.get("items", []) if isinstance(history_raw, dict) else (history_raw or [])

        fallback_result = {
            "patient_id": session.patient_id,
            "category": {
                "type": session.category_type,
                "code": session.category_code,
                "name": session.category_name,
            },
            "chief_complaint": chief_complaint or "Consultation Intake",
            "symptoms": symptoms_list or [{"name": chief_complaint or "Reported discomfort", "duration": "recent", "severity": "mild", "description": None}],
            "medical_history": history_list,
            "medications": meds_list,
            "allergies": allergies_list,
            "additional_information": None,
            "language": session.language,
            "onboarding_completed": True,
        }
        validated = PatientOnboardingResult.model_validate(fallback_result)
        return validated.model_dump()