# app/main.py
from __future__ import annotations
import shutil
import re
import uuid
from pathlib import Path

from fastapi import FastAPI, Request, Header, HTTPException, UploadFile, File
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

from openai import OpenAI
from openai import OpenAIError


from fastapi.responses import StreamingResponse
from io import BytesIO
from datetime import datetime

from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet

from ecdallm.app.rag import rag
from ecdallm.app.vector import DEFAULT_EMBED_MODEL

from ecdallm.app.paths import (
    UPLOAD_DIR,
    RAG_STORE_DIR,
    STATIC_DIR,
    TEMPLATES_DIR,
    ensure_dirs,
)

import os
import threading
import webbrowser
from contextlib import asynccontextmanager


ensure_dirs()




@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Runs on application startup/shutdown.
    Opens the browser only after startup completes (no FastAPI deprecation warning).
    """
    url = os.getenv("ECDALLM_URL")

    if url and os.getenv("ECDALLM_BROWSER_OPENED") != "1":
        os.environ["ECDALLM_BROWSER_OPENED"] = "1"

        def _open() -> None:
            try:
                webbrowser.open(url)
            except Exception:
                pass

        threading.Thread(target=_open, daemon=True).start()

    yield


app = FastAPI(lifespan=lifespan)

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))






ALLOWED_EXTS = {".pdf", ".txt", ".docx"}



# -----------------------------
# Helpers
# -----------------------------


def sanitize_filename(name: str) -> str:
    name = (name or "").strip().replace(" ", "_")
    name = Path(name).name
    name = re.sub(r"[^a-zA-Z0-9._-]", "", name)
    name = name.lstrip(".").strip()
    return name or "file"



def normalize_base_url(base_url: str) -> str:
    base = (base_url or "").strip().rstrip("/")
    if not base:
        raise ValueError("Base URL is required")

    if base.endswith("/v1"):
        return base
    return base + "/v1"


def make_openai_client(base_url: str, api_key: str | None) -> OpenAI:
    return OpenAI(base_url=normalize_base_url(base_url), api_key=(api_key or "lm-studio"))


def derive_original_from_stored(stored_name: str) -> str:
    p = Path(stored_name)
    if "__" in p.stem:
        stem = p.stem.split("__", 1)[0]
        return f"{stem}{p.suffix}"
    return stored_name


def choose_chat_model(model_ids: list[str]) -> str | None:
    """
    Choose a likely chat/instruct model from /v1/models.
    """
    if not model_ids:
        return None

    preferred = []
    for mid in model_ids:
        m = (mid or "").lower()
        if any(k in m for k in ["instruct", "chat", "llama", "mistral", "qwen", "gemma"]):
            preferred.append(mid)

    return preferred[0] if preferred else model_ids[0]


def get_cfg_from_headers(
    x_llm_base_url: str | None,
    x_llm_api_key: str | None,
    x_llm_model: str | None,
    x_embed_model: str | None,
):
    """
    Embeddings are LOCAL via FastEmbed, but we still accept X-Embed-Model so the user can override.
    """
    if not x_llm_base_url:
        raise HTTPException(400, detail="Missing X-LLM-Base-URL header")
    if not x_llm_model:
        raise HTTPException(400, detail="Missing X-LLM-Model header")

    return {
        "base_url": x_llm_base_url,
        "api_key": x_llm_api_key or "lm-studio",
        "chat_model": x_llm_model,
        "embed_model": (x_embed_model or DEFAULT_EMBED_MODEL),
    }


def is_greeting(text: str) -> bool:
    """
    True only if the message is basically JUST a greeting (optionally punctuation/emojis).
    Examples that match: "hi", "hi!", "hello 🙂", "hey!!!"
    Examples that do NOT match: "hi, give me the db", "hello can you explain...", "this is awesome hi"
    """
    t = (text or "").strip().lower()
    if not t:
        return False

    # remove common punctuation/emojis spacing at ends
    t = re.sub(r"[\s\!\.\,\?\:\;\-\_]+$", "", t)  # trailing punctuation
    t = re.sub(r"^[\s\!\.\,\?\:\;\-\_]+", "", t)  # leading punctuation

    # allow only single-token greetings (or two-token like "good morning" if you want)
    # Here we keep it simple: single word greetings only.
    greetings = {"hi", "hello", "hey", "hola", "salam"}

    # If it contains spaces => not a pure greeting
    if " " in t:
        return False

    return t in greetings

def is_positive_ack(text: str) -> bool:
    t = (text or "").strip().lower()
    t = re.sub(r"[!.,?]+$", "", t)
    positives = {"ok", "okay", "thanks", "thank you", "thx", "awesome", "great", "perfect", "cool", "nice"}
    return t in positives


def is_negative_ack(text: str) -> bool:
    t = (text or "").strip().lower()
    t = re.sub(r"[!.,?]+$", "", t)
    negatives = {"bad", "wrong", "not good", "terrible", "doesn't help", "useless", "not helpful"}
    return t in negatives



def is_dump_request(text: str) -> bool:
    t = (text or "").strip().lower()
    if not t:
        return False
    patterns = [
        "whole text",
        "all text",
        "full text",
        "entire text",
        "everything in the db",
        "everything in database",
        "all content",
        "dump the database",
        "show me all documents",
        "give me all documents",
        "give me the whole text which is saved in the db",
        "give me whole text which is saved in the db",
        "print the database",
        "show database",
    ]
    return any(p in t for p in patterns)


def has_uploaded_documents() -> bool:
    # Only count allowed file types
    for p in UPLOAD_DIR.glob("*"):
        if p.is_file() and p.suffix.lower() in ALLOWED_EXTS:
            return True
    return False



def word_count(text: str) -> int:
    return len((text or "").strip().split())

def truncate_to_words(text: str, max_words: int) -> str:
    words = (text or "").strip().split()
    if len(words) <= max_words:
        return text.strip()
    return " ".join(words[:max_words]).strip()


def user_requests_long_answer(text: str) -> bool:
    t = (text or "").lower()
    # simple patterns
    patterns = [
        "more than 250 words",
        "over 250 words",
        "at least 300 words",
        "300 words",
        "500 words",
        "1000 words",
        "very detailed",
        "as long as possible",
        "long answer",
        "in detail",
    ]
    return any(p in t for p in patterns)



# -----------------------------
# Pages
# -----------------------------





@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/upload", response_class=HTMLResponse)
async def upload_page(request: Request):
    return templates.TemplateResponse("upload.html", {"request": request})


@app.get("/chat", response_class=HTMLResponse)
async def chat_page(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})


# -----------------------------
# API: Upload + list
# -----------------------------
@app.post("/api/upload")
async def upload_files(
    files: list[UploadFile] = File(...),
    x_llm_base_url: str | None = Header(default=None),
    x_llm_api_key: str | None = Header(default=None),
    x_llm_model: str | None = Header(default=None),
    x_embed_model: str | None = Header(default=None),
):
    cfg = get_cfg_from_headers(x_llm_base_url, x_llm_api_key, x_llm_model, x_embed_model)

    saved = []
    for f in files:
        original_name = sanitize_filename(f.filename)
        ext = Path(original_name).suffix.lower()

        if ext not in ALLOWED_EXTS:
            raise HTTPException(status_code=400, detail=f"Unsupported file type: {f.filename}")

        stem = Path(original_name).stem
        unique_name = f"{stem}__{uuid.uuid4().hex}{ext}"
        out_path = UPLOAD_DIR / unique_name

        content = await f.read()
        out_path.write_bytes(content)

        saved.append({"original": original_name, "stored_as": unique_name, "size": len(content)})

    # Reindex into persistent ChromaDB using LOCAL embeddings (FastEmbed)
    rag_info = rag.index_new_files(embedding_model=cfg["embed_model"])
    return {"ok": True, "files": saved, "rag": rag_info}


@app.get("/api/uploads")
async def list_uploads():
    files_out = []
    for p in sorted(UPLOAD_DIR.glob("*"), key=lambda x: x.stat().st_mtime, reverse=True):
        if not p.is_file():
            continue
        files_out.append(
            {
                "stored_as": p.name,
                "original": derive_original_from_stored(p.name),
                "size": p.stat().st_size,
            }
        )
    return {"ok": True, "files": files_out}


@app.post("/api/reindex")
async def reindex(
    x_llm_base_url: str | None = Header(default=None),
    x_llm_api_key: str | None = Header(default=None),
    x_llm_model: str | None = Header(default=None),
    x_embed_model: str | None = Header(default=None),
):
    cfg = get_cfg_from_headers(x_llm_base_url, x_llm_api_key, x_llm_model, x_embed_model)
    info = rag.reindex(embedding_model=cfg["embed_model"])
    return {"ok": True, "rag": info}


# -----------------------------
# API: Export DB text (fast, no LLM)
# -----------------------------
@app.get("/api/db_text")
async def db_text(
    x_embed_model: str | None = Header(default=None),
    limit_chars: int = 200_000,
):
    """
    Export text stored in the vector DB (Chroma) as plain text.
    Bounded by limit_chars to avoid huge responses.
    """
    embed_model = x_embed_model or DEFAULT_EMBED_MODEL

    try:
        rag.load(embedding_model=embed_model)
        vs = rag.engine.vs  # Chroma instance

        if vs is None:
            return {"ok": False, "message": "Vector store not loaded."}

        col = getattr(vs, "_collection", None)
        if col is None:
            return {"ok": False, "message": "Could not access Chroma collection."}

        data = col.get(include=["documents", "metadatas"])
        docs = data.get("documents") or []
        metas = data.get("metadatas") or []
        if not docs:
            return {
                "ok": False,
                "message": "No documents indexed yet. Please upload documents first."
            }

        grouped: dict[str, list[str]] = {}
        for doc, meta in zip(docs, metas):
            meta = meta or {}
            display_file = meta.get("display_file") or meta.get("source_file") or "unknown"
            grouped.setdefault(display_file, []).append(doc or "")

        out_parts = []
        total = 0
        for fname, chunks in grouped.items():
            header = f"\n\n===== {fname} =====\n"
            if total + len(header) > limit_chars:
                break
            out_parts.append(header)
            total += len(header)

            for ch in chunks:
                if not ch:
                    continue
                if total + len(ch) + 1 > limit_chars:
                    out_parts.append("\n[TRUNCATED]\n")
                    total += len("\n[TRUNCATED]\n")
                    break
                out_parts.append(ch + "\n")
                total += len(ch) + 1

        return {
            "ok": True,
            "embed_model": embed_model,
            "files": list(grouped.keys()),
            "text": "".join(out_parts).strip(),
        }

    except Exception as e:
        return {"ok": False, "message": f"{type(e).__name__}: {e}"}



from typing import Literal

# -----------------------------
# API: Validate (CHAT ONLY)
# -----------------------------
class ValidateReq(BaseModel):
    base_url: str
    api_key: str | None = None
    test_prompt: str = "ping"

    # NEW
    provider: Literal["local", "external"] = "local"
    model: str | None = None 


class ValidateResp(BaseModel):
    ok: bool
    message: str
    chat_model: str | None = None


@app.post("/api/validate", response_model=ValidateResp)
async def validate(req: ValidateReq):
    """
    Validate chat only.
    - local: list /v1/models and auto-pick a chat model
    - external: use provided model if set; otherwise try to list models, then fallback to error
    """
    try:
        # IMPORTANT: for external, do NOT default to "lm-studio"
        api_key = req.api_key or ("lm-studio" if req.provider == "local" else None)
        if req.provider == "external" and not api_key:
            return ValidateResp(ok=False, message="API key is required for external LLM providers.")

        client = make_openai_client(req.base_url, api_key)

        chat_model = req.model  # if user provided
        model_ids: list[str] = []

        # Try to list models (works for LM Studio and most OpenAI-compatible providers)
        try:
            models = client.models.list()
            model_ids = [m.id for m in (models.data or []) if getattr(m, "id", None)]
        except Exception:
            # Some providers may block /models; that's OK if external model is provided
            model_ids = []

        if not chat_model:
            if not model_ids:
                return ValidateResp(
                    ok=False,
                    message="No model provided and /v1/models returned nothing. Please enter a model name."
                )
            chat_model = choose_chat_model(model_ids)
            if not chat_model:
                chat_model = model_ids[0]

        # Ping chat
        _ = client.chat.completions.create(
            model=chat_model,
            messages=[{"role": "user", "content": req.test_prompt}],
            max_tokens=8,
            temperature=0,
        )

        return ValidateResp(
            ok=True,
            message=f"Chat OK. Embeddings will use FastEmbed locally ({DEFAULT_EMBED_MODEL}).",
            chat_model=chat_model,
        )

    except ValueError as e:
        return ValidateResp(ok=False, message=str(e))
    except OpenAIError as e:
        return ValidateResp(ok=False, message=f"{type(e).__name__}: {e}")
    except Exception as e:
        return ValidateResp(ok=False, message=f"{type(e).__name__}: {e}")


# -----------------------------
# API: Chat (RAG)
# -----------------------------
class ChatReq(BaseModel):
    message: str
    system: str | None = None
    temperature: float = 0.7
    top_k: int = 5


class ChatResp(BaseModel):
    reply: str


@app.post("/api/chat", response_model=ChatResp)
async def chat_api(
    payload: ChatReq,
    x_llm_base_url: str | None = Header(default=None),
    x_llm_api_key: str | None = Header(default=None),
    x_llm_model: str | None = Header(default=None),
    x_embed_model: str | None = Header(default=None),
):
    cfg = get_cfg_from_headers(x_llm_base_url, x_llm_api_key, x_llm_model, x_embed_model)

    msg = (payload.message or "").strip()
    if not msg:
        return ChatResp(reply="Please type something 🙂")

    # Fast deterministic replies (no retrieval, no LLM)
    if is_greeting(msg):
        return ChatResp(reply="Hi! How can I help you today?")

    if is_positive_ack(msg):
        return ChatResp(reply="I'm glad you found the explanation helpful! ...")

    if is_negative_ack(msg):
        return ChatResp(reply="Please give me more information to give you correct or your desirable replies!")

    # Only now apply the short-message rule
    if len(msg) < 4:
        return ChatResp(reply="Please ask a question about your uploaded documents 🙂")

    if user_requests_long_answer(msg):
        return ChatResp(
            reply="I can answer, but I’m limited to 250 words per message. If you want a longer explanation, I can send it in multiple parts—just tell me: “continue”."
        )

    # No documents guard (perfect)
    if not has_uploaded_documents():
        return ChatResp(
            reply="I don't have any documents yet. Please upload PDF/TXT/DOCX files first (click Upload and then Proceed) and then ask your question 🙂"
        )

    

    # Fast "dump DB" (no LLM)
    if is_dump_request(msg):
        dump = await db_text(x_embed_model=cfg["embed_model"], limit_chars=200_000)
        if not dump.get("ok"):
            return ChatResp(reply=dump.get("message", "Could not dump DB text."))
        return ChatResp(reply=dump["text"])

    try:
        # Load ChromaDB + embeddings (local)
        rag.load(embedding_model=cfg["embed_model"])

        docs = rag.retrieve(msg, k=max(1, min(int(payload.top_k), 10)))
        messages = rag.build_prompt_messages(question=msg, docs=docs)

        client = make_openai_client(cfg["base_url"], cfg["api_key"])
        resp = client.chat.completions.create(
            model=cfg["chat_model"],
            messages=messages,
            temperature=payload.temperature,
            max_tokens=150,  # helps keep answers short + faster
        )

        content = resp.choices[0].message.content or ""
        reply = content.strip() or "(empty response)"

        max_words = 50
        wc = word_count(reply)
        if wc > max_words:
            reply = truncate_to_words(reply, max_words)
            # reply += "\n\n(I’m limited to  250 words per answer. If you want more, ask “continue” and I’ll respond in parts.)"

        return ChatResp(reply=reply)


    except OpenAIError as e:
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")



@app.post("/api/reset")
async def reset_all():
    """
    Deletes EVERYTHING:
      - app/uploads/*
      - app/rag_store/* (including chroma sqlite + wal/shm)
    Without restarting the server.
    """
    uploads_dir = UPLOAD_DIR
    rag_store_dir = RAG_STORE_DIR

    deleted_upload_files = 0

    # 1) HARD reset Chroma connections/caches FIRST (critical)
    try:
        rag.engine.hard_reset()
    except Exception:
        pass

    # 2) Delete uploaded files
    try:
        uploads_dir.mkdir(parents=True, exist_ok=True)
        for p in uploads_dir.glob("*"):
            if p.is_file():
                p.unlink(missing_ok=True)
                deleted_upload_files += 1
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to delete uploads: {type(e).__name__}: {e}")

    # 3) Delete rag_store entirely
    try:
        if rag_store_dir.exists() and rag_store_dir.is_dir():
            shutil.rmtree(rag_store_dir)

        # recreate required dirs
        (rag_store_dir / "chroma_db").mkdir(parents=True, exist_ok=True)

        # ensure writable (mac/linux)
        try:
            rag_store_dir.chmod(0o755)
            (rag_store_dir / "chroma_db").chmod(0o755)
        except Exception:
            pass

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to delete rag_store: {type(e).__name__}: {e}")

    return {
        "ok": True,
        "message": "All uploads and vector DB data deleted.",
        "deleted_upload_files": deleted_upload_files,
        "deleted_dir": str(rag_store_dir),
    }









class ExportChatPDFReq(BaseModel):
    title: str | None = "ECDA LLM Chat Export"
    meta: str | None = None
    messages: list[dict]  # each dict: {"role": "user"|"assistant", "text": "..."}


@app.post("/api/export_chat_pdf")
async def export_chat_pdf(payload: ExportChatPDFReq):
    buf = BytesIO()

    doc = SimpleDocTemplate(
        buf,
        pagesize=A4,
        leftMargin=36,
        rightMargin=36,
        topMargin=36,
        bottomMargin=36,
        title=payload.title or "Chat Export",
        author="ECDA LLM"
    )

    styles = getSampleStyleSheet()
    h = styles["Heading2"]
    normal = styles["BodyText"]

    story = []
    story.append(Paragraph(payload.title or "Chat Export", h))

    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    story.append(Paragraph(f"Generated: {ts}", normal))

    if payload.meta:
        story.append(Spacer(1, 8))
        safe_meta = (
            (payload.meta or "")
            .replace("&", "&amp;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
            .replace("\n", "<br/>")
        )
        story.append(Paragraph(safe_meta, normal))


    story.append(Spacer(1, 12))

    # Messages
    for i, m in enumerate(payload.messages or [], start=1):
        role = (m.get("role") or "").strip().lower()
        text = (m.get("text") or "").strip()

        if not text:
            continue

        label = "User" if role == "user" else "Assistant"
        # Escape basic HTML in text to avoid Paragraph parsing issues
        safe = (
            text.replace("&", "&amp;")
                .replace("<", "&lt;")
                .replace(">", "&gt;")
                .replace("\n", "<br/>")
        )

        story.append(Paragraph(f"<b>{i}. {label}:</b>", normal))
        story.append(Paragraph(safe, normal))
        story.append(Spacer(1, 10))

    doc.build(story)
    buf.seek(0)

    filename = f"ecda_chat_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
    headers = {"Content-Disposition": f'attachment; filename="{filename}"'}

    return StreamingResponse(buf, media_type="application/pdf", headers=headers)


