from __future__ import annotations

import json
import math
import re
import socket
import time
import uuid
from collections import Counter
from datetime import datetime, timezone
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlparse

import ollama


ROOT = Path(__file__).parent
DATA_DIR = ROOT / "data"
CHATS_FILE = DATA_DIR / "chats.json"
DEFAULT_MODEL = "qwen2.5:7b"
AVAILABLE_MODELS = ["qwen2.5:7b", "llama3.1:8b", "qwen2.5:14b", "qwen2.5:3b"]
HOST = "0.0.0.0"
PORT = 8787
RAG_EXTENSIONS = {".md", ".txt", ".py", ".js", ".html", ".css", ".ps1"}
RAG_EXCLUDED_DIRS = {".git", ".venv", "__pycache__", "data"}

rag_chunks: list[dict] = []
rag_indexed_at = ""


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def read_store() -> dict:
    if not CHATS_FILE.exists():
        return {"chats": []}

    try:
        return json.loads(CHATS_FILE.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {"chats": []}


def write_store(store: dict) -> None:
    DATA_DIR.mkdir(exist_ok=True)
    CHATS_FILE.write_text(json.dumps(store, ensure_ascii=False, indent=2), encoding="utf-8")


def chat_title(messages: list[dict]) -> str:
    for message in messages:
        if message.get("role") == "user" and message.get("content"):
            title = " ".join(str(message["content"]).split())
            return title[:40] or "New chat"
    return "New chat"


def installed_models() -> list[str]:
    try:
        response = ollama.list()
    except Exception:
        return AVAILABLE_MODELS

    names: list[str] = []
    for model in response.get("models", []):
        if isinstance(model, dict):
            name = model.get("name") or model.get("model")
        else:
            name = getattr(model, "name", None) or getattr(model, "model", None)
        if name:
            names.append(name)

    ordered = [name for name in AVAILABLE_MODELS if name in names]
    ordered.extend(name for name in names if name not in ordered)
    return ordered or AVAILABLE_MODELS


def tokenize(text: str) -> list[str]:
    return [token for token in re.findall(r"\w+", text.lower(), flags=re.UNICODE) if len(token) > 1]


def iter_rag_files() -> list[Path]:
    files: list[Path] = []
    for path in ROOT.rglob("*"):
        if any(part in RAG_EXCLUDED_DIRS for part in path.relative_to(ROOT).parts):
            continue
        if path.is_file() and path.suffix.lower() in RAG_EXTENSIONS:
            files.append(path)
    return sorted(files)


def chunk_text(path: Path, text: str, size: int = 1800, overlap: int = 250) -> list[dict]:
    chunks: list[dict] = []
    start = 0
    rel = path.relative_to(ROOT).as_posix()
    while start < len(text):
        end = min(start + size, len(text))
        chunk = text[start:end].strip()
        if chunk:
            chunks.append({"source": rel, "text": chunk, "tokens": Counter(tokenize(chunk))})
        if end == len(text):
            break
        start = max(0, end - overlap)
    return chunks


def build_rag_index() -> dict:
    global rag_chunks, rag_indexed_at

    chunks: list[dict] = []
    for path in iter_rag_files():
        try:
            text = path.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue
        chunks.extend(chunk_text(path, text))

    rag_chunks = chunks
    rag_indexed_at = now_iso()
    return rag_status()


def rag_status() -> dict:
    return {
        "chunks": len(rag_chunks),
        "indexedAt": rag_indexed_at,
        "files": len({chunk["source"] for chunk in rag_chunks}),
    }


def search_rag(query: str, limit: int = 5) -> list[dict]:
    if not rag_chunks:
        build_rag_index()

    query_tokens = Counter(tokenize(query))
    if not query_tokens:
        return []

    scored: list[tuple[float, dict]] = []
    total = max(1, len(rag_chunks))
    doc_freq = Counter()
    for chunk in rag_chunks:
        for token in chunk["tokens"]:
            doc_freq[token] += 1

    for chunk in rag_chunks:
        score = 0.0
        chunk_tokens = chunk["tokens"]
        for token, query_count in query_tokens.items():
            if token not in chunk_tokens:
                continue
            idf = math.log((total + 1) / (doc_freq[token] + 1)) + 1
            score += min(query_count, chunk_tokens[token]) * idf
        if score:
            scored.append((score, chunk))

    scored.sort(key=lambda item: item[0], reverse=True)
    return [
        {"source": chunk["source"], "text": chunk["text"], "score": round(score, 3)}
        for score, chunk in scored[:limit]
    ]


class Handler(SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=str(ROOT), **kwargs)

    def do_GET(self) -> None:
        path = urlparse(self.path).path

        if path == "/api/models":
            self._send_json(200, {"defaultModel": DEFAULT_MODEL, "models": installed_models()})
            return

        if path == "/api/rag/status":
            if not rag_chunks:
                build_rag_index()
            self._send_json(200, rag_status())
            return

        if path == "/api/chats":
            store = read_store()
            chats = [
                {
                    "id": chat["id"],
                    "title": chat.get("title", "New chat"),
                    "updatedAt": chat.get("updatedAt", ""),
                    "model": chat.get("model", DEFAULT_MODEL),
                    "useRag": chat.get("useRag", True),
                }
                for chat in store.get("chats", [])
            ]
            self._send_json(200, {"chats": chats})
            return

        if path.startswith("/api/chats/"):
            chat_id = path.rsplit("/", 1)[-1]
            store = read_store()
            chat = next((item for item in store.get("chats", []) if item["id"] == chat_id), None)
            if not chat:
                self._send_json(404, {"error": "Chat not found"})
                return

            self._send_json(200, {"chat": chat})
            return

        super().do_GET()

    def do_POST(self) -> None:
        path = urlparse(self.path).path
        if path == "/api/rag/reindex":
            self._send_json(200, build_rag_index())
            return

        if path != "/api/chat":
            self.send_error(404)
            return

        length = int(self.headers.get("content-length", 0))
        data = json.loads(self.rfile.read(length) or b"{}")
        chat_id = data.get("chatId")
        messages = data.get("messages", [])
        model = data.get("model") or DEFAULT_MODEL
        use_rag = bool(data.get("useRag", True))

        try:
            sources = search_rag(messages[-1]["content"]) if use_rag and messages else []
            system_prompt = (
                "You are a concise local assistant. Reply in Korean unless the user asks otherwise. "
                "If reference context is provided, use it and mention source filenames when useful. "
                "If the context is not enough, say so clearly."
            )
            if sources:
                context = "\n\n".join(
                    f"[{index + 1}] {item['source']}\n{item['text']}" for index, item in enumerate(sources)
                )
                system_prompt += f"\n\nReference context:\n{context}"

            started = time.perf_counter()
            response = ollama.chat(
                model=model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    *messages,
                ],
            )
            elapsed_ms = round((time.perf_counter() - started) * 1000)
            reply = response["message"]["content"]
            saved_messages = [*messages, {"role": "assistant", "content": reply}]
            chat = self._save_chat(chat_id, saved_messages, model, use_rag)
            payload = {
                "chatId": chat["id"],
                "reply": reply,
                "chat": chat,
                "model": model,
                "elapsedMs": elapsed_ms,
                "sources": [{"source": item["source"], "score": item["score"]} for item in sources],
            }
            self._send_json(200, payload)
        except Exception as exc:
            self._send_json(500, {"error": str(exc)})

    def do_DELETE(self) -> None:
        path = urlparse(self.path).path
        if not path.startswith("/api/chats/"):
            self.send_error(404)
            return

        chat_id = path.rsplit("/", 1)[-1]
        store = read_store()
        store["chats"] = [item for item in store.get("chats", []) if item["id"] != chat_id]
        write_store(store)
        self._send_json(200, {"ok": True})

    def _save_chat(self, chat_id: str | None, messages: list[dict], model: str, use_rag: bool) -> dict:
        store = read_store()
        chats = store.setdefault("chats", [])
        chat = next((item for item in chats if item["id"] == chat_id), None)

        if not chat:
            chat = {
                "id": str(uuid.uuid4()),
                "createdAt": now_iso(),
                "title": chat_title(messages),
                "messages": [],
            }
            chats.insert(0, chat)

        chat["title"] = chat_title(messages)
        chat["messages"] = messages
        chat["model"] = model
        chat["useRag"] = use_rag
        chat["updatedAt"] = now_iso()
        chats.sort(key=lambda item: item.get("updatedAt", ""), reverse=True)
        write_store(store)
        return chat

    def _send_json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("content-type", "application/json; charset=utf-8")
        self.send_header("content-length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)


def main() -> None:
    build_rag_index()
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    local_ip = socket.gethostbyname(socket.gethostname())
    print(f"LLM web is running locally at http://127.0.0.1:{PORT}")
    print(f"LAN access is available at http://{local_ip}:{PORT}")
    server.serve_forever()


if __name__ == "__main__":
    main()