"""LLM interaction, tool dispatch, and memory management."""

import os
import re
import json
import urllib.request
import urllib.error


def log(msg):
    print(f"[tools] {msg}", flush=True)


def set_logger(fn):
    global log
    log = fn


# ─── Memory ──────────────────────────────────────────────────────────

def load_memory(workspace):
    """Load all memory files from workspace."""
    memory = ""
    try:
        with open(f"{workspace}/MEMORY.md") as f:
            memory = f.read().strip()
        mem_dir = f"{workspace}/memory"
        if os.path.isdir(mem_dir):
            for fname in sorted(os.listdir(mem_dir)):
                if fname.endswith(".md"):
                    try:
                        with open(f"{mem_dir}/{fname}") as f:
                            topic = fname.replace(".md", "")
                            memory += f"\n\n## {topic}\n{f.read().strip()}"
                    except Exception:
                        pass
    except FileNotFoundError:
        pass
    return memory


# ─── Tool Call Parsing ───────────────────────────────────────────────

def try_parse_tool_call(text):
    """Parse text-based tool calls (model dumps JSON as text)."""
    text = re.sub(r"</?tool_call>", "", text).strip()
    for start in range(len(text)):
        if text[start] == "{":
            for end in range(len(text), start, -1):
                if text[end - 1] == "}":
                    try:
                        obj = json.loads(text[start:end])
                        name = obj.get("name")
                        args = obj.get("arguments", {})
                        if name and isinstance(args, dict):
                            return (name, args)
                    except json.JSONDecodeError:
                        continue
    return None


# ─── LLM Interaction ────────────────────────────────────────────────

def ollama_request(ollama_url, payload):
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{ollama_url}/api/chat",
        data=data,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        return json.loads(resp.read(2_000_000))


def query_ollama(messages, runtime, tools, skill_scripts, dispatch_fn, ollama_url, max_rounds):
    """Call Ollama chat API with skill-based tool support."""
    payload = {
        "model": runtime["model"],
        "messages": messages,
        "stream": False,
        "options": {"num_predict": 512},
    }

    if tools:
        payload["tools"] = tools

    for round_num in range(max_rounds):
        remaining = max_rounds - round_num
        try:
            data = ollama_request(ollama_url, payload)
        except (urllib.error.URLError, TimeoutError) as e:
            return f"[error: {e}]"

        msg = data.get("message", {})

        # Structured tool calls
        tool_calls = msg.get("tool_calls")
        if tool_calls:
            messages.append(msg)
            for tc in tool_calls:
                fn = tc.get("function", {})
                result = dispatch_fn(
                    fn.get("name", ""),
                    fn.get("arguments", {}),
                    round_num + 1,
                )
                if remaining <= 2:
                    result += f"\n[warning: {remaining - 1} tool rounds remaining — wrap up]"
                messages.append({"role": "tool", "content": result})
            payload["messages"] = messages
            continue

        # Text-based tool calls
        content = msg.get("content", "").strip()
        parsed_tool = try_parse_tool_call(content)
        if parsed_tool:
            fn_name, fn_args = parsed_tool
            if fn_name in skill_scripts:
                messages.append({"role": "assistant", "content": content})
                result = dispatch_fn(fn_name, fn_args, round_num + 1)
                if remaining <= 2:
                    result += f"\n[warning: {remaining - 1} tool rounds remaining — wrap up]"
                messages.append({
                    "role": "user",
                    "content": f"Tool result:\n{result}\n\nNow respond to the user based on this result.",
                })
                payload["messages"] = messages
                continue

        return content

    return "[max tool rounds reached]"