#!/usr/bin/env python3
"""
AI Knowledge Graph Chat Application (SQLite Edition)
====================================================

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)
  • SQLite storage (Free, Serverless, Persistent) with optional Neo4j 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
import sqlite3
import threading
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 (Optional, for fallback) ---
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)

    # SQLite (Default DB)
    SQLITE_DB_PATH: str = os.getenv("SQLITE_DB_PATH", "./data/knowledge_graph.db")

    # Neo4j (Optional)
    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.
    """

    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*."""
        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:
        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)
        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")
            for tag in soup(["script", "style"]):
                tag.decompose()
            return soup.get_text(separator="\n")
        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": []}}.
"""


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]:
        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]:
        content = re.sub(r"```(?:json)?\s*", "", content)
        content = content.strip().rstrip("`")
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass
        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]]:
        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)

        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)

        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 (SQLite Default + InMemory/Neo4j 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: ...


# --- SQLite backend (Default) ----------------------------------------------

class SQLiteGraphStore(GraphStoreBase):
    """
    SQLite-backed graph store.
    Stores nodes and relationships in relational tables with JSON arrays
    for metadata (source_files, etc.) to allow incremental updates.
    """

    def __init__(self, db_path: str):
        self._db_path = db_path
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._lock = threading.Lock()
        self._init_db()
        logger.info("Connected to SQLite at %s", db_path)

    def _init_db(self):
        with self._lock:
            cursor = self._conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    file_id TEXT PRIMARY KEY,
                    filename TEXT,
                    file_hash TEXT,
                    upload_time TEXT
                )
            """)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS entities (
                    name TEXT PRIMARY KEY,
                    source_files TEXT,
                    filenames TEXT,
                    upload_times TEXT,
                    versions TEXT,
                    confidences TEXT,
                    created_at TEXT
                )
            """)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS relationships (
                    hash TEXT PRIMARY KEY,
                    subject TEXT,
                    relation TEXT,
                    object TEXT,
                    source_files TEXT,
                    filenames TEXT,
                    upload_times TEXT,
                    versions TEXT,
                    confidences TEXT,
                    created_at TEXT
                )
            """)
            self._conn.commit()

    @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):
        with self._lock:
            self._conn.execute(
                "INSERT OR REPLACE INTO documents VALUES (?, ?, ?, ?)",
                (file_id, filename, file_hash, upload_time)
            )
            self._conn.commit()

    def add_triple(self, subject, relation, object_, file_id, filename,
                   upload_time, version, confidence):
        h = self._triple_hash(subject, relation, object_)
        now = datetime.now(timezone.utc).isoformat()

        with self._lock:
            cursor = self._conn.cursor()
            cursor.execute("SELECT * FROM relationships WHERE hash = ?", (h,))
            row = cursor.fetchone()

            if row:
                sfs = json.loads(row["source_files"])
                if file_id not in sfs:
                    sfs.append(file_id)
                    fns = json.loads(row["filenames"]) + [filename]
                    uts = json.loads(row["upload_times"]) + [upload_time]
                    vers = json.loads(row["versions"]) + [version]
                    confs = json.loads(row["confidences"]) + [confidence]

                    self._conn.execute("""
                        UPDATE relationships SET
                        source_files=?, filenames=?, upload_times=?, versions=?, confidences=?
                        WHERE hash=?
                    """, (json.dumps(sfs), json.dumps(fns), json.dumps(uts),
                          json.dumps(vers), json.dumps(confs), h))
                    self._conn.commit()
                return False  # Skipped

            self._conn.execute("""
                INSERT INTO relationships VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                h, subject, relation, object_,
                json.dumps([file_id]), json.dumps([filename]),
                json.dumps([upload_time]), json.dumps([version]),
                json.dumps([confidence]), now
            ))

            for name in (subject, object_):
                cursor.execute("SELECT * FROM entities WHERE name = ?", (name,))
                ent_row = cursor.fetchone()
                if ent_row:
                    sfs = json.loads(ent_row["source_files"])
                    if file_id not in sfs:
                        sfs.append(file_id)
                        fns = json.loads(ent_row["filenames"]) + [filename]
                        uts = json.loads(ent_row["upload_times"]) + [upload_time]
                        vers = json.loads(ent_row["versions"]) + [version]
                        confs = json.loads(ent_row["confidences"]) + [confidence]

                        self._conn.execute("""
                            UPDATE entities SET
                            source_files=?, filenames=?, upload_times=?, versions=?, confidences=?
                            WHERE name=?
                        """, (json.dumps(sfs), json.dumps(fns), json.dumps(uts),
                              json.dumps(vers), json.dumps(confs), name))
                else:
                    self._conn.execute("""
                        INSERT INTO entities VALUES (?, ?, ?, ?, ?, ?, ?)
                    """, (
                        name,
                        json.dumps([file_id]), json.dumps([filename]),
                        json.dumps([upload_time]), json.dumps([version]),
                        json.dumps([confidence]), now
                    ))
            self._conn.commit()
            return True  # Inserted

    def delete_by_file(self, file_id):
        deleted = 0
        with self._lock:
            cursor = self._conn.cursor()
            
            cursor.execute("SELECT hash, source_files, filenames, upload_times, versions, confidences FROM relationships")
            for row in cursor.fetchall():
                sfs = json.loads(row["source_files"])
                if file_id in sfs:
                    idx = sfs.index(file_id)
                    for col in ["source_files", "filenames", "upload_times", "versions", "confidences"]:
                        arr = json.loads(row[col])
                        arr.pop(idx)
                        self._conn.execute(f"UPDATE relationships SET {col}=? WHERE hash=?", (json.dumps(arr), row["hash"]))

                    if not sfs:
                        self._conn.execute("DELETE FROM relationships WHERE hash=?", (row["hash"],))
                        deleted += 1

            cursor.execute("SELECT name, source_files, filenames, upload_times, versions, confidences FROM entities")
            for row in cursor.fetchall():
                sfs = json.loads(row["source_files"])
                if file_id in sfs:
                    idx = sfs.index(file_id)
                    for col in ["source_files", "filenames", "upload_times", "versions", "confidences"]:
                        arr = json.loads(row[col])
                        arr.pop(idx)
                        self._conn.execute(f"UPDATE entities SET {col}=? WHERE name=?", (json.dumps(arr), row["name"]))

                    if not sfs:
                        self._conn.execute("DELETE FROM entities WHERE name=?", (row["name"],))

            self._conn.execute("DELETE FROM documents WHERE file_id=?", (file_id,))
            self._conn.commit()
        return deleted

    def get_stats(self):
        with self._lock:
            cursor = self._conn.cursor()
            docs = cursor.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
            ents = cursor.execute("SELECT COUNT(*) FROM entities").fetchone()[0]
            rels = cursor.execute("SELECT COUNT(*) FROM relationships").fetchone()[0]
        return GraphStats(documents=docs, entities=ents, relationships=rels)

    def search(self, query, limit=20):
        keywords = [w.lower().strip() for w in re.split(r"\s+", query) if len(w) > 2]
        if not keywords:
            return []
        
        conditions = " OR ".join(
            [f"(LOWER(subject) LIKE '%{kw}%' OR LOWER(relation) LIKE '%{kw}%' OR LOWER(object) LIKE '%{kw}%')"
             for kw in keywords]
        )
        sql = f"SELECT subject, relation, object, filenames FROM relationships WHERE {conditions} LIMIT 500"
        
        with self._lock:
            cursor = self._conn.cursor()
            cursor.execute(sql)
            rows = cursor.fetchall()

        results = []
        for row in rows:
            text = f"{row['subject']} {row['relation']} {row['object']}".lower()
            score = sum(1 for kw in keywords if kw in text)
            results.append({
                "subject": row["subject"],
                "relation": row["relation"],
                "object": row["object"],
                "sources": json.loads(row["filenames"]),
                "score": score
            })
        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:limit]

    def get_file_triple_count(self, file_id):
        with self._lock:
            cursor = self._conn.cursor()
            cursor.execute("SELECT source_files FROM relationships")
            count = 0
            for row in cursor.fetchall():
                if file_id in json.loads(row["source_files"]):
                    count += 1
        return count

    def close(self):
        with self._lock:
            self._conn.close()


# --- Neo4j backend (Optional Fallback) -------------------------------------

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

    def _init_constraints(self):
        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:
            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:
                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

            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

    def delete_by_file(self, file_id):
        deleted = 0
        with self._driver.session() as session:
            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

            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,
            )

            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):
        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 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."""
    # 1. Use Neo4j if explicitly configured
    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 SQLite.", e)
    
    # 2. Default to SQLite (Free & Serverless)
    logger.info("Using SQLite graph store.")
    return SQLiteGraphStore(Config.SQLITE_DB_PATH)


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

class ConversationMemory:
    """
    Maintains per-session conversation history.
    This is SEPARATE from the knowledge graph.
    """

    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]:
        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:
        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()
            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


# ============================================================
# 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 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()

    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:
        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

        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

    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"

    def run(self, user_input: str, provider: str) -> AgentState:
        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 tying all subsystems together."""

    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] = {}
        self._stop_flags: Dict[str, bool] = {}

    @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]:
        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:
        start = time.time()

        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).")

        file_hash = self._file_hash(filepath)

        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.",
            )

        version = 1
        for fid, meta in self.file_registry.items():
            if meta["filename"] == original_name:
                version = max(version, meta["version"] + 1)

        try:
            text = FileProcessor.extract_text(filepath)
        except Exception as e:
            raise RuntimeError(f"Text extraction failed: {e}")

        try:
            entities, relationships = self.extractor.extract(text, provider)
        except Exception as e:
            raise RuntimeError(f"Knowledge extraction failed: {e}")

        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)

        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

        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 time=%.2fs",
            original_name, inserted, skipped, 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]:
        target_id = None
        for fid, meta in self.file_registry.items():
            if meta["filename"] == filename:
                target_id = fid

        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]

        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]:
        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()

    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]:
        self._stop_flags[session_id] = False

        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", [])

        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."
        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."

        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))

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

        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

        self.memory.add_message(session_id, "user", message)
        self.memory.add_message(session_id, "assistant", response_text)

    def chat(self, message: str, provider: str,
             history: List[Dict[str, str]]) -> ChatResponse:
        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():
    return {
        "status": "healthy",
        "db_type": type(app_core.graph_store).__name__,
        "providers": app_core.provider_mgr.list_providers(),
    }

@api.post("/chat", response_model=ChatResponse)
async def chat_endpoint(req: ChatRequest):
    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):
    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}")

    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="File too large.")
        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):
    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():
    return app_core.get_file_list()

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

@api.post("/provider")
async def set_provider_endpoint(req: ProviderRequest):
    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 "No API key set.",
    }


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

def _stats_markdown() -> str:
    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():
    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 [f["filename"] for f in app_core.get_file_list()]

def _provider_choices():
    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:
    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):
    provider = _provider_value_to_key(provider_label)
    async for updated_history, intent, exec_path in app_core.chat_stream(
        message, history, provider, session_id
    ):
        yield updated_history, intent, exec_path, _stats_markdown()

def _send_handler(message, history, provider_label):
    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()
    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):
    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} inserted, {result.triples_skipped} skipped.")
        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):
    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():
    app_core.memory.clear("gradio_session")
    return [], "", ""

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

def build_ui() -> gr.Blocks:
    custom_css = """
    .main { max-width: 1400px; margin: auto; }
    .stats-box { background: #f0f4ff; padding: 12px; border-radius: 8px; }
    """
    with gr.Blocks(
        title="AI Knowledge Graph Chat",
    ) as demo:
        gr.HTML(f"<style>{custom_css}</style>")
        gr.Markdown("# 🧠 AI Knowledge Graph Chat (SQLite Edition)")
        gr.Markdown(
            "Upload documents to build a knowledge graph, then ask questions. "
            "Data is stored locally in a free SQLite database."
        )

        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)

        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=480,
                    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)

            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)

        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,
        )

        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_btn.click(
            fn=_stop_handler,
            outputs=upload_status,
            cancels=[click_event],
        )

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

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

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

        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.
    """
    os.system("")  # ensure terminal output on Windows

    logger.info("=" * 60)
    logger.info("AI Knowledge Graph Chat Application (SQLite Edition)")
    logger.info("=" * 60)
    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)

    demo = build_ui()
    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()