import os
import json
import httpx
from dotenv import load_dotenv
from supabase import create_client

load_dotenv()

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

supabase = None
if all([GROQ_API_KEY, SUPABASE_URL, SUPABASE_KEY]):
    try:
        supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
    except Exception as e:
        print(f"WARNING: Could not initialize Supabase: {e}")


# ─────────────────────────────────────────
# TOOL FUNCTIONS
# ─────────────────────────────────────────

def get_patient_by_name(name: str):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("patients").select("*").ilike("name", f"%{name}%").execute()
    return res.data if res.data else {"error": "Patient not found"}

def get_patients_by_doctor(doctor: str):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("patients").select("*").ilike("doctor", f"%{doctor}%").execute()
    return res.data if res.data else {"error": "No patients found for this doctor"}

def get_admitted_patients():
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("patients").select("*").eq("status", "admitted").execute()
    return res.data if res.data else {"message": "No admitted patients currently"}

def get_appointments_by_date(date: str):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("appointments").select("*").eq("date", date).execute()
    return res.data if res.data else {"message": f"No appointments on {date}"}

def get_patient_appointments(patient_name: str):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("appointments").select("*").ilike("patient_name", f"%{patient_name}%").execute()
    return res.data if res.data else {"error": "No appointments found for this patient"}

def add_patient(name: str, age: int, gender: str, disease: str, doctor: str, admission_date: str):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    data = {
        "name": name,
        "age": age,
        "gender": gender,
        "disease": disease,
        "doctor": doctor,
        "admission_date": admission_date,
        "status": "admitted"
    }
    res = supabase.table("patients").insert(data).execute()
    return {"success": True, "patient": res.data[0]} if res.data else {"error": "Failed to add patient"}

def discharge_patient(patient_id: int):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("patients").update({"status": "discharged"}).eq("id", patient_id).execute()
    return {"success": True, "message": f"Patient {patient_id} discharged"} if res.data else {"error": "Patient not found"}

def update_appointment_status(appointment_id: int, status: str):
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    if status not in ["scheduled", "completed", "cancelled"]:
        return {"error": "Invalid status. Use: scheduled, completed, cancelled"}
    res = supabase.table("appointments").update({"status": status}).eq("id", appointment_id).execute()
    return {"success": True, "message": f"Appointment {appointment_id} updated to {status}"} if res.data else {"error": "Appointment not found"}

def get_disease_statistics():
    if not supabase:
        return {"error": "Database not connected. Check SUPABASE_URL and SUPABASE_KEY."}
    res = supabase.table("patients").select("disease").execute()
    if not res.data:
        return {"message": "No data available"}
    stats = {}
    for row in res.data:
        d = row["disease"]
        stats[d] = stats.get(d, 0) + 1
    return stats

# ─────────────────────────────────────────
# FUNCTION MAP
# ─────────────────────────────────────────

available_functions = {
    "get_patient_by_name": get_patient_by_name,
    "get_patients_by_doctor": get_patients_by_doctor,
    "get_admitted_patients": get_admitted_patients,
    "get_appointments_by_date": get_appointments_by_date,
    "get_patient_appointments": get_patient_appointments,
    "add_patient": add_patient,
    "discharge_patient": discharge_patient,
    "update_appointment_status": update_appointment_status,
    "get_disease_statistics": get_disease_statistics,
}

# ─────────────────────────────────────────
# TOOLS SCHEMA
# ─────────────────────────────────────────

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_patient_by_name",
            "description": "Search and get patient details by name",
            "parameters": {
                "type": "object",
                "properties": {"name": {"type": "string"}},
                "required": ["name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_patients_by_doctor",
            "description": "Get all patients assigned to a specific doctor",
            "parameters": {
                "type": "object",
                "properties": {"doctor": {"type": "string"}},
                "required": ["doctor"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_admitted_patients",
            "description": "Get all currently admitted patients",
            "parameters": {"type": "object", "properties": {}}
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_appointments_by_date",
            "description": "Get all appointments on a specific date (format: YYYY-MM-DD)",
            "parameters": {
                "type": "object",
                "properties": {"date": {"type": "string"}},
                "required": ["date"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_patient_appointments",
            "description": "Get all appointments for a specific patient by name",
            "parameters": {
                "type": "object",
                "properties": {"patient_name": {"type": "string"}},
                "required": ["patient_name"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "add_patient",
            "description": "Register a new patient in the hospital system",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "age": {"type": "integer"},
                    "gender": {"type": "string"},
                    "disease": {"type": "string"},
                    "doctor": {"type": "string"},
                    "admission_date": {"type": "string", "description": "Format: YYYY-MM-DD"}
                },
                "required": ["name", "age", "gender", "disease", "doctor", "admission_date"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "discharge_patient",
            "description": "Discharge a patient by their ID (change status to discharged)",
            "parameters": {
                "type": "object",
                "properties": {"patient_id": {"type": "integer"}},
                "required": ["patient_id"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "update_appointment_status",
            "description": "Update the status of an appointment (scheduled, completed, or cancelled)",
            "parameters": {
                "type": "object",
                "properties": {
                    "appointment_id": {"type": "integer"},
                    "status": {"type": "string", "enum": ["scheduled", "completed", "cancelled"]}
                },
                "required": ["appointment_id", "status"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_disease_statistics",
            "description": "Get count of patients grouped by disease/condition",
            "parameters": {"type": "object", "properties": {}}
        }
    }
]

# ─────────────────────────────────────────
# MAIN AGENT
# ─────────────────────────────────────────

SYSTEM_PROMPT = """You are MedBot, a hospital admin assistant for City General Hospital.

ABSOLUTE RULES - NEVER BREAK THESE:
1. YOU MUST ALWAYS RESPOND IN ENGLISH ONLY. This is mandatory. Never use Hindi, Urdu, Hinglish, or any other language. Not even one word in another language.
2. You are ONLY a hospital admin assistant. Reject all non-hospital queries politely in English.

You help admins with:
- Patient records (search, add, discharge)
- Doctor-patient assignments  
- Appointment management
- Disease statistics

Always use available tools for database queries.
Format responses cleanly with bullet points for multiple records.
"""

async def process_agent_chat(user_message: str):
    try:
        if not GROQ_API_KEY:
            return {"response": "Error: GROQ_API_KEY not configured. Please set it in Space Variables."}
        
        headers = {
            "Authorization": f"Bearer {GROQ_API_KEY}",
            "Content-Type": "application/json"
        }

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message}
        ]

        payload = {
            "model": "llama-3.3-70b-versatile",
            "messages": messages,
            "tools": tools,
            "tool_choice": "auto",
            "temperature": 0.2
        }

        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers=headers,
                json=payload
            )

        # Handle non-200 responses from Groq (e.g. invalid key, rate limit)
        if response.status_code != 200:
            try:
                error_data = response.json()
                error_msg = error_data.get("error", {}).get("message", response.text)
            except Exception:
                error_msg = response.text
            return {"response": f"⚠️ Groq API Error (Status {response.status_code}): {error_msg}"}

        try:
            data = response.json()
        except Exception as e:
            return {"response": f"⚠️ Failed to parse Groq response as JSON. Status code: {response.status_code}. Raw response: {response.text}"}

        print("GROQ RESPONSE:", data)

        if "choices" not in data or not data["choices"]:
            return {"response": f"⚠️ Error: Groq response did not contain 'choices'. Full response: {data}"}

        assistant_message = data["choices"][0]["message"]
        tool_calls = assistant_message.get("tool_calls")

        if not tool_calls:
            return {"response": assistant_message.get("content", "No response text received.")}

        messages.append(assistant_message)

        for tool_call in tool_calls:
            function_name = tool_call["function"]["name"]
            function_args = tool_call["function"].get("arguments")
            try:
                function_args = json.loads(function_args) if function_args else {}
            except Exception as e:
                function_args = {}
                print(f"Error parsing tool arguments: {e}")

            if function_name not in available_functions:
                function_response = {"error": f"Function '{function_name}' is not supported."}
            else:
                function_to_call = available_functions[function_name]
                try:
                    try:
                        function_response = function_to_call(**function_args)
                    except TypeError:
                        function_response = function_to_call()
                except Exception as e:
                    import traceback
                    print(f"Error executing database tool '{function_name}': {e}")
                    traceback.print_exc()
                    function_response = {"error": f"Database query failed in {function_name}: {str(e)}"}

            messages.append({
                "role": "tool",
                "tool_call_id": tool_call["id"],
                "name": function_name,
                "content": json.dumps(function_response)
            })

        second_payload = {
            "model": "llama-3.3-70b-versatile",
            "messages": messages,
            "temperature": 0.2
        }

        async with httpx.AsyncClient(timeout=30.0) as client:
            final_response = await client.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers=headers,
                json=second_payload
            )

        if final_response.status_code != 200:
            try:
                error_data = final_response.json()
                error_msg = error_data.get("error", {}).get("message", final_response.text)
            except Exception:
                error_msg = final_response.text
            return {"response": f"⚠️ Groq Final Response Generation Error (Status {final_response.status_code}): {error_msg}"}

        try:
            final_data = final_response.json()
        except Exception as e:
            return {"response": f"⚠️ Failed to parse Groq final response as JSON. Status code: {final_response.status_code}. Raw response: {final_response.text}"}

        if "choices" not in final_data or not final_data["choices"]:
            return {"response": f"⚠️ Error: Groq final response did not contain 'choices'. Full response: {final_data}"}

        final_answer = final_data["choices"][0]["message"].get("content", "No final answer generated.")
        return {"response": final_answer}

    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        print("RUNTIME ERROR IN process_agent_chat:\n", error_trace)
        return {"response": f"⚠️ Server Runtime Error: {str(e)}\n\nTraceback:\n{error_trace}"}