#!/usr/bin/env python3
"""
RAG Cytoscape Streamlit Visualization

Usage:
    streamlit run rag_cytoscape_streamlit.py -- --query_name Alzheimer_CaseStudy
"""

import streamlit as st
import json
import os
import glob
import sys
import argparse

# ----------------------------
# Argument Parsing
# ----------------------------
parser = argparse.ArgumentParser()
parser.add_argument("--query_name", type=str, required=True, help="Query name / folder inside output/")
args, unknown = parser.parse_known_args()  # Streamlit ignores unknown args

query_name = args.query_name
output_dir = os.path.join("output", query_name)
if not os.path.exists(output_dir):
    st.error(f"Output folder for query_name '{query_name}' does not exist: {output_dir}")
    st.stop()

# ----------------------------
# Helper Function
# ----------------------------
def load_jsons_to_cy_elements(output_dir, node_types=None, search=None):
    json_files = glob.glob(os.path.join(output_dir, "*.json"))
    if not json_files:
        st.warning(f"No JSON files found in {output_dir}")
        return []

    nodes_dict = {}
    edges = []

    for jf in json_files:
        with open(jf, "r") as f:
            try:
                data = json.load(f)
            except Exception as e:
                st.warning(f"Failed to load {jf}: {e}")
                continue

        for entry in data:
            # Target node
            target_name = entry.get("target") or entry.get("Target") or entry.get("gene") or entry.get("entity")
            if not target_name or str(target_name).strip() == "":
                continue
            target_name = str(target_name).strip()
            if not node_types or "target" in node_types:
                nodes_dict[target_name] = {"data": {"id": target_name, "label": target_name, "type": "target"}}

            # Drug node
            drug_name = entry.get("drug") or entry.get("Drug")
            if drug_name and str(drug_name).strip() != "":
                drug_name = str(drug_name).strip()
                if not node_types or "drug" in node_types:
                    nodes_dict[drug_name] = {"data": {"id": drug_name, "label": drug_name, "type": "drug"}}
                    edges.append({"data": {"source": target_name, "target": drug_name, "label": "INTERACTS_WITH"}})

            # Disease node
            disease_field = entry.get("disease") or entry.get("cancer_association") or entry.get("Cancer")
            if disease_field and str(disease_field).strip() != "":
                for d in str(disease_field).split(","):
                    d = d.strip()
                    if d == "":
                        continue
                    if not node_types or "disease" in node_types:
                        nodes_dict[d] = {"data": {"id": d, "label": d, "type": "disease"}}
                        edges.append({"data": {"source": target_name, "target": d, "label": "ASSOCIATED_WITH"}})

            # PMID node
            pmid = entry.get("pmid")
            if pmid and str(pmid).strip() != "":
                pmid = str(pmid).strip()
                if not node_types or "pmid" in node_types:
                    nodes_dict[pmid] = {"data": {"id": pmid, "label": pmid, "type": "pmid"}}
                    edges.append({"data": {"source": target_name, "target": pmid, "label": "CITED_IN"}})

    # Optional search filter
    if search and search.strip() != "":
        search = search.lower()
        filtered_nodes = {k: v for k, v in nodes_dict.items() if search in k.lower()}
        filtered_ids = set(filtered_nodes.keys())
        filtered_edges = [e for e in edges if e["data"]["source"] in filtered_ids or e["data"]["target"] in filtered_ids]
        nodes_dict = filtered_nodes
        edges = filtered_edges

    return list(nodes_dict.values()) + edges

# ----------------------------
# Streamlit UI
# ----------------------------
st.title(f"RAG Network Visualization — {query_name}")

all_node_types = ["target", "drug", "disease", "pmid"]
selected_types = st.multiselect("Select node types to display:", options=all_node_types, default=all_node_types)
search_term = st.text_input("Search target, drug, or disease:", "")

if st.button("Load Graph from JSONs"):
    cy_elements = load_jsons_to_cy_elements(output_dir=output_dir, node_types=selected_types, search=search_term)
    st.success(f"Loaded {len(cy_elements)} nodes and edges from JSON files for '{query_name}'.")

    st.components.v1.html(f"""
    <html>
    <head>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/cytoscape/3.25.0/cytoscape.min.js"></script>
    </head>
    <body>
    <div style="width:100%; height:450px; border:1px solid #ccc; overflow:auto;">
        <div id="cy" style="width:700px; height:450px;"></div>
    </div>
    <script>
    var cy = cytoscape({{
      container: document.getElementById('cy'),
      elements: {json.dumps(cy_elements)},
      layout: {{ name: 'cose', animate: true }},
      style: [
        {{ selector: 'node[type="target"]', style: {{ 'background-color': '#61bffc', 'label': 'data(label)' }} }},
        {{ selector: 'node[type="drug"]', style: {{ 'background-color': '#ff6666', 'label': 'data(label)' }} }},
        {{ selector: 'node[type="disease"]', style: {{ 'background-color': '#ffcc66', 'label': 'data(label)' }} }},
        {{ selector: 'node[type="pmid"]', style: {{ 'background-color': '#99cc99', 'label': 'data(label)' }} }},
        {{ selector: 'edge', style: {{ 'width': 2, 'line-color': '#9dbaea', 'target-arrow-color': '#9dbaea', 'target-arrow-shape': 'triangle', 'curve-style': 'bezier' }} }}
      ]
    }});
    </script>
    </body>
    </html>
    """, height=650)
