import gradio as gr
from ai_logic.intent_parser import parse_intent
from data.data_loader import list_services, load_services_data, HOSPITALS
from insurance import plans, cost_estimator
from insurance.coverage_explainer import CoverageExplainer

# Valid options for hospitals and plans (used in the conversational flow)
HOSPITAL_CHOICES = list(HOSPITALS.keys())
PLAN_CHOICES = list(plans.SAMPLE_PLANS.keys()) + ["No Insurance"]


def _extract_hospital_and_plan(text):
    """Best-effort extraction of hospital and plan names from free text."""
    if not text:
        return None, None

    # Handle case where text might be a list (multimodal Gradio format)
    if isinstance(text, list):
        # Extract text content from list of content parts
        text_parts = []
        for part in text:
            if isinstance(part, str):
                text_parts.append(part)
            elif isinstance(part, dict) and part.get("type") == "text":
                text_parts.append(part.get("text", ""))
        text = " ".join(text_parts)
    
    if not text:
        return None, None

    lower = text.lower()
    hospital_name = None
    plan_name = None

    # Match hospitals by substring
    for name in HOSPITAL_CHOICES:
        if name.lower() in lower:
            hospital_name = name

    # Match plans by key
    for plan_key in plans.SAMPLE_PLANS.keys():
        if plan_key.lower() in lower:
            plan_name = plan_key

    # Allow "no insurance" as a phrase
    if "no insurance" in lower:
        plan_name = "No Insurance"

    return hospital_name, plan_name


def _get_current_hospital_and_plan(message, history):
    """Look through the chat history (and latest message) for the most recent
    hospital & plan mentioned by the user."""
    hospital_name = None
    plan_name = None

    # History is a list of {"role": ..., "content": ...} dicts (new Gradio format)
    for msg in history or []:
        # Only look at user messages for selections
        if msg.get("role") == "user":
            user_text = msg.get("content", "")
            h, p = _extract_hospital_and_plan(user_text)
            if h:
                hospital_name = h
            if p:
                plan_name = p

    # Also extract from the current message
    h, p = _extract_hospital_and_plan(message)
    if h:
        hospital_name = h
    if p:
        plan_name = p

    return hospital_name, plan_name


def _normalize_message(text):
    """Convert message content to a plain string (handles Gradio's multimodal format)."""
    if not text:
        return ""
    
    if isinstance(text, list):
        text_parts = []
        for part in text:
            if isinstance(part, str):
                text_parts.append(part)
            elif isinstance(part, dict) and part.get("type") == "text":
                text_parts.append(part.get("text", ""))
        return " ".join(text_parts)
    
    return text


def respond(message, history):
    """Main response function - now asks for hospital/plan conversationally."""

    # Normalize message in case it's a list (multimodal Gradio format)
    message = _normalize_message(message)

    # Determine the user's current hospital & plan from history/message
    hospital_name, plan_name = _get_current_hospital_and_plan(message, history)

    # If either is missing, show the options and ask the user to choose
    if hospital_name is None or plan_name is None:
        hospitals_list = "\n".join(f"• {name}" for name in HOSPITAL_CHOICES)
        plans_list = "\n".join(f"• {name}" for name in PLAN_CHOICES)

        # Case 1: Both missing - show full welcome message
        if hospital_name is None and plan_name is None:
            return f"""Hi, I'm the THICC Cost Chatbot. 🏥

Before I can estimate your costs, tell me **which hospital** you're using and **what insurance plan** you have.

**Hospitals I currently support:**
{hospitals_list}

**Insurance options I support:**
{plans_list}

Please reply with something like:
• "I'm going to UCLA Medical Center and I have a PPO plan."
• "Cedars-Sinai Medical Center with No Insurance."
"""

        # Case 2: Have hospital, need insurance plan
        if hospital_name is not None and plan_name is None:
            return f"""Great, I see you're going to **{hospital_name}**! 🏥

Now, what **insurance plan** do you have?

**Insurance options I support:**
{plans_list}

Please reply with your plan, like "PPO" or "No Insurance".
"""

        # Case 3: Have insurance plan, need hospital
        if hospital_name is None and plan_name is not None:
            return f"""Got it, you have **{plan_name}**! 📋

Now, which **hospital** are you going to?

**Hospitals I currently support:**
{hospitals_list}

Please reply with the hospital name, like "UCLA Medical Center".
"""

    # --- 1) Coverage questions first ---
    if CoverageExplainer.identify_coverage_question(message):
        # Figure out which term (deductible, copay, coinsurance, etc.)
        term = CoverageExplainer.get_matching_term(message)

        if term:
            # Explain the specific term (NOT the whole message)
            explanation = CoverageExplainer.explain_term(term)

            # Add plan-specific context if a sample plan is selected
            if plan_name != "No Insurance" and plan_name in plans.SAMPLE_PLANS:
                plan = plans.SAMPLE_PLANS[plan_name]
                plan_details = {
                    "deductible": plan.deductible,
                    "copay": plan.copay,
                    "coinsurance": plan.coinsurance,
                }
                explanation += "\n\n---\n\n"
                explanation += CoverageExplainer.format_plan_coverage_summary(
                    plan_name, plan_details
                )

            return explanation
        else:
            # If we can't match a specific term, give the full coverage explainer
            return CoverageExplainer.explain_all_terms()

    # --- 2) Service cost estimation path ---
    hospital_data_path = HOSPITALS.get(hospital_name)
    services_data = load_services_data(hospital_data_path)
    requested_info = parse_intent(message, services_data, hospital_name=hospital_name)

    if requested_info is None or requested_info == "list_services":
        services_list = list_services(services_data)
        response = (
            "**Available services:**\n"
            + "\n".join(f"• {service}" for service in services_list)
        )
        response += (
            "\n\n💡 **Tip**: You can ask me about insurance terms like "
            "'What is a deductible?' or 'Explain coinsurance'."
        )
        return response

    service_data = services_data[
        services_data["intent"].str.contains(requested_info, case=False, na=False)
    ]

    if service_data.empty:
        return (
            "Sorry, no information found for your request.\n\nYou can:\n"
            "• Ask about available services\n"
            "• Ask about insurance terms (e.g., 'What is a copay?')\n"
            "• Get cost estimates for specific procedures"
        )

    service_description = service_data.iloc[0]["description"]
    price = service_data.iloc[0]["negotiated_rate"]

    # Map plan name -> InsurancePlan object
    if plan_name == "No Insurance":
        plan = plans.NO_INSURANCE_PLAN
    else:
        plan = plans.SAMPLE_PLANS.get(plan_name, plans.NO_INSURANCE_PLAN)

    cost = cost_estimator.estimate_cost(price, plan, deductible_met=True)

    # --- 3) Format response with cost breakdown ---

    # Special handling for No Insurance so messaging isn't confusing
    if plan_name == "No Insurance":
        response = f"""**Cost Estimate for {service_description}**

• Hospital: {hospital_name}
• Insurance Plan: {plan_name}
• Estimated Cost: **${cost:.2f}**

Because you selected **No Insurance**, this demo assumes you pay the full negotiated rate.

💡 **Understanding your cost**:
• Negotiated rate: ${price:.2f}
• Your insurance covers: $0.00
• You pay: ${cost:.2f}

If you want to see how deductibles, copays, and coinsurance work, try asking:
• "What is a deductible?"
• "Explain coinsurance"
"""
        return response

    # For actual plans
    response = f"""**Cost Estimate for {service_description}**

• Hospital: {hospital_name}
• Insurance Plan: {plan_name}
• Estimated Cost: **${cost:.2f}**

This estimate assumes your deductible has been met. 

💡 **Understanding your cost**:
• Negotiated rate: ${price:.2f}
• Your insurance covers: ${price - cost:.2f}
• You pay: ${cost:.2f}"""

    # Explain payment type
    if getattr(plan, "copay", None) and cost == plan.copay:
        response += (
            f"\n\n*You're paying a fixed copay of ${plan.copay:.2f} for this service.*"
        )
    elif getattr(plan, "coinsurance", None) and plan.coinsurance > 0:
        response += (
            f"\n\n*You're paying {plan.coinsurance*100:.0f}% coinsurance "
            f"({plan.coinsurance*100:.0f}% of ${price:.2f}).*"
        )

    response += (
        "\n\n**Need help?** Ask me 'What is coinsurance?' "
        "or any other insurance term!"
    )

    return response


# Gradio interface
demo = gr.ChatInterface(
    fn=respond,
    title="THICC Cost Chatbot 🏥",
    description="""Get healthcare cost estimates and understand your insurance coverage.
    
**What you can ask:**
• Cost estimates: "How much does an MRI cost?"
• Coverage terms: "What is a deductible?" or "Explain coinsurance"
• Available services: "What services are available?"
""",
    chatbot=gr.Chatbot(height="70vh"),
)

if __name__ == "__main__":
    demo.launch(ssr_mode=False)