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

# 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.
_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_PATH = Path(__file__).parent / "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 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)