import os
import re
import glob
import shutil
import threading
import logging
from typing import List, Tuple, Optional

import faiss
import numpy as np
import yaml
from sentence_transformers import SentenceTransformer
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel

# ─── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# ─── Config ───────────────────────────────────────────────────────────────────
_CONFIG_PATH = "./config.yaml"

def load_config(path: str = _CONFIG_PATH) -> dict:
    defaults = {
        "kb": {"directory": "./kb", "index_directory": "./.cache/faiss"},
        "models": {
            "embedding": "sentence-transformers/all-MiniLM-L6-v2",
            "qa": "google/flan-t5-small",
        },
        "chunking": {"chunk_size": 500, "overlap": 100},
        "retrieval": {"top_k": 5, "min_similarity": 0.25, "max_sentences": 6},
        "client": {"name": "Self-Service KB Assistant"},
        "messages": {
            "welcome": "Ask me anything about the documents in the knowledge base!",
        },
        "quick_actions": [],
    }
    if os.path.exists(path):
        try:
            with open(path, "r", encoding="utf-8") as f:
                user_cfg = yaml.safe_load(f) or {}
            for key, val in user_cfg.items():
                if isinstance(val, dict) and key in defaults:
                    defaults[key].update(val)
                else:
                    defaults[key] = val
            logger.info(f"Config loaded from {path}")
        except Exception as e:
            logger.warning(f"Could not load {path}: {e}. Using defaults.")
    return defaults

CONFIG               = load_config()
KB_DIR               = CONFIG["kb"]["directory"]
INDEX_DIR            = CONFIG["kb"]["index_directory"]
EMBEDDING_MODEL_NAME = CONFIG["models"]["embedding"]
QA_MODEL_NAME        = CONFIG["models"]["qa"]
CHUNK_SIZE           = CONFIG["chunking"]["chunk_size"]
CHUNK_OVERLAP        = CONFIG["chunking"]["overlap"]
TOP_K                = CONFIG["retrieval"]["top_k"]
MIN_SIMILARITY       = CONFIG["retrieval"]["min_similarity"]
MAX_SENTENCES        = CONFIG["retrieval"]["max_sentences"]
MAX_QUERY_CHARS      = 800
ALLOWED_EXTENSIONS   = {".txt", ".md", ".pdf", ".docx", ".doc"}
MAX_FILE_SIZE_MB     = 10
_IDX_PATH            = os.path.join(INDEX_DIR, "kb.index")
_META_PATH           = os.path.join(INDEX_DIR, "kb_meta.npy")

# ─── File loading ──────────────────────────────────────────────────────────────
def load_file_text(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".pdf":
        from PyPDF2 import PdfReader
        reader = PdfReader(path)
        return "\n".join(p.extract_text() or "" for p in reader.pages)
    elif ext in (".docx", ".doc"):
        import docx as python_docx
        doc = python_docx.Document(path)
        return "\n".join(p.text for p in doc.paragraphs if p.text.strip())
    else:
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            return f.read()

def load_kb_documents(kb_dir: str) -> List[Tuple[str, str]]:
    docs = []
    os.makedirs(kb_dir, exist_ok=True)
    for pat in ["*.txt", "*.md", "*.pdf", "*.docx", "*.doc"]:
        for path in sorted(glob.glob(os.path.join(kb_dir, pat))):
            try:
                text = load_file_text(path)
                if text.strip():
                    docs.append((os.path.basename(path), text))
            except Exception as e:
                logger.warning(f"Skipped {path}: {e}")
    return docs

# ─── Chunking ──────────────────────────────────────────────────────────────────
def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[str]:
    if not text:
        return []
    overlap = min(overlap, chunk_size - 1)
    sentences = re.compile(r'(?<=[.!?])\s+').split(text.strip())
    chunks, current, current_len = [], [], 0
    for sent in sentences:
        sent = sent.strip()
        if not sent:
            continue
        slen = len(sent)
        if slen > chunk_size:
            if current:
                chunks.append(" ".join(current))
                current, current_len = [], 0
            start = 0
            while start < slen:
                chunks.append(sent[start:min(start + chunk_size, slen)])
                start += chunk_size - overlap
            continue
        if current_len + slen + 1 > chunk_size:
            if current:
                chunks.append(" ".join(current))
            carry, carry_len = [], 0
            for s in reversed(current):
                if carry_len + len(s) + 1 > overlap:
                    break
                carry.insert(0, s)
                carry_len += len(s) + 1
            current, current_len = carry, carry_len
        current.append(sent)
        current_len += slen + 1
    if current:
        chunks.append(" ".join(current))
    return [c for c in chunks if len(c.strip()) > 20]

def clean_context(text: str) -> str:
    lines, cleaned, seen = text.splitlines(), [], set()
    for line in lines:
        l = re.sub(r'^#+\s*', '', line.strip())
        l = re.sub(r'^\d+[\.\)]\s*', '', l)
        l = re.sub(r'^[-*]\s*', '', l)
        if len(l) < 5 or l in seen:
            continue
        words = l.split()
        if words and len(words) <= 10 and sum(1 for w in words if w[:1].isupper()) >= len(words) - 1:
            continue
        seen.add(l)
        cleaned.append(l)
    return "\n".join(cleaned)

# ─── KB Index (FAISS) ──────────────────────────────────────────────────────────
class KBIndex:
    def __init__(self):
        self.embedder: Optional[SentenceTransformer] = None
        self.chunks:   List[str] = []
        self.sources:  List[str] = []
        self.index:    Optional[faiss.Index] = None
        self._lock = threading.Lock()
        logger.info("Loading embedding model…")
        self.embedder = SentenceTransformer(EMBEDDING_MODEL_NAME)
        logger.info("Embedding model loaded.")
        os.makedirs(INDEX_DIR, exist_ok=True)
        self._load_or_build()

    def _load_or_build(self):
        if os.path.exists(_IDX_PATH) and os.path.exists(_META_PATH):
            try:
                self.index = faiss.read_index(_IDX_PATH)
                meta = np.load(_META_PATH, allow_pickle=True).item()
                self.chunks  = list(meta["chunks"])
                self.sources = list(meta["sources"])
                logger.info(f"✅ FAISS index loaded — {len(self.chunks)} chunks.")
                return
            except Exception as e:
                logger.warning(f"Could not load index: {e}. Rebuilding…")
        self._build()

    def _build(self):
        docs = load_kb_documents(KB_DIR)
        if not docs:
            demo = ("demo_content.txt",
                    "Welcome to the Self-Service KB Assistant. "
                    "Upload your own documents via the Knowledge Base panel. "
                    "The assistant will index them and answer questions from their content.")
            docs = [demo]
        all_chunks, all_sources = [], []
        for source, text in docs:
            for chunk in chunk_text(text):
                all_chunks.append(chunk)
                all_sources.append(source)
        if not all_chunks:
            self.index = None; self.chunks = []; self.sources = []
            return
        logger.info(f"Encoding {len(all_chunks)} chunks…")
        embeddings = self.embedder.encode(
            all_chunks, show_progress_bar=False,
            convert_to_numpy=True, batch_size=64
        ).astype("float32")
        faiss.normalize_L2(embeddings)
        index = faiss.IndexFlatIP(embeddings.shape[1])
        index.add(embeddings)
        self.index = index; self.chunks = all_chunks; self.sources = all_sources
        try:
            faiss.write_index(index, _IDX_PATH)
            np.save(_META_PATH, {"chunks": np.array(all_chunks, dtype=object),
                                 "sources": np.array(all_sources, dtype=object)})
            logger.info("FAISS index saved.")
        except Exception as e:
            logger.warning(f"Could not save index: {e}")

    def rebuild(self) -> str:
        with self._lock:
            for p in (_IDX_PATH, _META_PATH):
                if os.path.exists(p): os.remove(p)
            self._build()
        if not self.chunks:
            return "Index rebuilt (empty — no documents found)."
        return f"Index rebuilt — {len(self.chunks)} chunks from {len(set(self.sources))} files."

    def retrieve(self, query: str) -> List[Tuple[str, str, float]]:
        if not query.strip() or self.index is None or not self.chunks:
            return []
        q = self.embedder.encode([query], convert_to_numpy=True).astype("float32")
        faiss.normalize_L2(q)
        scores, idxs = self.index.search(q, min(TOP_K, len(self.chunks)))
        return [
            (self.chunks[i], self.sources[i], float(s))
            for s, i in zip(scores[0], idxs[0])
            if 0 <= i < len(self.chunks) and float(s) >= MIN_SIMILARITY
        ]

    def best_sentences(self, query: str, results: List[Tuple[str, str, float]]) -> str:
        sents = []
        for chunk, _, _ in results:
            for s in re.split(r'(?<=[.!?])\s+|\n+', clean_context(chunk)):
                if len(s.strip()) >= 25:
                    sents.append(s.strip())
        if not sents:
            return ""
        q_emb = self.embedder.encode([query], convert_to_numpy=True).astype("float32")
        s_emb = self.embedder.encode(sents,  convert_to_numpy=True).astype("float32")
        faiss.normalize_L2(q_emb); faiss.normalize_L2(s_emb)
        sims = np.dot(s_emb, q_emb.T).reshape(-1)
        q_words = {w.lower() for w in re.findall(r"\w+", query) if len(w) > 3}
        lex = np.array([len(q_words & {w.lower() for w in re.findall(r"\w+", s) if len(w) > 3})
                        for s in sents], dtype=float)
        top = np.argsort(-(1.5 * sims + 0.5 * lex))
        chosen, seen = [], set()
        for i in top:
            s = sents[i].strip()
            if s and s not in seen:
                chosen.append(s); seen.add(s)
            if len(chosen) >= MAX_SENTENCES:
                break
        return "\n".join(chosen)

# ─── QA Model ─────────────────────────────────────────────────────────────────
qa_tokenizer = None
qa_model     = None
qa_available = False
_qa_lock     = threading.Lock()

try:
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
    logger.info(f"Loading QA model: {QA_MODEL_NAME}…")
    qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_NAME)
    qa_model     = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_NAME)
    qa_model.eval()
    qa_available = True
    logger.info(f"✅ QA model loaded ({QA_MODEL_NAME})")
except Exception as e:
    logger.warning(f"⚠️ Could not load QA model: {e}. Using extractive fallback.")

def generate_answer(question: str, context: str) -> Optional[str]:
    if not qa_available:
        return None
    prompt = (
        "You are a warm, helpful assistant. Answer using ONLY the context below. "
        "Open with a friendly sentence. Introduce any steps with a sentence first. "
        "Keep tone conversational. If context is insufficient, say so honestly.\n\n"
        f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:"
    )
    try:
        with _qa_lock:
            inputs = qa_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=768)
            outputs = qa_model.generate(**inputs, max_new_tokens=220, temperature=1.0, do_sample=False)
            answer = qa_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
        return answer or None
    except Exception as e:
        logger.error(f"QA error: {e}")
        return None

# ─── Init index ───────────────────────────────────────────────────────────────
logger.info("Initialising KB index…")
kb_index = KBIndex()

# ─── RAG pipeline ─────────────────────────────────────────────────────────────
def build_answer(query: str, history: list) -> dict:
    """Returns {"answer": str, "sources": [str]}"""
    results = kb_index.retrieve(query)
    if not results:
        return {"answer": "I couldn't find anything relevant for your question. Try rephrasing or using different keywords.", "sources": []}

    context = kb_index.best_sentences(query, results)
    if not context:
        return {"answer": "I found some content but couldn't extract relevant sentences. Please try rephrasing.", "sources": []}

    # Collect unique sources above threshold
    seen_sources: List[str] = []
    for _, src, score in results:
        if src not in seen_sources and score >= MIN_SIMILARITY:
            seen_sources.append(src)
        if len(seen_sources) == 3:
            break

    answer = generate_answer(query, context)
    if not answer:
        answer = "Here's the most relevant information I found:\n\n" + "\n".join(
            f"• {s}" for s in context.split("\n") if s.strip()
        )
    return {"answer": answer, "sources": seen_sources}

# ─── FastAPI app ───────────────────────────────────────────────────────────────
app = FastAPI(title="Self-Service KB Assistant")
app.mount("/static", StaticFiles(directory="static"), name="static")

# ── Pydantic models ────────────────────────────────────────────────────────────
class ChatRequest(BaseModel):
    message: str
    history: list = []

class DeleteRequest(BaseModel):
    filename: str

class RenameRequest(BaseModel):
    old_name: str
    new_name: str

# ── Routes ─────────────────────────────────────────────────────────────────────
@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.get("/api/config")
async def get_config():
    quick = [q["query"] for q in CONFIG.get("quick_actions", []) if "query" in q] or [
        "What makes a good knowledge base article?",
        "How could a KB assistant help support agents?",
        "Why is self-service important for customer support?",
    ]
    return {
        "name":         CONFIG["client"]["name"],
        "welcome":      CONFIG["messages"].get("welcome", "Ask me anything!"),
        "quick_actions": quick,
    }

@app.post("/api/chat")
async def chat(req: ChatRequest):
    if not req.message.strip():
        return {"answer": "Please ask a question.", "sources": []}
    try:
        return build_answer(req.message.strip()[:MAX_QUERY_CHARS], req.history)
    except Exception as e:
        logger.error(f"Chat error: {e}")
        raise HTTPException(status_code=500, detail="Something went wrong. Please try again.")

@app.get("/api/files")
async def list_files():
    files = []
    for pat in ["*.txt", "*.md", "*.pdf", "*.docx", "*.doc"]:
        for p in sorted(glob.glob(os.path.join(KB_DIR, pat))):
            name = os.path.basename(p)
            size = os.path.getsize(p)
            files.append({"name": name, "size": size})
    return {"files": files, "count": len(files)}

@app.post("/api/files/upload")
async def upload_files(files: List[UploadFile] = File(...)):
    os.makedirs(KB_DIR, exist_ok=True)
    saved, skipped = [], []
    for file in files:
        name = os.path.basename(file.filename or "")
        ext  = os.path.splitext(name)[1].lower()
        if ext not in ALLOWED_EXTENSIONS:
            skipped.append(f"{name} (unsupported type)")
            continue
        content = await file.read()
        if len(content) / (1024 * 1024) > MAX_FILE_SIZE_MB:
            skipped.append(f"{name} (exceeds {MAX_FILE_SIZE_MB} MB)")
            continue
        try:
            with open(os.path.join(KB_DIR, name), "wb") as f:
                f.write(content)
            saved.append(name)
        except Exception as e:
            skipped.append(f"{name} (error: {e})")
    msg = ""
    if saved:
        msg = kb_index.rebuild()
    return {"saved": saved, "skipped": skipped, "message": msg}

@app.delete("/api/files/{filename}")
async def delete_file(filename: str):
    filename = os.path.basename(filename)
    if os.path.splitext(filename)[1].lower() not in ALLOWED_EXTENSIONS:
        raise HTTPException(status_code=400, detail="Invalid file type")
    target = os.path.join(KB_DIR, filename)
    if not os.path.exists(target):
        raise HTTPException(status_code=404, detail="File not found")
    os.remove(target)
    msg = kb_index.rebuild()
    return {"deleted": filename, "message": msg}

@app.put("/api/files/rename")
async def rename_file(req: RenameRequest):
    old = os.path.basename(req.old_name.strip())
    new = os.path.basename(req.new_name.strip())
    if not old or not new:
        raise HTTPException(status_code=400, detail="Both names required")
    if os.path.splitext(new)[1].lower() not in ALLOWED_EXTENSIONS:
        raise HTTPException(status_code=400, detail="New name must have a supported extension")
    src = os.path.join(KB_DIR, old)
    if not os.path.exists(src):
        raise HTTPException(status_code=404, detail=f"File not found: {old}")
    os.rename(src, os.path.join(KB_DIR, new))
    msg = kb_index.rebuild()
    return {"renamed_to": new, "message": msg}

# ─── Entry point ──────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)