import json
import os
import re
import ssl
import time
import urllib.request
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

import certifi

# CUDA 13.x moved its runtime DLLs to bin/x64, but llama_cpp's loader only
# checks CUDA_PATH/bin, so point it at the right folder before importing.
# Only relevant when running from source against a system CUDA toolkit - a
# packaged build ships the runtime DLLs next to llama_cpp's own, so this is a
# no-op there (CUDA_PATH won't be set on a machine with just an NVIDIA driver).
_cuda_path = os.environ.get("CUDA_PATH")
if _cuda_path:
    _cuda_bin_x64 = Path(_cuda_path) / "bin" / "x64"
    if _cuda_bin_x64.is_dir():
        os.environ["PATH"] = str(_cuda_bin_x64) + os.pathsep + os.environ["PATH"]

from ddgs import DDGS  # noqa: E402
from llama_cpp import Llama  # noqa: E402

# Model lives outside the install/source directory - a packaged build's own
# folder may be read-only (Program Files) or discarded on update, and the
# model is far too large to embed in an installer, so it's fetched into the
# user's profile on first run instead.
MODEL_FILENAME = "Qwen3.5-27B-IQ4_XS.gguf"
MODELS_DIR = Path.home() / ".dev_assist_chatbot" / "models"
MODEL_PATH = MODELS_DIR / MODEL_FILENAME
MODEL_DOWNLOAD_URL = (
    "https://huggingface.co/unsloth/Qwen3.5-27B-GGUF/resolve/main/Qwen3.5-27B-IQ4_XS.gguf"
)

SYSTEM_PROMPT = (
    "You are a knowledgeable developer assistant. Answer questions about "
    "programming languages, libraries, frameworks, tools, debugging, and "
    "software design with clear, accurate, and concise explanations. Use "
    "code snippets when they help illustrate an answer, and always place "
    "them in fenced Markdown code blocks with a language tag right after "
    "the opening backticks (e.g. ```python), never inline in a paragraph "
    "or left unfenced. Use Markdown throughout: numbered or bulleted "
    "lists for steps, **bold** for emphasis, and inline `code` spans for "
    "identifiers, filenames, and commands mentioned in prose. Use the "
    "web_search tool when a question depends on current information you "
    "may not know (recent releases, changelogs, current best practices, "
    "error messages you don't recognize)."
)

# Number of prior user/assistant exchanges to keep in the prompt. Bounds the
# prompt size so a long chat session doesn't overflow the model's context.
MAX_HISTORY_TURNS = 10

# Caps how many times the model can call a tool while answering a single
# question, so a confused model can't loop searches forever.
MAX_TOOL_ROUNDS = 3

WEB_SEARCH_TOOL = {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "Search the web and return brief result snippets.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The search query."},
            },
            "required": ["query"],
        },
    },
}

# Matches the <tool_call><function=name>...</function></tool_call> blocks
# Qwen3.5's chat template trains it to emit when it wants to invoke a tool -
# an XML-style syntax, not Qwen2.5's JSON-in-<tool_call> format.
TOOL_CALL_RE = re.compile(
    r"<tool_call>\s*<function=([^>]+)>(.*?)</function>\s*</tool_call>", re.DOTALL
)
# Matches each <parameter=name>value</parameter> inside a tool call's body.
PARAMETER_RE = re.compile(r"<parameter=([^>]+)>\n?(.*?)\n?</parameter>", re.DOTALL)

# Qwen3.5 thinks by default, wrapping its reasoning trace in these tags
# before the actual answer (or tool call). _route_stream_piece splits a
# streamed response between the two so a caller can render them separately.
THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"


def web_search(query: str, max_results: int = 5) -> str:
    # ddgs is a metasearch client and defaults to aggregating several backends
    # (Mojeek, Brave, ...); pin it to duckduckgo specifically, since some of
    # the others time out from this machine. The duckduckgo backend itself
    # occasionally returns zero results transiently, so retry once.
    for attempt in range(2):
        try:
            results = DDGS(timeout=10).text(query, max_results=max_results, backend="duckduckgo")
            break
        except Exception as e:
            if attempt == 1:
                return f"Search failed: {e}"
            time.sleep(1)
    if not results:
        return "No results found."
    return "\n\n".join(f"{r['title']}\n{r['href']}\n{r['body']}" for r in results)


TOOLS = {"web_search": web_search}


def extract_tool_calls(text: str) -> list[dict]:
    calls = []
    for call_match in TOOL_CALL_RE.finditer(text):
        name = call_match.group(1).strip()
        arguments = {
            param_match.group(1).strip(): param_match.group(2)
            for param_match in PARAMETER_RE.finditer(call_match.group(2))
        }
        calls.append({"name": name, "arguments": arguments})
    return calls


# Chat sessions are persisted outside the repo, one JSON file per session, so
# history survives closing the app and isn't lost on an app update/reinstall.
CHATS_DIR = Path.home() / ".dev_assist_chatbot" / "chats"

# Placeholder title for a session that hasn't received its first message yet;
# used to detect when a title still needs to be derived from that message.
DEFAULT_TITLE = "New chat"


@dataclass
class ChatSession:
    id: str
    title: str
    created_at: str
    updated_at: str
    history: list[tuple[str, str]] = field(default_factory=list)


def new_session() -> ChatSession:
    now = datetime.now(timezone.utc).isoformat()
    return ChatSession(id=uuid.uuid4().hex, title=DEFAULT_TITLE, created_at=now, updated_at=now)


def derive_title(question: str, max_len: int = 40) -> str:
    """First line of a session's opening question, trimmed to max_len - used
    as its display title once it gets its first message."""
    title = question.strip().splitlines()[0]
    return title if len(title) <= max_len else title[: max_len - 1].rstrip() + "…"


def _session_path(session_id: str) -> Path:
    return CHATS_DIR / f"{session_id}.json"


def save_session(session: ChatSession) -> None:
    session.updated_at = datetime.now(timezone.utc).isoformat()
    CHATS_DIR.mkdir(parents=True, exist_ok=True)
    payload = {
        "id": session.id,
        "title": session.title,
        "created_at": session.created_at,
        "updated_at": session.updated_at,
        "history": [[question, answer] for question, answer in session.history],
    }
    _session_path(session.id).write_text(json.dumps(payload, indent=2), encoding="utf-8")


def delete_session(session_id: str) -> None:
    _session_path(session_id).unlink(missing_ok=True)


def list_sessions() -> list[ChatSession]:
    """All persisted chat sessions, most recently updated first. A session
    file that fails to parse is skipped rather than failing the whole list."""
    if not CHATS_DIR.is_dir():
        return []
    sessions = []
    for path in CHATS_DIR.glob("*.json"):
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            sessions.append(
                ChatSession(
                    id=data["id"],
                    title=data["title"],
                    created_at=data["created_at"],
                    updated_at=data["updated_at"],
                    history=[tuple(turn) for turn in data["history"]],
                )
            )
        except (json.JSONDecodeError, OSError, KeyError, TypeError):
            continue
    sessions.sort(key=lambda s: s.updated_at, reverse=True)
    return sessions


def build_messages(history: list[tuple[str, str]], question: str) -> list[dict[str, str]]:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_turn, assistant_turn in history[-MAX_HISTORY_TURNS:]:
        messages.append({"role": "user", "content": user_turn})
        messages.append({"role": "assistant", "content": assistant_turn})
    messages.append({"role": "user", "content": question})
    return messages


def download_model(on_progress: Callable[[int, int], None] | None = None) -> None:
    """Fetch MODEL_PATH from MODEL_DOWNLOAD_URL if it isn't already present.
    Downloads to a .part file alongside the target and renames on success, so
    a crash or interrupted download can't leave a truncated file that
    load_model() would then try (and fail) to use. If given, on_progress is
    invoked with (bytes_downloaded, total_bytes) as each chunk arrives -
    total_bytes is 0 if the server doesn't report a Content-Length."""
    if MODEL_PATH.exists():
        return
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    part_path = MODEL_PATH.with_suffix(MODEL_PATH.suffix + ".part")

    request = urllib.request.Request(
        MODEL_DOWNLOAD_URL, headers={"User-Agent": "dev-assist-chatbot"}
    )
    # A fresh/sandboxed Windows install may not have fetched the intermediate
    # CA cert for huggingface.co's chain yet (normally done transparently via
    # Windows Update on first need, which a locked-down or offline sandbox
    # can't reach) - use certifi's bundled CA list instead of the OS trust
    # store so this doesn't depend on that happening.
    ssl_context = ssl.create_default_context(cafile=certifi.where())
    with urllib.request.urlopen(request, context=ssl_context) as response:
        total = int(response.headers.get("Content-Length", 0))
        downloaded = 0
        with open(part_path, "wb") as f:
            while chunk := response.read(1024 * 1024):
                f.write(chunk)
                downloaded += len(chunk)
                if on_progress is not None:
                    on_progress(downloaded, total)

    part_path.rename(MODEL_PATH)


def load_model() -> Llama:
    return Llama(
        model_path=str(MODEL_PATH),
        # Qwen3.5-27B's hybrid attention keeps KV cache small (most layers
        # carry a fixed-size state rather than one that grows with context),
        # so this fits alongside the IQ4_XS weights on a 16GB card with room
        # to spare - unlike a same-sized dense-attention model.
        n_ctx=32768,
        # -1 offloads every layer to the GPU
        n_gpu_layers=-1,
        # number of tokens in the prompt that are fed into the model at a time
        n_batch=1024,
        verbose=False,
    )


def _route_stream_piece(
    piece: str,
    state: dict,
    on_token: Callable[[str], None] | None,
    on_thinking_token: Callable[[str], None] | None,
) -> None:
    """Forward piece to on_thinking_token while inside a <think>...</think>
    block and to on_token otherwise, using state (a per-call {"in_think",
    "buffer"} dict) to track which side of a tag boundary we're on and to
    hold back text that might be the start of a tag split across pieces."""
    buffer = state["buffer"] + piece
    while True:
        tag = THINK_CLOSE if state["in_think"] else THINK_OPEN
        idx = buffer.find(tag)
        if idx == -1:
            break
        before, buffer = buffer[:idx], buffer[idx + len(tag) :]
        callback = on_thinking_token if state["in_think"] else on_token
        if before and callback is not None:
            callback(before)
        state["in_think"] = not state["in_think"]

    # The buffer may end with a prefix of the tag we're now watching for
    # (e.g. a chunk boundary landing mid-"</think>"), so hold that part back
    # rather than misrouting it as ordinary content.
    tag = THINK_CLOSE if state["in_think"] else THINK_OPEN
    keep = 0
    for length in range(min(len(tag), len(buffer)), 0, -1):
        if tag.startswith(buffer[-length:]):
            keep = length
            break
    emit, state["buffer"] = buffer[: len(buffer) - keep], buffer[len(buffer) - keep :]
    callback = on_thinking_token if state["in_think"] else on_token
    if emit and callback is not None:
        callback(emit)


def _stream_content(
    llm: Llama,
    messages: list[dict],
    tools: list[dict] | None,
    on_token: Callable[[str], None] | None,
    on_thinking_token: Callable[[str], None] | None = None,
) -> str:
    """Run one streaming create_chat_completion call, forwarding each piece
    of generated text to on_token (or on_thinking_token, while inside a
    <think> block) as it arrives, and return the full raw text, thinking
    tags included, so the caller's message history matches what the model
    actually produced."""
    kwargs = {}
    if tools is not None:
        kwargs["tools"] = tools
    stream = llm.create_chat_completion(
        messages,
        max_tokens=8192,
        temperature=0.4,
        top_p=0.9,
        repeat_penalty=1.1,
        stream=True,
        **kwargs,
    )
    content = ""
    # Qwen3.5's chat template appends "<think>\n" to the prompt itself before
    # generation starts (see its tokenizer.chat_template's add_generation_prompt
    # block), so the model's own streamed output never contains that opening
    # tag - it starts already "inside" one, and only emits "</think>" once
    # reasoning is done.
    route_state = {"in_think": True, "buffer": ""}
    for chunk in stream:
        piece = chunk["choices"][0]["delta"].get("content")
        if piece:
            content += piece
            _route_stream_piece(piece, route_state, on_token, on_thinking_token)
    if route_state["buffer"]:
        callback = on_thinking_token if route_state["in_think"] else on_token
        if callback is not None:
            callback(route_state["buffer"])
    return content


def _strip_thinking(text: str) -> str:
    """Drop a leading reasoning trace from raw model output, the same way
    the chat template splits it back out when re-rendering a past turn:
    everything before the first </think> is reasoning (there's no opening
    tag to look for - see _stream_content), everything after is the answer."""
    before, sep, after = text.partition(THINK_CLOSE)
    return after.lstrip("\n") if sep else before


def ask(
    llm: Llama,
    messages: list[dict],
    on_tool_call: Callable[[str], None] | None = None,
    on_usage: Callable[[dict], None] | None = None,
    on_token: Callable[[str], None] | None = None,
    on_thinking_token: Callable[[str], None] | None = None,
) -> str:
    """Run the chat loop for one question, letting the model call web_search
    up to MAX_TOOL_ROUNDS times before returning its final answer text.
    If given, on_tool_call is invoked with the search query each time a tool
    is about to run, so a caller (CLI, GUI, ...) can surface progress.
    If given, on_usage is invoked with a {total_tokens} usage dict after each
    model call, so a caller can track how much of the context window is
    filled. If given, on_token is invoked with each piece of generated text
    as it streams from the model, so a caller can render the answer live
    instead of waiting for it to finish. If given, on_thinking_token is
    invoked the same way but for text inside the model's <think> block,
    letting a caller render its reasoning trace separately from the answer."""
    for _ in range(MAX_TOOL_ROUNDS):
        content = _stream_content(llm, messages, [WEB_SEARCH_TOOL], on_token, on_thinking_token)
        if on_usage is not None:
            on_usage({"total_tokens": llm.n_tokens})
        calls = extract_tool_calls(content)
        if not calls:
            return _strip_thinking(content)

        messages.append({"role": "assistant", "content": content})
        for call in calls:
            name = call.get("name")
            args = call.get("arguments") or {}
            if isinstance(args, str):
                args = json.loads(args)
            func = TOOLS.get(name)
            if func is None:
                result = f"Unknown tool: {name}"
            else:
                if on_tool_call is not None:
                    on_tool_call(args.get("query", ""))
                result = func(**args)
            messages.append({"role": "tool", "content": result})

    # Ran out of tool-call rounds without a plain answer; ask one more time
    # without offering the tool, forcing it to answer from what it has
    # instead of returning a leftover tool-call as if it were the answer.
    content = _stream_content(llm, messages, None, on_token, on_thinking_token)
    if on_usage is not None:
        on_usage({"total_tokens": llm.n_tokens})
    return _strip_thinking(content)