import os

# --- Environment Setup ---
os.environ["HF_HOME"] = "/app/cache"
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/app/cache"
os.environ["CHROMA_DB_DIR"] = "/app/chroma_db"
os.environ["XDG_CONFIG_HOME"] = "/app/.streamlit"

import streamlit as st
st.set_page_config(layout="wide")  # ✅ Must come first before any other st.* calls

# --- Libraries ---
import networkx as nx
import spacy
import subprocess
from docx import Document
import pdfplumber
from pyvis.network import Network
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings
from transformers import pipeline
import tempfile
import streamlit.components.v1 as components

# --- Setup ---
@st.cache_resource
def ensure_spacy_model():
    nlp = spacy.load("en_core_web_sm")
    nlp.max_length = 2_000_000
    return nlp

nlp = ensure_spacy_model()
embedder = SentenceTransformer("all-MiniLM-L6-v2")

os.makedirs("/app/chroma_db", exist_ok=True)
client = chromadb.PersistentClient(path="/app/chroma_db")
collection = client.get_or_create_collection(name="graphrag")

llm = pipeline("text2text-generation", model="google/flan-t5-large")

# --- Utils ---
def extract_text(upload):
    if upload is None:
        return ""
    ext = os.path.splitext(upload.name)[1].lower()
    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
        tmp.write(upload.read())
        tmp_path = tmp.name

    if ext == ".txt":
        return open(tmp_path, "r", encoding="utf-8").read()
    elif ext == ".pdf":
        with pdfplumber.open(tmp_path) as pdf:
            return "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
    elif ext == ".docx":
        doc = Document(tmp_path)
        return "\n".join(p.text for p in doc.paragraphs)
    return ""

def build_graph(text):
    doc = nlp(text)
    G = nx.Graph()

    entity_data = [(ent.text.strip(), ent.label_) for ent in doc.ents if ent.text.strip()]
    unique_entities = list(set(entity_data))
    entity_texts = [ent[0] for ent in unique_entities]
    embeddings = embedder.encode(entity_texts, convert_to_numpy=True).tolist()

    for (text, label), emb in zip(unique_entities, embeddings):
        G.add_node(text, label=label, embedding=emb)

    unique_ids = [f"{ent}_{i}" for i, ent in enumerate(entity_texts)]
    collection.add(documents=entity_texts, embeddings=embeddings, ids=unique_ids)

    for sent in doc.sents:
        ents = [ent.text.strip() for ent in sent.ents if ent.text.strip() in G.nodes]
        for i in range(len(ents)):
            for j in range(i + 1, len(ents)):
                G.add_edge(ents[i], ents[j])

    return G

def retrieve_relevant_entities(query, top_k=5):
    query_emb = embedder.encode([query], convert_to_numpy=True).tolist()
    results = collection.query(query_embeddings=query_emb, n_results=top_k)
    return [match for match in results['ids'][0] if match in collection.peek()['ids']]

def expand_graph_context(graph, entities, depth=1):
    subgraph_nodes = set()
    for entity in entities:
        if entity in graph:
            neighbors = nx.single_source_shortest_path_length(graph, entity, cutoff=depth).keys()
            subgraph_nodes.update(neighbors)
    return graph.subgraph(subgraph_nodes)

def format_graph_context(subgraph):
    lines = []
    for node in subgraph.nodes(data=True):
        lines.append(f"Entity: {node[0]}, Type: {node[1].get('label', 'N/A')}")
    for edge in subgraph.edges():
        lines.append(f"Connection: {edge[0]} ↔ {edge[1]}")
    return "\n".join(lines)

def generate_answer(question, context_text):
    prompt = f"""You are a compliance assistant. Use the following knowledge graph context to answer the question.

{context_text}

Question: {question}
Answer:"""
    return llm(prompt, max_length=512, do_sample=False)[0]['generated_text']

def plot_graph(G, name="graph"):
    import shutil
    output_dir = "/app/tmp"
    lib_dir = os.path.join(output_dir, "lib")

    os.makedirs(output_dir, exist_ok=True)
    original_cwd = os.getcwd()
    os.chdir(output_dir)

    try:
        net = Network(
            height="600px",
            width="100%",
            bgcolor="#222222",
            font_color="white",
            cdn_resources="local"
        )
        for node, data in G.nodes(data=True):
            net.add_node(node, label=f"{node}\n[{data.get('label', '')}]")
        for src, dst in G.edges():
            net.add_edge(src, dst)

        net.repulsion(node_distance=150, spring_length=200)
        net.show_buttons(filter_=['physics'])

        html_path = f"{name}.html"
        net.save_graph(html_path)
        return os.path.join(output_dir, html_path)
    finally:
        os.chdir(original_cwd)

# --- Streamlit App ---
st.title("\U0001F4D8 GraphRAG Agent")

uploaded_file = st.file_uploader("Upload document", type=["pdf", "txt", "docx"])

if uploaded_file:
    with st.spinner("Processing document..."):
        text = extract_text(uploaded_file)
        G = build_graph(text)

        html_file = plot_graph(G)

        st.subheader("Entity Graph")
        with open(html_file, 'r', encoding='utf-8') as f:
            components.html(f.read(), height=650, scrolling=True)

        st.download_button("Download Graph HTML", data=open(html_file, 'rb'), file_name="graph.html")

        question = st.text_input("Ask a question (Graph-RAG):")
        if question:
            with st.spinner("Querying..."):
                relevant_entities = retrieve_relevant_entities(question)
                subgraph = expand_graph_context(G, relevant_entities)
                context_text = format_graph_context(subgraph)
                answer = generate_answer(question, context_text)
                st.markdown(f"**Answer:** {answer}")