import os
import json
import operator
from typing import TypedDict, Annotated

import gradio as gr
from gradio_client import Client, handle_file
from neo4j import GraphDatabase
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langchain_core.documents import Document
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langgraph.graph import StateGraph, END

# ==========================================
# 1. CONFIGURATION
# ==========================================
AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_API_BASE = os.environ.get("AZURE_OPENAI_API_BASE_URL")
AZURE_DEPLOYMENT_NAME = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME")
NEO4J_URI = os.environ.get("NEO4J_URI")
NEO4J_USERNAME = os.environ.get("NEO4J_USERNAME", "neo4j")
NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD")

# ==========================================
# 2. SETUP
# ==========================================

llm = ChatOpenAI(
    api_key=AZURE_OPENAI_API_KEY,
    base_url=AZURE_OPENAI_API_BASE,
    model=AZURE_DEPLOYMENT_NAME,
    temperature=0
)
try: ChatOpenAI.model_rebuild()
except: pass

def ingest_data():
    if not NEO4J_URI or not NEO4J_PASSWORD:
        print("⚠️ Neo4j Credentials missing. Skipping ingestion.")
        return
    print("--- 🔄 Initializing Knowledge Graph ---")
    try:
        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
        prime_kg_data = [
            ("Ischemic Stroke", "associated_with", "Hypertension"),
            ("Ischemic Stroke", "associated_with", "Atrial Fibrillation"),
            ("Ischemic Stroke", "treated_by", "Alteplase"),
            ("Ischemic Stroke", "treated_by", "Aspirin"),
            ("Ischemic Stroke", "treated_by", "Clopidogrel"),
            ("Alteplase", "mechanism", "Fibrinolytic Agent"),
            ("Alteplase", "contraindicated_in", "Active Bleeding"),
            ("Intracranial Hemorrhage", "associated_with", "Hypertension"),
            ("Intracranial Hemorrhage", "associated_with", "Cerebral Amyloid Angiopathy"),
            ("Intracranial Hemorrhage", "treated_by", "Mannitol"),
            ("Intracranial Hemorrhage", "treated_by", "Nicardipine"),
            ("Mannitol", "mechanism", "Osmotic Diuretic"),
            ("Ischemic Stroke", "presents_as", "Hemiparesis"),
            ("Ischemic Stroke", "presents_as", "Aphasia"),
            ("Intracranial Hemorrhage", "presents_as", "Severe Headache"),
            ("Intracranial Hemorrhage", "presents_as", "Vomiting")
        ]
        query = """
        UNWIND $data as row
        MERGE (h:MedicalEntity {name: row[0]})
        MERGE (t:MedicalEntity {name: row[2]})
        MERGE (h)-[:PRIMEKG_REL {type: row[1]}]->(t)
        """
        with driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            session.run(query, data=prime_kg_data)
        
        aha_guideline_text = """
        For patients eligible for intravenous fibrinolytic therapy, Alteplase (0.9 mg/kg) is recommended within 3 hours of onset of ischemic stroke.
        Mechanical Thrombectomy is indicated for patients with large vessel occlusion (LVO) in the anterior circulation within 6 hours.
        For patients with acute Intracranial Hemorrhage, patients with systolic blood pressure >220 mmHg should be treated with aggressive blood pressure reduction.
        Aspirin is recommended for secondary prevention 24 to 48 hours after thrombolysis.
        """
        docs = [Document(page_content=aha_guideline_text)]
        llm_transformer = LLMGraphTransformer(llm=llm)
        graph_documents = llm_transformer.convert_to_graph_documents(docs)
        with driver.session() as session:
            for doc in graph_documents:
                for node in doc.nodes:
                    session.run("MERGE (n:GuidelineEntity {id: $id, type: $type})", id=node.id, type=node.type)
                for edge in doc.relationships:
                    session.run("""
                        MATCH (s:GuidelineEntity {id: $source})
                        MATCH (t:GuidelineEntity {id: $target})
                        MERGE (s)-[:GUIDELINE_REL {type: $type}]->(t)
                        """, source=edge.source.id, target=edge.target.id, type=edge.type)
        driver.close()
        print("✅ Knowledge Graph Built Successfully.")
    except Exception as e:
        print(f"❌ KG Initialization Failed: {e}")

ingest_data()

# ==========================================
# 3. TOOLS
# ==========================================

@tool
def detect_stroke_type(image_path: str) -> str:
    """Analyzes a patient's CT scan using the remote Vision AI model."""
    HF_SPACE_URL = "Siyad798/stroke-diagnosis-ai" 
    print(f"--- 🧠 Calling Vision AI for {image_path} ---")
    if not os.path.exists(image_path): return json.dumps({"error": f"File not found locally: {image_path}"})
    try:
        client = Client(HF_SPACE_URL)
        result = client.predict(files=[handle_file(image_path)], api_name="/predict")
        ai_output = {}
        if isinstance(result, (list, tuple)) and len(result) > 1:
            json_data = result[1]
            if isinstance(json_data, dict): ai_output = json_data
            elif isinstance(json_data, str) and os.path.exists(json_data):
                with open(json_data, 'r') as f: ai_output = json.load(f)
            elif isinstance(json_data, str):
                try: ai_output = json.loads(json_data)
                except: ai_output = {"raw_output": json_data}
        else: ai_output = result
        if not ai_output: ai_output = {"error": "Empty response", "diagnosis": "Unknown"}
        return json.dumps(ai_output)
    except Exception as e: return json.dumps({"error": f"Inference Failed: {str(e)}", "diagnosis": "Unknown"})

@tool
def get_treatment_guideline(stroke_type: str) -> str:
    """Queries AHA/ASA Guidelines from Neo4j Knowledge Graph."""
    search_term = "Ischemic Stroke" if "Ischemic" in stroke_type else "Intracranial Hemorrhage"
    print(f"--- 🕸️ Querying KG for: {search_term} ---")
    if not NEO4J_URI or not NEO4J_PASSWORD: return json.dumps({"error": "Neo4j credentials not configured."})
    try:
        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
        with driver.session() as session:
            q1 = "MATCH (s:StrokeType {name: $st})-[r:INDICATED_FOR]->(t:Treatment) RETURN t.name as Treatment, t.class as Class, t.dosage as Dosage, r.criteria as Criteria"
            treatments = session.run(q1, st=search_term).data()
            q2 = "MATCH (c:ClinicalCondition)-[:CONTRAINDICATES]->(t:Treatment)<-[:INDICATED_FOR]-(s:StrokeType {name: $st}) RETURN t.name as Drug, c.name as Contraindication"
            safety = session.run(q2, st=search_term).data()
        driver.close()
        return json.dumps({"Guidelines_Source": "AHA/ASA 2019", "Recommended_Treatments": treatments, "Safety_Warnings": safety})
    except Exception as e: return json.dumps({"error": str(e)})

# ==========================================
# 4. AGENT WORKFLOW
# ==========================================

class AgentState(TypedDict):
    messages: Annotated[list, operator.add]
    image_path: str
    stroke_diagnosis: dict
    guidelines: dict
    final_report: str

def detector_node(state: AgentState):
    print("--- Node: Stroke Detector ---")
    result_json = detect_stroke_type.invoke(state["image_path"])
    diagnosis = json.loads(result_json)
    msg = f"Imaging Analysis Complete. Diagnosis: {diagnosis.get('diagnosis')}."
    return {"messages": [HumanMessage(content=msg)], "stroke_diagnosis": diagnosis}

def advisor_node(state: AgentState):
    print("--- Node: Guideline Advisor ---")
    diagnosis = state["stroke_diagnosis"].get("diagnosis", "Unknown")
    guidelines_json = get_treatment_guideline.invoke(diagnosis)
    guidelines = json.loads(guidelines_json)
    return {"messages": [HumanMessage(content=f"Guidelines Retrieved for {diagnosis}.")], "guidelines": guidelines}

def synthesizer_node(state: AgentState):
    print("--- Node: Report Synthesizer ---")
    prompt = f"""
    You are a Stroke Neurologist Assistant. Write a clinical report based on:
    1. **Imaging AI Result**: {json.dumps(state['stroke_diagnosis'], indent=2)}
    2. **Clinical Guidelines**: {json.dumps(state['guidelines'], indent=2)}
    3. **Patient Context**: {state['messages'][0].content}
    **Structure:**
    - Diagnosis & Severity
    - Imaging Findings
    - Recommended Interventions (Cite Evidence)
    - Safety Checks & Contraindications
    """
    response = llm.invoke(prompt)
    return {"final_report": response.content}

workflow = StateGraph(AgentState)
workflow.add_node("detector", detector_node)
workflow.add_node("advisor", advisor_node)
workflow.add_node("synthesizer", synthesizer_node)
workflow.set_entry_point("detector")
workflow.add_edge("detector", "advisor")
workflow.add_edge("advisor", "synthesizer")
workflow.add_edge("synthesizer", END)
app = workflow.compile()

# ==========================================
# 5. GRADIO UI
# ==========================================

def predict_stroke_report(uploaded_file, patient_context_input):
    if not uploaded_file: return "Please upload a CT scan file."
    image_path = uploaded_file.name
    inputs = {"messages": [HumanMessage(content=patient_context_input)], "image_path": image_path}
    final_report_content = "Processing..."
    try:
        for event in app.stream(inputs, stream_mode="values"): pass 
        final_state = event
        final_report_content = final_state.get("final_report", "Error generating report.")
    except Exception as e: final_report_content = f"Critical Error: {str(e)}"
    return final_report_content

# Safe initialization without arguments
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 AI Stroke Diagnosis & Guideline Agent")
    gr.Markdown("Integrates Vision AI (DenseNet) + Knowledge Graph (Neo4j) + LLM (Azure OpenAI)")
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload CT (ZIP/DCM)", file_types=[".zip", ".dcm"])
            patient_context = gr.Textbox(label="Patient Context", lines=4, value="Patient 55M, onset 2 hours ago. BP 150/90. No history of bleeding.")
            submit_btn = gr.Button("Generate Report", variant="primary")
        with gr.Column():
            report_output = gr.Markdown(label="Clinical Report")
    submit_btn.click(predict_stroke_report, inputs=[file_input, patient_context], outputs=report_output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)