import json
import os
import re
import time
from collections.abc import Callable
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 / "Qwen2.5-Coder-14B-Instruct-Q5_K_M.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. 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>{...}</tool_call> blocks Qwen's chat template trains
# it to emit when it wants to invoke a tool.
TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)


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 _iter_balanced_braces(text: str):
    """Yield each top-level {...} substring of text, respecting nesting."""
    depth = 0
    start = None
    for i, ch in enumerate(text):
        if ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}" and depth > 0:
            depth -= 1
            if depth == 0 and start is not None:
                yield text[start : i + 1]


def extract_tool_calls(text: str) -> list[dict]:
    calls = []
    for match in TOOL_CALL_RE.finditer(text):
        try:
            calls.append(json.loads(match.group(1)))
        except json.JSONDecodeError:
            continue
    if calls:
        return calls

    # This model often skips the <tool_call> wrapper it's instructed to use
    # and just emits a bare JSON object instead - sometimes with leading
    # prose alongside it (e.g. "Let's try again. {...}") - so scan for any
    # JSON object in the text rather than requiring the whole message to be one.
    for candidate in _iter_balanced_braces(text):
        try:
            call = json.loads(candidate)
        except json.JSONDecodeError:
            continue
        if isinstance(call, dict) and call.get("name") in TOOLS:
            calls.append(call)
    return calls


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),
        # the model supports up to 32k context natively; give plenty of room
        # for chat history plus a long generated response
        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 ask(
    llm: Llama,
    messages: list[dict],
    on_tool_call: Callable[[str], None] | None = None,
    on_usage: Callable[[dict], 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 the {prompt_tokens, completion_tokens,
    total_tokens} usage dict after each model call, so a caller can track how
    much of the context window is filled."""
    for _ in range(MAX_TOOL_ROUNDS):
        response = llm.create_chat_completion(
            messages,
            tools=[WEB_SEARCH_TOOL],
            max_tokens=8192,
            temperature=0.4,
            top_p=0.9,
            repeat_penalty=1.1,
        )
        if on_usage is not None:
            on_usage(response["usage"])
        content = response["choices"][0]["message"]["content"] or ""
        calls = extract_tool_calls(content)
        if not calls:
            return 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.
    response = llm.create_chat_completion(
        messages,
        max_tokens=8192,
        temperature=0.4,
        top_p=0.9,
        repeat_penalty=1.1,
    )
    if on_usage is not None:
        on_usage(response["usage"])
    return response["choices"][0]["message"]["content"] or ""