app.py

#!/usr/bin/env python3
"""
AI Knowledge Graph Chat Application
====================================

A production-quality single-file AI chat application that builds a knowledge
graph from uploaded documents and answers questions using LangGraph orchestration.

Features:
  • Multi-provider LLM support (Groq, Gemini, Cohere, Cerebras)
  • Document upload & processing (PDF, DOCX, TXT, CSV, XLSX, JSON, MD, HTML)
  • LLM-driven knowledge-graph extraction (entities + relationships → triples)
  • Neo4j storage with automatic in-memory fallback
  • SHA-256-based triple deduplication and incremental updates
  • File deletion with surgical graph cleanup
  • Intent detection and intelligent routing via LangGraph
  • Streaming responses with stop/clear controls
  • Modern Gradio UI + REST API on a single port

Deployment:
    python app.py
"""

# ============================================================
# SECTION 1: IMPORTS
# ============================================================

import os
import re
import json
import time
import uuid
import asyncio
import hashlib
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import (
    Any, AsyncGenerator, Dict, List, Optional, Tuple, Union, Sequence
)

# --- Third-party core ---
import pandas as pd
from pydantic import BaseModel, Field

# --- FastAPI ---
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# --- Gradio ---
import gradio as gr

# --- LangChain core ---
from langchain_core.messages import (
    HumanMessage, AIMessage, SystemMessage, BaseMessage
)
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.outputs import ChatResult, ChatGeneration

# --- LangGraph ---
from langgraph.graph import StateGraph, END
from typing_extensions import TypedDict

# --- Optional LLM provider imports (gracefully degrade) ---
try:
    from langchain_groq import ChatGroq
    _GROQ_OK = True
except Exception:
    _GROQ_OK = False

try:
    from langchain_google_genai import ChatGoogleGenerativeAI
    _GEMINI_OK = True
except Exception:
    _GEMINI_OK = False

try:
    from langchain_cohere import ChatCohere
    _COHERE_OK = True
except Exception:
    _COHERE_OK = False

try:
    from langchain_cerebras import ChatCerebras
    _CEREBRAS_OK = True
except Exception:
    _CEREBRAS_OK = False

# --- Optional file-processing imports ---
try:
    import fitz  # PyMuPDF
    _PYMUPDF_OK = True
except Exception:
    _PYMUPDF_OK = False

try:
    from docx import Document as DocxDocument
    _DOCX_OK = True
except Exception:
    _DOCX_OK = False

try:
    from bs4 import BeautifulSoup
    _BS4_OK = True
except Exception:
    _BS4_OK = False

# --- Neo4j ---
try:
    from neo4j import GraphDatabase
    from neo4j.exceptions import ServiceUnavailable, AuthError, CypherSyntaxError
    _NEO4J_OK = True
except Exception:
    _NEO4J_OK = False


# ============================================================
# SECTION 2: CONFIGURATION & CONSTANTS
# ============================================================

class Config:
    """Central configuration. Override via environment variables."""

    # Server
    HOST: str = os.getenv("HOST", "0.0.0.0")
    PORT: int = int(os.getenv("PORT", "7860"))

    # Upload limits
    MAX_UPLOAD_SIZE_MB: int = int(os.getenv("MAX_UPLOAD_SIZE_MB", "50"))
    UPLOAD_DIR: Path = Path(os.getenv("UPLOAD_DIR", "./data/uploads"))
    UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

    # Neo4j
    NEO4J_URI: str = os.getenv("NEO4J_URI", "")
    NEO4J_USERNAME: str = os.getenv("NEO4J_USERNAME", "neo4j")
    NEO4J_PASSWORD: str = os.getenv("NEO4J_PASSWORD", "")

    # LLM defaults
    DEFAULT_PROVIDER: str = os.getenv("DEFAULT_PROVIDER", "groq")
    DEFAULT_MODELS: Dict[str, str] = {
        "groq":     os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile"),
        "gemini":   os.getenv("GEMINI_MODEL", "gemini-1.5-flash"),
        "cohere":   os.getenv("COHERE_MODEL", "command-r-plus"),
        "cerebras": os.getenv("CEREBRAS_MODEL", "llama-3.3-70b"),
    }

    # Extraction
    CHUNK_SIZE: int = int(os.getenv("CHUNK_SIZE", "4000"))
    MAX_TRIPLES_PER_CHUNK: int = int(os.getenv("MAX_TRIPLES_PER_CHUNK", "50"))

    # Supported file types
    SUPPORTED_EXTENSIONS: Tuple[str, ...] = (
        ".pdf", ".docx", ".txt", ".csv", ".xlsx", ".json", ".md", ".html"
    )


# Provider display metadata
PROVIDERS: Dict[str, Dict[str, Any]] = {
    "groq":     {"label": "Groq",     "available": _GROQ_OK,     "env_key": "GROQ_API_KEY"},
    "gemini":   {"label": "Gemini",   "available": _GEMINI_OK,   "env_key": "GOOGLE_API_KEY"},
    "cohere":   {"label": "Cohere",   "available": _COHERE_OK,   "env_key": "COHERE_API_KEY"},
    "cerebras": {"label": "Cerebras", "available": _CEREBRAS_OK, "env_key": "CEREBRAS_API_KEY"},
}

# Intent constants
INTENT_GENERAL_CHAT = "GENERAL_CHAT"
INTENT_KG_QUERY = "KNOWLEDGE_GRAPH_QUERY"
INTENT_DOC_SEARCH = "DOCUMENT_SEARCH"
INTENT_GREETING = "GREETING"
INTENT_PROGRAMMING = "PROGRAMMING"
INTENT_EXPLANATION = "EXPLANATION"

ALL_INTENTS = [
    INTENT_GENERAL_CHAT, INTENT_KG_QUERY, INTENT_DOC_SEARCH,
    INTENT_GREETING, INTENT_PROGRAMMING, INTENT_EXPLANATION,
]


# ============================================================
# SECTION 3: PYDANTIC MODELS
# ============================================================

class ChatRequest(BaseModel):
    """Request body for /chat endpoint."""
    message: str = Field(..., min_length=1, max_length=10000)
    provider: str = Field(default=Config.DEFAULT_PROVIDER)
    history: List[Dict[str, str]] = Field(default_factory=list)


class ChatResponse(BaseModel):
    """Response body for /chat endpoint."""
    reply: str
    intent: str
    execution_path: List[str]
    sources: List[str]


class UploadResponse(BaseModel):
    """Response body for /upload endpoint."""
    filename: str
    triples_inserted: int
    triples_skipped: int
    entities: int
    relationships: int
    processing_time: float
    message: str = ""


class FileMetadata(BaseModel):
    """Metadata for an uploaded file."""
    filename: str
    upload_date: str
    file_hash: str
    size_bytes: int
    triples: int


class GraphStats(BaseModel):
    """Knowledge-graph statistics."""
    documents: int
    entities: int
    relationships: int


class ProviderRequest(BaseModel):
    """Request body for /provider endpoint."""
    provider: str


class Triple(BaseModel):
    """A knowledge-graph triple."""
    subject: str
    relation: str
    object: str
    confidence: float = 0.8


# ============================================================
# SECTION 4: LOGGING SETUP
# ============================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("kg_app")


# ============================================================
# SECTION 5: LLM PROVIDER MANAGER
# ============================================================

class MockChatModel(BaseChatModel):
    """Fallback chat model used when no API key is configured."""

    def _generate(self, messages, stop=None, run_manager=None, **kwargs):
        last = messages[-1].content if messages else ""
        text = (
            f"[Mock LLM] No API key configured for the selected provider.\n\n"
            f"Your message was: {last[:200]}\n\n"
            f"Set the appropriate environment variable (e.g. GROQ_API_KEY) "
            f"to enable real LLM responses."
        )
        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=text))])

    @property
    def _llm_type(self) -> str:
        return "mock"


class LLMProviderManager:
    """
    Single abstraction over all supported LLM providers.

    Caches client instances so repeated calls reuse the same client.
    Switching providers is transparent to the rest of the application.
    """

    def __init__(self):
        self._cache: Dict[str, BaseChatModel] = {}

    def _normalise(self, provider: str) -> str:
        return provider.lower().strip()

    def is_available(self, provider: str) -> bool:
        """Check whether *provider* is installed AND has an API key."""
        p = self._normalise(provider)
        meta = PROVIDERS.get(p)
        if not meta or not meta["available"]:
            return False
        return bool(os.getenv(meta["env_key"]))

    def get_llm(self, provider: str, **kwargs) -> BaseChatModel:
        """
        Return a (possibly cached) chat model for *provider*.
        Falls back to MockChatModel when the real provider is unavailable.
        """
        p = self._normalise(provider)
        model_name = kwargs.get("model", Config.DEFAULT_MODELS.get(p, ""))
        cache_key = f"{p}::{model_name}"

        if cache_key in self._cache:
            return self._cache[cache_key]

        llm = self._create(p, model_name=model_name, **kwargs)
        self._cache[cache_key] = llm
        logger.info("LLM provider initialised: %s (model=%s, type=%s)",
                    p, model_name, type(llm).__name__)
        return llm

    def _create(self, provider: str, *, model_name: str, **kwargs) -> BaseChatModel:
        temperature = kwargs.get("temperature", 0.7)

        if provider == "groq" and _GROQ_OK:
            key = os.getenv("GROQ_API_KEY")
            if key:
                return ChatGroq(model=model_name, temperature=temperature, api_key=key)

        if provider == "gemini" and _GEMINI_OK:
            key = os.getenv("GOOGLE_API_KEY")
            if key:
                return ChatGoogleGenerativeAI(
                    model=model_name, temperature=temperature, google_api_key=key
                )

        if provider == "cohere" and _COHERE_OK:
            key = os.getenv("COHERE_API_KEY")
            if key:
                return ChatCohere(model=model_name, temperature=temperature, cohere_api_key=key)

        if provider == "cerebras" and _CEREBRAS_OK:
            key = os.getenv("CEREBRAS_API_KEY")
            if key:
                return ChatCerebras(model=model_name, temperature=temperature, api_key=key)

        logger.warning("Provider '%s' unavailable – using MockChatModel", provider)
        return MockChatModel()

    def list_providers(self) -> List[Dict[str, Any]]:
        """Return provider info for UI rendering."""
        result = []
        for key, meta in PROVIDERS.items():
            result.append({
                "key": key,
                "label": meta["label"],
                "available": self.is_available(key),
            })
        return result


# Singleton
provider_manager = LLMProviderManager()


# ============================================================
# SECTION 6: FILE PROCESSORS
# ============================================================

class FileProcessor:
    """Detect file type and extract clean text."""

    @staticmethod
    def detect_type(filepath: str) -> str:
        ext = Path(filepath).suffix.lower()
        if ext not in Config.SUPPORTED_EXTENSIONS:
            raise ValueError(f"Unsupported file type: {ext}")
        return ext

    @staticmethod
    def extract_text(filepath: str) -> str:
        """Dispatch to the correct extractor based on file extension."""
        ext = FileProcessor.detect_type(filepath)
        extractors = {
            ".pdf":   FileProcessor._extract_pdf,
            ".docx":  FileProcessor._extract_docx,
            ".txt":   FileProcessor._extract_text,
            ".csv":   FileProcessor._extract_csv,
            ".xlsx":  FileProcessor._extract_xlsx,
            ".json":  FileProcessor._extract_json,
            ".md":    FileProcessor._extract_text,
            ".html":  FileProcessor._extract_html,
        }
        extractor = extractors.get(ext, FileProcessor._extract_text)
        text = extractor(filepath)
        # Clean whitespace
        text = re.sub(r"[ \t]+", " ", text)
        text = re.sub(r"\n{3,}", "\n\n", text).strip()
        if not text:
            raise ValueError("No readable text found in file.")
        return text

    @staticmethod
    def _extract_pdf(filepath: str) -> str:
        if not _PYMUPDF_OK:
            raise RuntimeError("PyMuPDF not installed – cannot process PDF.")
        doc = fitz.open(filepath)
        pages = [page.get_text("text") for page in doc]
        doc.close()
        return "\n\n".join(pages)

    @staticmethod
    def _extract_docx(filepath: str) -> str:
        if not _DOCX_OK:
            raise RuntimeError("python-docx not installed – cannot process DOCX.")
        doc = DocxDocument(filepath)
        return "\n".join(p.text for p in doc.paragraphs if p.text.strip())

    @staticmethod
    def _extract_text(filepath: str) -> str:
        with open(filepath, "r", encoding="utf-8", errors="replace") as f:
            return f.read()

    @staticmethod
    def _extract_csv(filepath: str) -> str:
        df = pd.read_csv(filepath)
        return df.to_string(index=False)

    @staticmethod
    def _extract_xlsx(filepath: str) -> str:
        xl = pd.ExcelFile(filepath, engine="openpyxl")
        parts = []
        for sheet in xl.sheet_names:
            df = xl.parse(sheet)
            parts.append(f"## Sheet: {sheet}\n{df.to_string(index=False)}")
        return "\n\n".join(parts)

    @staticmethod
    def _extract_json(filepath: str) -> str:
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        return json.dumps(data, indent=2, ensure_ascii=False)

    @staticmethod
    def _extract_html(filepath: str) -> str:
        with open(filepath, "r", encoding="utf-8", errors="replace") as f:
            raw = f.read()
        if _BS4_OK:
            soup = BeautifulSoup(raw, "html.parser")
            # Remove script/style
            for tag in soup(["script", "style"]):
                tag.decompose()
            return soup.get_text(separator="\n")
        # Crude fallback
        return re.sub(r"<[^>]+>", " ", raw)


# ============================================================
# SECTION 7: KNOWLEDGE GRAPH EXTRACTOR
# ============================================================

EXTRACTION_SYSTEM_PROMPT = """\
You are a knowledge-graph extraction engine.
Extract entities and relationships from the user-provided text.

Return STRICT JSON with this schema:
{{
  "entities": [
    {{"name": "canonical entity name", "type": "Person|Organization|Technology|Location|Event|Concept|Date|Other"}}
  ],
  "relationships": [
    {{"subject": "entity name", "relation": "lowercase_snake_case", "object": "entity name", "confidence": 0.0-1.0}}
  ]
}}

Rules:
• Entity names MUST be normalised (canonical form, Title Case where appropriate).
• Relations MUST be lowercase snake_case verbs or short phrases (e.g. "works_for", "located_in").
• Confidence is a float between 0 and 1.
• Extract ONLY clear, factual relationships stated in the text.
• Do NOT invent information.
• Return at most {max_triples} relationships.
• If no knowledge can be extracted, return {{"entities": [], "relationships": []}}.
"""

CYPHER_SYSTEM_PROMPT = """\
You are a Cypher query generator for a Neo4j knowledge graph.

Graph schema:
  • Nodes:    (:Entity {{name: string, type: string}})
  • Relationships: (:Entity)-[:RELATES_TO {{relation: string}}]->(:Entity)

Instructions:
  1. Generate a Cypher query that retrieves information relevant to the user's question.
  2. Use fuzzy matching where helpful: `toLower(n.name) CONTAINS toLower(keyword)`.
  3. Limit results to 20 rows.
  4. Return ONLY the Cypher query — no markdown, no explanation.

Example:
  Question: "Who works at Acme?"
  Query: MATCH (s:Entity)-[r:RELATES_TO]->(o:Entity)
         WHERE toLower(r.relation) CONTAINS 'work' AND toLower(o.name) CONTAINS 'acme'
         RETURN s.name, r.relation, o.name LIMIT 20
"""


class KnowledgeExtractor:
    """Use an LLM to extract structured knowledge from raw text."""

    def __init__(self, provider_manager: LLMProviderManager):
        self._pm = provider_manager

    def _chunk_text(self, text: str, chunk_size: int = Config.CHUNK_SIZE) -> List[str]:
        """Split text into character chunks (NOT for RAG retrieval — for extraction)."""
        return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

    def _parse_json_response(self, content: str) -> Dict[str, Any]:
        """Robustly extract JSON from an LLM response that may contain markdown fences."""
        # Strip markdown code fences
        content = re.sub(r"```(?:json)?\s*", "", content)
        content = content.strip().rstrip("`")

        # Try direct parse
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        # Try to find first { ... } block
        match = re.search(r"\{.*\}", content, re.DOTALL)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                pass

        logger.warning("Failed to parse LLM JSON response. Returning empty.")
        return {"entities": [], "relationships": []}

    def extract(self, text: str, provider: str) -> Tuple[List[Dict], List[Dict]]:
        """
        Extract entities and relationships from *text*.
        Returns (entities, relationships) where each relationship is a triple dict.
        """
        llm = self._pm.get_llm(provider, temperature=0.1)
        chunks = self._chunk_text(text)

        all_entities: List[Dict] = []
        all_rels: List[Dict] = []

        for idx, chunk in enumerate(chunks):
            logger.info("Extracting knowledge from chunk %d/%d (%d chars)",
                        idx + 1, len(chunks), len(chunk))

            messages = [
                SystemMessage(content=EXTRACTION_SYSTEM_PROMPT.format(
                    max_triples=Config.MAX_TRIPLES_PER_CHUNK
                )),
                HumanMessage(content=f"Extract knowledge from this text:\n\n{chunk}"),
            ]

            try:
                response = llm.invoke(messages)
                data = self._parse_json_response(response.content)
                all_entities.extend(data.get("entities", []))
                all_rels.extend(data.get("relationships", []))
            except Exception as e:
                logger.error("Extraction failed on chunk %d: %s", idx + 1, e)

        # Deduplicate entities by name (case-insensitive)
        seen_names: set = set()
        unique_entities: List[Dict] = []
        for ent in all_entities:
            key = ent["name"].lower().strip()
            if key and key not in seen_names:
                seen_names.add(key)
                unique_entities.append(ent)

        # Normalise relationships
        normalised_rels: List[Dict] = []
        for rel in all_rels:
            s = str(rel.get("subject", "")).strip()
            r = str(rel.get("relation", "")).strip().lower().replace(" ", "_")
            o = str(rel.get("object", "")).strip()
            conf = float(rel.get("confidence", 0.8))
            if s and r and o:
                normalised_rels.append({
                    "subject": s, "relation": r, "object": o,
                    "confidence": min(max(conf, 0.0), 1.0),
                })

        return unique_entities, normalised_rels


# ============================================================
# SECTION 8: GRAPH STORE  (Neo4j with in-memory fallback)
# ============================================================

class GraphStoreBase:
    """Abstract interface for graph storage backends."""

    def register_document(self, file_id: str, filename: str,
                          file_hash: str, upload_time: str) -> None: ...

    def add_triple(self, subject: str, relation: str, object_: str,
                   file_id: str, filename: str, upload_time: str,
                   version: int, confidence: float) -> bool: ...

    def delete_by_file(self, file_id: str) -> int: ...

    def get_stats(self) -> GraphStats: ...

    def search(self, query: str, limit: int = 20) -> List[Dict]: ...

    def get_file_triple_count(self, file_id: str) -> int: ...

    def close(self) -> None: ...


# --- In-memory backend ----------------------------------------------------

class InMemoryGraphStore(GraphStoreBase):
    """
    Pure-Python graph store used when Neo4j is not configured.
    Implements the same interface as the Neo4j backend.
    """

    def __init__(self):
        self._documents: Dict[str, Dict] = {}   # file_id -> doc meta
        self._nodes: Dict[str, Dict] = {}       # normalised name -> node
        self._rels: Dict[str, Dict] = {}        # triple_hash -> relationship

    @staticmethod
    def _triple_hash(subject: str, relation: str, object_: str) -> str:
        raw = f"{subject.lower().strip()}|{relation.lower().strip()}|{object_.lower().strip()}"
        return hashlib.sha256(raw.encode("utf-8")).hexdigest()

    def register_document(self, file_id, filename, file_hash, upload_time):
        self._documents[file_id] = {
            "filename": filename,
            "file_hash": file_hash,
            "upload_time": upload_time,
        }

    def add_triple(self, subject, relation, object_, file_id, filename,
                   upload_time, version, confidence):
        h = self._triple_hash(subject, relation, object_)

        if h in self._rels:
            # Triple exists — just add source if not already recorded
            rel = self._rels[h]
            if file_id not in rel["source_files"]:
                rel["source_files"].append(file_id)
                rel["filenames"].append(filename)
                rel["upload_times"].append(upload_time)
                rel["versions"].append(version)
                rel["confidences"].append(confidence)
            return False  # skipped (already existed)

        # Merge nodes
        for name in (subject, object_):
            key = name.lower().strip()
            if key not in self._nodes:
                self._nodes[key] = {
                    "name": name,
                    "source_files": [],
                    "filenames": [],
                    "upload_times": [],
                    "versions": [],
                    "confidences": [],
                    "created_at": datetime.now(timezone.utc).isoformat(),
                }
            node = self._nodes[key]
            if file_id not in node["source_files"]:
                node["source_files"].append(file_id)
                node["filenames"].append(filename)
                node["upload_times"].append(upload_time)
                node["versions"].append(version)
                node["confidences"].append(confidence)

        self._rels[h] = {
            "subject": subject,
            "relation": relation,
            "object": object_,
            "hash": h,
            "source_files": [file_id],
            "filenames": [filename],
            "upload_times": [upload_time],
            "versions": [version],
            "confidences": [confidence],
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
        return True  # inserted

    def delete_by_file(self, file_id):
        deleted = 0

        # Delete relationships
        to_del_rels = []
        for h, rel in self._rels.items():
            if file_id in rel["source_files"]:
                idx = rel["source_files"].index(file_id)
                rel["source_files"].pop(idx)
                rel["filenames"].pop(idx)
                rel["upload_times"].pop(idx)
                rel["versions"].pop(idx)
                rel["confidences"].pop(idx)
                if not rel["source_files"]:
                    to_del_rels.append(h)
                    deleted += 1
        for h in to_del_rels:
            del self._rels[h]

        # Delete nodes whose source_files no longer include any file
        to_del_nodes = []
        for key, node in self._nodes.items():
            if file_id in node["source_files"]:
                idx = node["source_files"].index(file_id)
                node["source_files"].pop(idx)
                node["filenames"].pop(idx)
                node["upload_times"].pop(idx)
                node["versions"].pop(idx)
                node["confidences"].pop(idx)
                if not node["source_files"]:
                    to_del_nodes.append(key)
        for key in to_del_nodes:
            del self._nodes[key]

        # Remove document registration
        self._documents.pop(file_id, None)
        return deleted

    def get_stats(self):
        return GraphStats(
            documents=len(self._documents),
            entities=len(self._nodes),
            relationships=len(self._rels),
        )

    def search(self, query, limit=20):
        """Keyword-based search over triples."""
        keywords = [w.lower().strip() for w in re.split(r"\s+", query) if len(w) > 2]
        results = []
        for rel in self._rels.values():
            text = f"{rel['subject']} {rel['relation']} {rel['object']}".lower()
            score = sum(1 for kw in keywords if kw in text)
            if score > 0:
                results.append({
                    "subject": rel["subject"],
                    "relation": rel["relation"],
                    "object": rel["object"],
                    "sources": rel["filenames"],
                    "score": score,
                })
        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:limit]

    def get_file_triple_count(self, file_id):
        return sum(
            1 for rel in self._rels.values()
            if file_id in rel["source_files"]
        )

    def close(self):
        pass


# --- Neo4j backend --------------------------------------------------------

class Neo4jGraphStore(GraphStoreBase):
    """Neo4j-backed graph store. Falls back gracefully on connection errors."""

    def __init__(self, uri: str, username: str, password: str):
        self._driver = GraphDatabase.driver(uri, auth=(username, password))
        # Verify connectivity
        self._driver.verify_connectivity()
        logger.info("Connected to Neo4j at %s", uri)
        self._init_constraints()

    def _init_constraints(self):
        """Create uniqueness constraint on Entity.name if not exists."""
        with self._driver.session() as session:
            try:
                session.run(
                    "CREATE CONSTRAINT entity_name_unique IF NOT EXISTS "
                    "FOR (n:Entity) REQUIRE n.name IS UNIQUE"
                )
            except Exception as e:
                logger.warning("Could not create constraint: %s", e)

    @staticmethod
    def _triple_hash(subject, relation, object_):
        raw = f"{subject.lower().strip()}|{relation.lower().strip()}|{object_.lower().strip()}"
        return hashlib.sha256(raw.encode("utf-8")).hexdigest()

    def register_document(self, file_id, filename, file_hash, upload_time):
        with self._driver.session() as session:
            session.run(
                "MERGE (d:Document {file_id: $fid}) "
                "SET d.filename = $fn, d.file_hash = $fh, d.upload_time = $ut",
                fid=file_id, fn=filename, fh=file_hash, ut=upload_time,
            )

    def add_triple(self, subject, relation, object_, file_id, filename,
                   upload_time, version, confidence):
        h = self._triple_hash(subject, relation, object_)
        with self._driver.session() as session:
            # Check if triple already exists
            result = session.run(
                "MATCH ()-[r:RELATES_TO {hash: $h}]->() "
                "RETURN r.source_files AS sfs",
                h=h,
            )
            record = result.single()

            if record is not None:
                # Triple exists — add source file if not present
                sfs = record["sfs"] or []
                if file_id not in sfs:
                    session.run(
                        "MATCH ()-[r:RELATES_TO {hash: $h}]->() "
                        "SET r.source_files = coalesce(r.source_files, []) + $fid, "
                        "    r.filenames = coalesce(r.filenames, []) + $fn, "
                        "    r.upload_times = coalesce(r.upload_times, []) + $ut, "
                        "    r.versions = coalesce(r.versions, []) + $ver, "
                        "    r.confidences = coalesce(r.confidences, []) + $conf",
                        h=h, fid=file_id, fn=filename, ut=upload_time,
                        ver=version, conf=confidence,
                    )
                return False  # skipped

            # Create new triple
            session.run(
                """
                MERGE (s:Entity {name: $subject})
                MERGE (o:Entity {name: $object})
                CREATE (s)-[r:RELATES_TO {hash: $h}]->(o)
                SET r.relation = $rel,
                    r.source_files = [$fid],
                    r.filenames = [$fn],
                    r.upload_times = [$ut],
                    r.versions = [$ver],
                    r.confidences = [$conf],
                    r.created_at = $now
                WITH s, o
                WHERE NOT $fid IN coalesce(s.source_files, [])
                SET s.source_files = coalesce(s.source_files, []) + $fid,
                    s.filenames = coalesce(s.filenames, []) + $fn,
                    s.upload_times = coalesce(s.upload_times, []) + $ut,
                    s.versions = coalesce(s.versions, []) + $ver,
                    s.confidences = coalesce(s.confidences, []) + $conf,
                    s.created_at = coalesce(s.created_at, $now)
                WITH o
                WHERE NOT $fid IN coalesce(o.source_files, [])
                SET o.source_files = coalesce(o.source_files, []) + $fid,
                    o.filenames = coalesce(o.filenames, []) + $fn,
                    o.upload_times = coalesce(o.upload_times, []) + $ut,
                    o.versions = coalesce(o.versions, []) + $ver,
                    o.confidences = coalesce(o.confidences, []) + $conf,
                    o.created_at = coalesce(o.created_at, $now)
                """,
                subject=subject, object_=object_, h=h, rel=relation,
                fid=file_id, fn=filename, ut=upload_time,
                ver=version, conf=confidence,
                now=datetime.now(timezone.utc).isoformat(),
            )
            return True  # inserted

    def delete_by_file(self, file_id):
        deleted = 0
        with self._driver.session() as session:
            # Count and remove file from relationships, delete orphaned rels
            result = session.run(
                """
                MATCH ()-[r:RELATES_TO]-()
                WHERE $fid IN r.source_files
                WITH r, r.source_files AS sfs
                SET r.source_files = [x IN sfs WHERE x <> $fid]
                WITH r WHERE size(r.source_files) = 0
                DELETE r
                RETURN count(*) AS cnt
                """,
                fid=file_id,
            )
            rec = result.single()
            deleted = rec["cnt"] if rec else 0

            # Remove file from nodes, delete orphaned nodes
            session.run(
                """
                MATCH (n:Entity)
                WHERE $fid IN n.source_files
                SET n.source_files = [x IN n.source_files WHERE x <> $fid]
                WITH n WHERE size(n.source_files) = 0
                DETACH DELETE n
                """,
                fid=file_id,
            )

            # Remove document node
            session.run("MATCH (d:Document {file_id: $fid}) DELETE d", fid=file_id)
        return deleted

    def get_stats(self):
        with self._driver.session() as session:
            docs = session.run("MATCH (d:Document) RETURN count(d) AS c").single()["c"]
            ents = session.run("MATCH (n:Entity) RETURN count(n) AS c").single()["c"]
            rels = session.run("MATCH ()-[r:RELATES_TO]->() RETURN count(r) AS c").single()["c"]
        return GraphStats(documents=docs, entities=ents, relationships=rels)

    def search(self, query, limit=20):
        """Keyword search over triples."""
        keywords = [w.lower().strip() for w in re.split(r"\s+", query) if len(w) > 2]
        if not keywords:
            return []
        conditions = " OR ".join(
            [f"toLower(s.name) CONTAINS '{kw}' OR toLower(r.relation) CONTAINS '{kw}' OR toLower(o.name) CONTAINS '{kw}'"
             for kw in keywords]
        )
        cypher = (
            f"MATCH (s:Entity)-[r:RELATES_TO]->(o:Entity) "
            f"WHERE {conditions} "
            f"RETURN s.name AS subject, r.relation AS relation, o.name AS object, "
            f"r.filenames AS sources "
            f"LIMIT {limit}"
        )
        with self._driver.session() as session:
            result = session.run(cypher)
            return [dict(r) for r in result]

    def execute_cypher(self, cypher: str, limit: int = 20) -> List[Dict]:
        """Execute a raw Cypher query (used by LLM-generated queries)."""
        # Safety: only allow READ queries
        stripped = cypher.strip().upper()
        if not stripped.startswith("MATCH") and not stripped.startswith("RETURN"):
            raise ValueError("Only MATCH/RETURN queries are allowed.")
        if "DELETE" in stripped or "REMOVE" in stripped or "DROP" in stripped:
            raise ValueError("Destructive queries are not allowed.")
        with self._driver.session() as session:
            result = session.run(cypher)
            return [dict(r) for r in result]

    def get_file_triple_count(self, file_id):
        with self._driver.session() as session:
            result = session.run(
                "MATCH ()-[r:RELATES_TO]-() WHERE $fid IN r.source_files "
                "RETURN count(r) AS c",
                fid=file_id,
            )
            return result.single()["c"]

    def close(self):
        self._driver.close()


# --- Factory ---------------------------------------------------------------

def create_graph_store() -> GraphStoreBase:
    """Create the best available graph store."""
    if _NEO4J_OK and Config.NEO4J_URI:
        try:
            return Neo4jGraphStore(
                Config.NEO4J_URI, Config.NEO4J_USERNAME, Config.NEO4J_PASSWORD
            )
        except (ServiceUnavailable, AuthError, Exception) as e:
            logger.warning("Neo4j connection failed (%s) — falling back to in-memory.", e)
    logger.info("Using in-memory graph store.")
    return InMemoryGraphStore()


# ============================================================
# SECTION 9: CONVERSATION MEMORY
# ============================================================

class ConversationMemory:
    """
    Maintains per-session conversation history.
    This is SEPARATE from the knowledge graph — it is never persisted to Neo4j.
    """

    def __init__(self, max_messages: int = 50):
        self._sessions: Dict[str, List[Dict[str, str]]] = {}
        self._max = max_messages

    def get_history(self, session_id: str) -> List[Dict[str, str]]:
        return self._sessions.get(session_id, [])

    def add_message(self, session_id: str, role: str, content: str):
        hist = self._sessions.setdefault(session_id, [])
        hist.append({"role": role, "content": content})
        if len(hist) > self._max:
            self._sessions[session_id] = hist[-self._max:]

    def clear(self, session_id: str):
        self._sessions.pop(session_id, None)

    def to_langchain_messages(self, session_id: str) -> List[BaseMessage]:
        """Convert stored history to LangChain message objects."""
        msgs: List[BaseMessage] = []
        for m in self.get_history(session_id):
            if m["role"] == "user":
                msgs.append(HumanMessage(content=m["content"]))
            elif m["role"] == "assistant":
                msgs.append(AIMessage(content=m["content"]))
        return msgs


# ============================================================
# SECTION 10: INTENT DETECTION
# ============================================================

INTENT_SYSTEM_PROMPT = """\
You are an intent classifier for an AI knowledge-graph assistant.
Classify the user's message into EXACTLY one of these intents:

  • GENERAL_CHAT          — casual conversation, opinions, general questions
  • KNOWLEDGE_GRAPH_QUERY — questions that require information stored in the knowledge graph
  • DOCUMENT_SEARCH       — questions about uploaded documents
  • GREETING              — hello, hi, greetings
  • PROGRAMMING           — code, programming, technical implementation
  • EXPLANATION           — explain a concept, how something works

Return ONLY the intent name (one of the above), nothing else.
"""


class IntentDetector:
    """Lightweight LLM-based intent classifier."""

    def __init__(self, provider_manager: LLMProviderManager):
        self._pm = provider_manager

    def detect(self, message: str, provider: str) -> str:
        """Return one of the ALL_INTENTS constants."""
        llm = self._pm.get_llm(provider, temperature=0.0)
        try:
            response = llm.invoke([
                SystemMessage(content=INTENT_SYSTEM_PROMPT),
                HumanMessage(content=message),
            ])
            intent = response.content.strip().upper()
            # Validate
            for valid in ALL_INTENTS:
                if valid in intent:
                    return valid
        except Exception as e:
            logger.error("Intent detection failed: %s", e)
        return INTENT_GENERAL_CHAT  # safe fallback


# ============================================================
# SECTION 11: LANGGRAPH WORKFLOW
# ============================================================

class AgentState(TypedDict, total=False):
    """State object passed through the LangGraph workflow."""
    user_input: str
    provider: str
    intent: str
    context: str
    sources: List[str]
    execution_path: List[str]


class KnowledgeGraphWorkflow:
    """
    LangGraph-based orchestration:
      START → detect_intent → (route) → general_chat | kg_search → END
    """

    def __init__(
        self,
        intent_detector: IntentDetector,
        graph_store: GraphStoreBase,
        provider_manager: LLMProviderManager,
    ):
        self._intent = intent_detector
        self._store = graph_store
        self._pm = provider_manager
        self._graph = self._build()

    def _build(self):
        workflow = StateGraph(AgentState)

        workflow.add_node("detect_intent", self._detect_intent_node)
        workflow.add_node("general_chat", self._general_chat_node)
        workflow.add_node("kg_search", self._kg_search_node)

        workflow.set_entry_point("detect_intent")
        workflow.add_conditional_edges(
            "detect_intent",
            self._route,
            {
                "general_chat": "general_chat",
                "kg_search": "kg_search",
            },
        )
        workflow.add_edge("general_chat", END)
        workflow.add_edge("kg_search", END)

        return workflow.compile()

    # --- Nodes ---

    def _detect_intent_node(self, state: AgentState) -> AgentState:
        intent = self._intent.detect(state["user_input"], state["provider"])
        state["intent"] = intent
        state["execution_path"] = state.get("execution_path", []) + ["intent_detection"]
        logger.info("Intent detected: %s", intent)
        return state

    def _general_chat_node(self, state: AgentState) -> AgentState:
        state["context"] = ""
        state["sources"] = []
        state["execution_path"] = state.get("execution_path", []) + ["general_chat"]
        return state

    def _kg_search_node(self, state: AgentState) -> AgentState:
        """Search the knowledge graph and assemble context for the LLM."""
        state["execution_path"] = state.get("execution_path", []) + ["kg_search"]

        results = self._store.search(state["user_input"], limit=20)
        if not results:
            state["context"] = "No relevant information found in the knowledge graph."
            state["sources"] = []
            return state

        # Build context text
        lines = []
        sources_set = set()
        for r in results:
            lines.append(f"• {r['subject']} —[{r['relation']}]-> {r['object']}")
            for s in r.get("sources", []):
                sources_set.add(s)

        state["context"] = "\n".join(lines)
        state["sources"] = sorted(sources_set)
        return state

    # --- Routing ---

    def _route(self, state: AgentState) -> str:
        intent = state.get("intent", INTENT_GENERAL_CHAT)
        if intent in (INTENT_KG_QUERY, INTENT_DOC_SEARCH):
            return "kg_search"
        return "general_chat"

    # --- Public API ---

    def run(self, user_input: str, provider: str) -> AgentState:
        """Execute the workflow and return the final state."""
        initial: AgentState = {
            "user_input": user_input,
            "provider": provider,
            "intent": "",
            "context": "",
            "sources": [],
            "execution_path": [],
        }
        return self._graph.invoke(initial)


# ============================================================
# SECTION 12: APPLICATION ORCHESTRATOR
# ============================================================

class Application:
    """
    Central orchestrator that ties together all subsystems:
    file processing, knowledge extraction, graph storage,
    conversation memory, and the LangGraph workflow.
    """

    def __init__(self):
        self.graph_store: GraphStoreBase = create_graph_store()
        self.provider_mgr: LLMProviderManager = provider_manager
        self.extractor: KnowledgeExtractor = KnowledgeExtractor(self.provider_mgr)
        self.intent_detector: IntentDetector = IntentDetector(self.provider_mgr)
        self.workflow: KnowledgeGraphWorkflow = KnowledgeGraphWorkflow(
            self.intent_detector, self.graph_store, self.provider_mgr
        )
        self.memory: ConversationMemory = ConversationMemory()
        self.file_registry: Dict[str, Dict] = {}  # file_id -> metadata
        self._stop_flags: Dict[str, bool] = {}

    # --- File helpers ---

    @staticmethod
    def _file_hash(filepath: str) -> str:
        h = hashlib.sha256()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                h.update(chunk)
        return h.hexdigest()

    def _find_existing_file(self, filename: str, file_hash: str) -> Optional[str]:
        """Check if a file with same name AND hash is already indexed."""
        for fid, meta in self.file_registry.items():
            if meta["filename"] == filename and meta["file_hash"] == file_hash:
                return fid
        return None

    def upload_file(self, filepath: str, original_name: str,
                    provider: str) -> UploadResponse:
        """Process an uploaded file: extract text → extract knowledge → store triples."""
        start = time.time()

        # Validate
        ext = Path(original_name).suffix.lower()
        if ext not in Config.SUPPORTED_EXTENSIONS:
            raise ValueError(f"Unsupported file type: {ext}")

        size = os.path.getsize(filepath)
        if size > Config.MAX_UPLOAD_SIZE_MB * 1024 * 1024:
            raise ValueError(
                f"File too large ({size / 1024 / 1024:.1f} MB). "
                f"Max: {Config.MAX_UPLOAD_SIZE_MB} MB."
            )

        file_hash = self._file_hash(filepath)

        # Duplicate detection: same filename + same hash
        existing = self._find_existing_file(original_name, file_hash)
        if existing:
            return UploadResponse(
                filename=original_name,
                triples_inserted=0,
                triples_skipped=0,
                entities=0,
                relationships=0,
                processing_time=0.0,
                message="File already indexed.",
            )

        # Determine version (if same filename but different hash → new version)
        version = 1
        for fid, meta in self.file_registry.items():
            if meta["filename"] == original_name:
                version = max(version, meta["version"] + 1)

        # Extract text
        try:
            text = FileProcessor.extract_text(filepath)
        except Exception as e:
            logger.error("Text extraction failed for %s: %s", original_name, e)
            raise RuntimeError(f"Text extraction failed: {e}")

        # Extract knowledge via LLM
        try:
            entities, relationships = self.extractor.extract(text, provider)
        except Exception as e:
            logger.error("Knowledge extraction failed for %s: %s", original_name, e)
            raise RuntimeError(f"Knowledge extraction failed: {e}")

        # Register document in graph
        file_id = str(uuid.uuid4())
        upload_time = datetime.now(timezone.utc).isoformat()

        self.graph_store.register_document(file_id, original_name, file_hash, upload_time)

        # Insert triples (with deduplication)
        inserted = 0
        skipped = 0
        for rel in relationships:
            try:
                was_inserted = self.graph_store.add_triple(
                    subject=rel["subject"],
                    relation=rel["relation"],
                    object_=rel["object"],
                    file_id=file_id,
                    filename=original_name,
                    upload_time=upload_time,
                    version=version,
                    confidence=rel["confidence"],
                )
                if was_inserted:
                    inserted += 1
                else:
                    skipped += 1
            except Exception as e:
                logger.error("Triple insert failed: %s", e)
                skipped += 1

        # Register in local registry
        self.file_registry[file_id] = {
            "filename": original_name,
            "file_hash": file_hash,
            "upload_time": upload_time,
            "size_bytes": size,
            "version": version,
            "file_id": file_id,
        }

        elapsed = time.time() - start
        stats = self.graph_store.get_stats()

        logger.info(
            "Upload complete: %s | inserted=%d skipped=%d entities=%d rels=%d time=%.2fs",
            original_name, inserted, skipped, stats.entities, stats.relationships, elapsed
        )

        return UploadResponse(
            filename=original_name,
            triples_inserted=inserted,
            triples_skipped=skipped,
            entities=stats.entities,
            relationships=stats.relationships,
            processing_time=round(elapsed, 2),
            message=f"Successfully processed '{original_name}' (v{version}).",
        )

    def delete_file(self, filename: str) -> Dict[str, Any]:
        """Delete a file and all graph elements that originated ONLY from it."""
        # Find file by filename (use latest version if multiple)
        target_id = None
        for fid, meta in self.file_registry.items():
            if meta["filename"] == filename:
                target_id = fid  # keep searching to get latest

        if not target_id:
            return {"success": False, "message": f"File '{filename}' not found."}

        deleted_triples = self.graph_store.delete_by_file(target_id)
        del self.file_registry[target_id]

        # Optionally remove the physical file
        # (not strictly necessary since we use temp paths)

        stats = self.graph_store.get_stats()
        logger.info("Deleted %s: %d triples removed", filename, deleted_triples)
        return {
            "success": True,
            "message": f"Deleted '{filename}'. {deleted_triples} triples removed.",
            "stats": stats.model_dump(),
        }

    def get_file_list(self) -> List[Dict]:
        """Return metadata for all uploaded files."""
        result = []
        for meta in self.file_registry.values():
            triple_count = self.graph_store.get_file_triple_count(meta["file_id"])
            result.append({
                "filename": meta["filename"],
                "upload_date": meta["upload_time"][:19].replace("T", " "),
                "size_bytes": meta["size_bytes"],
                "triples": triple_count,
                "version": meta["version"],
            })
        return result

    def get_stats(self) -> GraphStats:
        return self.graph_store.get_stats()

    # --- Chat ---

    def request_stop(self, session_id: str):
        self._stop_flags[session_id] = True

    def _should_stop(self, session_id: str) -> bool:
        return self._stop_flags.get(session_id, False)

    async def chat_stream(
        self,
        message: str,
        history: List[Dict[str, str]],
        provider: str,
        session_id: str = "default",
    ) -> AsyncGenerator[Tuple[List[Dict], str, str], None]:
        """
        Stream a chat response.

        Yields tuples of (updated_history, intent_label, execution_info).
        The caller updates the Gradio chatbot with updated_history.
        """
        self._stop_flags[session_id] = False

        # 1. Run LangGraph workflow to determine intent and gather context
        state = self.workflow.run(message, provider)
        intent = state.get("intent", INTENT_GENERAL_CHAT)
        exec_path = " → ".join(state.get("execution_path", []))
        context = state.get("context", "")
        sources = state.get("sources", [])

        logger.info("Chat: intent=%s path=%s", intent, exec_path)

        # 2. Build system prompt
        if intent in (INTENT_KG_QUERY, INTENT_DOC_SEARCH) and context:
            system_content = (
                "You are an AI assistant with access to a knowledge graph.\n"
                "Use the following retrieved knowledge to answer the user's question.\n"
                "If the knowledge is insufficient, say so clearly.\n\n"
                f"## Knowledge Graph Context\n{context}\n"
            )
            if sources:
                system_content += f"\n## Sources: {', '.join(sources)}\n"
        elif intent == INTENT_PROGRAMMING:
            system_content = (
                "You are an expert programmer. Provide clear, well-structured "
                "code with explanations. Use markdown code blocks."
            )
        elif intent == INTENT_EXPLANATION:
            system_content = (
                "You are an expert educator. Explain concepts clearly with examples."
            )
        elif intent == INTENT_GREETING:
            system_content = "You are a friendly AI assistant. Greet the user warmly."
        else:
            system_content = "You are a helpful AI assistant."

        # 3. Build LangChain messages
        lc_messages: List[BaseMessage] = [SystemMessage(content=system_content)]
        for msg in history[-10:]:  # last 10 messages for context
            if msg["role"] == "user":
                lc_messages.append(HumanMessage(content=msg["content"]))
            elif msg["role"] == "assistant":
                lc_messages.append(AIMessage(content=msg["content"]))
        lc_messages.append(HumanMessage(content=message))

        # 4. Update history with user message
        updated = list(history) + [{"role": "user", "content": message}]
        updated.append({"role": "assistant", "content": ""})

        # 5. Stream LLM response
        llm = self.provider_mgr.get_llm(provider)
        response_text = ""

        try:
            async for chunk in llm.astream(lc_messages):
                if self._should_stop(session_id):
                    break
                token = chunk.content if hasattr(chunk, "content") else str(chunk)
                response_text += token
                updated[-1]["content"] = response_text
                yield updated, intent, exec_path
        except Exception as e:
            response_text = f"⚠️ Error generating response: {e}"
            updated[-1]["content"] = response_text
            yield updated, intent, exec_path

        # 6. Persist to conversation memory
        self.memory.add_message(session_id, "user", message)
        self.memory.add_message(session_id, "assistant", response_text)

    # --- Non-streaming chat (for REST API) ---

    def chat(self, message: str, provider: str,
             history: List[Dict[str, str]]) -> ChatResponse:
        """Non-streaming chat for the REST API."""
        state = self.workflow.run(message, provider)
        intent = state.get("intent", INTENT_GENERAL_CHAT)
        context = state.get("context", "")
        sources = state.get("sources", [])

        if intent in (INTENT_KG_QUERY, INTENT_DOC_SEARCH) and context:
            system_content = (
                "You are an AI assistant with access to a knowledge graph.\n"
                "Use the following retrieved knowledge to answer.\n\n"
                f"## Knowledge Graph Context\n{context}\n"
            )
        else:
            system_content = "You are a helpful AI assistant."

        lc_messages: List[BaseMessage] = [SystemMessage(content=system_content)]
        for msg in history[-10:]:
            if msg["role"] == "user":
                lc_messages.append(HumanMessage(content=msg["content"]))
            elif msg["role"] == "assistant":
                lc_messages.append(AIMessage(content=msg["content"]))
        lc_messages.append(HumanMessage(content=message))

        llm = self.provider_mgr.get_llm(provider)
        try:
            response = llm.invoke(lc_messages)
            reply = response.content
        except Exception as e:
            reply = f"⚠️ Error: {e}"

        return ChatResponse(
            reply=reply,
            intent=intent,
            execution_path=state.get("execution_path", []),
            sources=sources,
        )


# Create the global application instance
app_core = Application()


# ============================================================
# SECTION 13: FASTAPI ENDPOINTS
# ============================================================

api = FastAPI(title="AI Knowledge Graph Chat API", version="1.0.0")
api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@api.get("/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "neo4j": isinstance(app_core.graph_store, Neo4jGraphStore),
        "providers": app_core.provider_mgr.list_providers(),
    }


@api.post("/chat", response_model=ChatResponse)
async def chat_endpoint(req: ChatRequest):
    """Chat endpoint — returns full response (non-streaming)."""
    try:
        result = app_core.chat(req.message, req.provider, req.history)
        return result
    except Exception as e:
        logger.error("Chat error: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@api.post("/upload", response_model=UploadResponse)
async def upload_endpoint(file: UploadFile = File(...),
                          provider: str = Config.DEFAULT_PROVIDER):
    """Upload and process a document."""
    # Validate extension
    ext = Path(file.filename or "").suffix.lower()
    if ext not in Config.SUPPORTED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type: {ext}. Supported: {Config.SUPPORTED_EXTENSIONS}"
        )

    # Save to temp file
    tmp_path = Config.UPLOAD_DIR / f"{uuid.uuid4().hex}_{file.filename}"
    try:
        content = await file.read()
        if not content:
            raise HTTPException(status_code=400, detail="Empty file.")
        if len(content) > Config.MAX_UPLOAD_SIZE_MB * 1024 * 1024:
            raise HTTPException(
                status_code=413,
                detail=f"File too large. Max: {Config.MAX_UPLOAD_SIZE_MB} MB."
            )
        tmp_path.write_bytes(content)

        result = app_core.upload_file(str(tmp_path), file.filename, provider)
        return result
    except HTTPException:
        raise
    except Exception as e:
        logger.error("Upload error: %s", e)
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if tmp_path.exists():
            tmp_path.unlink(missing_ok=True)


@api.delete("/file/{filename}")
async def delete_file_endpoint(filename: str):
    """Delete a file and its graph data."""
    result = app_core.delete_file(filename)
    if not result["success"]:
        raise HTTPException(status_code=404, detail=result["message"])
    return result


@api.get("/files")
async def list_files_endpoint():
    """List all uploaded files."""
    return app_core.get_file_list()


@api.get("/graph/stats", response_model=GraphStats)
async def graph_stats_endpoint():
    """Return knowledge-graph statistics."""
    return app_core.get_stats()


@api.post("/provider")
async def set_provider_endpoint(req: ProviderRequest):
    """Validate a provider choice."""
    if req.provider.lower() not in PROVIDERS:
        raise HTTPException(status_code=400, detail="Unknown provider.")
    available = app_core.provider_mgr.is_available(req.provider)
    return {
        "provider": req.provider,
        "available": available,
        "message": "Provider is available." if available
                   else "Provider library installed but no API key set.",
    }


# ============================================================
# SECTION 14: GRADIO UI
# ============================================================

# Track the current Gradio click event for cancellation
_current_click_event = None


def _stats_markdown() -> str:
    """Render graph statistics as markdown."""
    stats = app_core.get_stats()
    return (
        f"### 📊 Knowledge Graph Stats\n"
        f"| Metric | Count |\n|---|---|\n"
        f"| 📄 Documents | **{stats.documents}** |\n"
        f"| 🔵 Entities | **{stats.entities}** |\n"
        f"| 🔗 Relationships | **{stats.relationships}** |\n"
    )


def _file_list_df():
    """Return file list as a pandas DataFrame for gr.Dataframe."""
    files = app_core.get_file_list()
    if not files:
        return pd.DataFrame(columns=["Filename", "Upload Date", "Size (KB)", "Triples", "Version"])
    return pd.DataFrame([
        {
            "Filename": f["filename"],
            "Upload Date": f["upload_date"],
            "Size (KB)": round(f["size_bytes"] / 1024, 1),
            "Triples": f["triples"],
            "Version": f["version"],
        }
        for f in files
    ])


def _file_choices():
    """Return list of filenames for the delete dropdown."""
    return [f["filename"] for f in app_core.get_file_list()]


def _provider_choices():
    """Return provider choices with availability indicators."""
    return [
        f"{PROVIDERS[k]['label']}{' ✅' if app_core.provider_mgr.is_available(k) else ' ⚠️'}"
        for k in PROVIDERS
    ]


def _provider_value_to_key(label: str) -> str:
    """Convert a display label back to a provider key."""
    for k, v in PROVIDERS.items():
        if v["label"] in label:
            return k
    return Config.DEFAULT_PROVIDER


async def _stream_response(message, history, provider_label, session_id):
    """Async generator that streams chat responses to Gradio."""
    provider = _provider_value_to_key(provider_label)
    intent_label = ""
    exec_info = ""

    async for updated_history, intent, exec_path in app_core.chat_stream(
        message, history, provider, session_id
    ):
        intent_label = intent
        exec_info = exec_path
        yield updated_history, intent_label, exec_info, _stats_markdown()


def _send_handler(message, history, provider_label):
    """Wrapper for the send button (sync generator for Gradio)."""
    session_id = "gradio_session"
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    async def gen():
        async for item in _stream_response(message, history, provider_label, session_id):
            yield item

    async_gen = gen()

    # Convert async generator to sync
    try:
        while True:
            try:
                item = loop.run_until_complete(async_gen.__anext__())
                yield item
            except StopAsyncIteration:
                break
    finally:
        loop.close()


def _upload_handler(files, provider_label):
    """Handle file uploads."""
    if not files:
        return "No files selected.", _file_list_df(), _stats_markdown(), gr.update(choices=_file_choices())

    provider = _provider_value_to_key(provider_label)
    results = []

    for f in files:
        try:
            result = app_core.upload_file(f.name, os.path.basename(f.name), provider)
            results.append(
                f"✅ **{result.filename}**: {result.triples_inserted} triples inserted, "
                f"{result.triples_skipped} skipped ({result.processing_time}s) — {result.message}"
            )
        except Exception as e:
            results.append(f"❌ **{os.path.basename(f.name)}**: {e}")

    return "\n".join(results), _file_list_df(), _stats_markdown(), gr.update(choices=_file_choices())


def _delete_handler(filename):
    """Handle file deletion."""
    if not filename:
        return "No file selected.", _file_list_df(), _stats_markdown(), gr.update(choices=_file_choices())
    result = app_core.delete_file(filename)
    return result["message"], _file_list_df(), _stats_markdown(), gr.update(choices=_file_choices())


def _clear_handler():
    """Clear the chat."""
    app_core.memory.clear("gradio_session")
    return [], "", ""


def _stop_handler():
    """Stop generation."""
    app_core.request_stop("gradio_session")
    return "Generation stopped."


def build_ui() -> gr.Blocks:
    """Build the Gradio interface."""
    with gr.Blocks(
        title="AI Knowledge Graph Chat",
        theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue"),
        css="""
        .main { max-width: 1400px; margin: auto; }
        .stats-box { background: #f0f4ff; padding: 12px; border-radius: 8px; }
        """
    ) as demo:
        gr.Markdown("# 🧠 AI Knowledge Graph Chat")
        gr.Markdown(
            "Upload documents to build a knowledge graph, then ask questions. "
            "The AI extracts entities and relationships, stores them in Neo4j, "
            "and routes your questions intelligently."
        )

        # --- Top bar: provider + stats ---
        with gr.Row():
            provider_dd = gr.Dropdown(
                choices=_provider_choices(),
                value=_provider_choices()[0] if _provider_choices() else "Groq",
                label="LLM Provider",
                scale=1,
                interactive=True,
            )
            stats_md = gr.Markdown(_stats_markdown(), elem_classes=["stats-box"], scale=2)
            refresh_btn = gr.Button("🔄 Refresh", scale=0)

        # --- Main area: chat + sidebar ---
        with gr.Row():
            # Chat column
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=480,
                    show_copy_button=True,
                    type="messages",
                    render_markdown=True,
                    avatar_images=("👤", "🤖"),
                )
                intent_md = gr.Markdown("", label="Intent")
                exec_md = gr.Markdown("", label="Execution Path")

                with gr.Row():
                    msg_input = gr.Textbox(
                        placeholder="Type your message... (Enter to send)",
                        show_label=False,
                        scale=4,
                        lines=2,
                    )
                    send_btn = gr.Button("📤 Send", variant="primary", scale=1)
                    stop_btn = gr.Button("⏹️ Stop", variant="stop", scale=1)
                    clear_btn = gr.Button("🗑️ Clear", scale=1)

            # Sidebar: file management
            with gr.Column(scale=2):
                gr.Markdown("### 📁 Upload Documents")
                file_upload = gr.File(
                    label="Drop files here",
                    file_count="multiple",
                    file_types=[ext.lstrip(".") for ext in Config.SUPPORTED_EXTENSIONS],
                )
                upload_status = gr.Markdown("")
                upload_btn = gr.Button("Process Files", variant="primary")

                gr.Markdown("---")
                gr.Markdown("### 📋 Uploaded Files")
                file_df = gr.Dataframe(
                    value=_file_list_df(),
                    headers=["Filename", "Upload Date", "Size (KB)", "Triples", "Version"],
                    datatype=["str", "str", "str", "number", "number"],
                    interactive=False,
                    wrap=True,
                )

                with gr.Row():
                    delete_dd = gr.Dropdown(
                        choices=_file_choices(),
                        label="Select file to delete",
                        scale=3,
                    )
                    delete_btn = gr.Button("🗑️ Delete", variant="stop", scale=1)

        # --- Event wiring ---

        # Send message (streaming)
        click_event = send_btn.click(
            fn=_send_handler,
            inputs=[msg_input, chatbot, provider_dd],
            outputs=[chatbot, intent_md, exec_md, stats_md],
        ).then(
            fn=lambda: "",
            outputs=msg_input,
        )

        # Enter key sends
        msg_input.submit(
            fn=_send_handler,
            inputs=[msg_input, chatbot, provider_dd],
            outputs=[chatbot, intent_md, exec_md, stats_md],
        ).then(
            fn=lambda: "",
            outputs=msg_input,
        )

        # Stop generation
        stop_btn.click(
            fn=_stop_handler,
            outputs=upload_status,
            cancels=[click_event],
        )

        # Clear conversation
        clear_btn.click(
            fn=_clear_handler,
            outputs=[chatbot, intent_md, exec_md],
        )

        # Upload files
        upload_btn.click(
            fn=_upload_handler,
            inputs=[file_upload, provider_dd],
            outputs=[upload_status, file_df, stats_md, delete_dd],
        )

        # Delete file
        delete_btn.click(
            fn=_delete_handler,
            inputs=[delete_dd],
            outputs=[upload_status, file_df, stats_md, delete_dd],
        )

        # Refresh stats
        refresh_btn.click(
            fn=lambda: (_stats_markdown(), _file_list_df(), gr.update(choices=_file_choices())),
            outputs=[stats_md, file_df, delete_dd],
        )

    return demo


# ============================================================
# SECTION 15: MAIN ENTRYPOINT
# ============================================================

def main():
    """
    Entry point: mount Gradio on FastAPI and run with uvicorn.
    This serves both the UI and REST API from a single port.
    """
    os.system("")  # ensure terminal output on Windows

    logger.info("=" * 60)
    logger.info("AI Knowledge Graph Chat Application")
    logger.info("=" * 60)
    logger.info("Neo4j available: %s", _NEO4J_OK)
    logger.info("Neo4j configured: %s", bool(Config.NEO4J_URI))
    logger.info("Graph store: %s", type(app_core.graph_store).__name__)

    for key, meta in PROVIDERS.items():
        avail = app_core.provider_mgr.is_available(key)
        logger.info("Provider %s: installed=%s, api_key=%s",
                    key, meta["available"], avail)

    # Build Gradio UI
    demo = build_ui()

    # Mount Gradio on FastAPI at root
    gr.mount_gradio_app(api, demo, path="/")

    logger.info("Starting server on %s:%d", Config.HOST, Config.PORT)
    uvicorn.run(api, host=Config.HOST, port=Config.PORT, log_level="info")


if __name__ == "__main__":
    main()