"""A deterministic troubleshooting agent with bounded, testable tools."""

from __future__ import annotations

import math
import queue
import re
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TypeVar, cast

PLAYBOOKS = {
    "network": "Check the timeout, retry transient failures with backoff, and log the final status code.",
    "dependency": "Recreate the environment, pin the dependency, and record the failing version.",
    "general": "Reproduce the issue with the smallest input, capture the error, and add a regression test.",
}


@dataclass(frozen=True)
class Action:
    name: str
    arguments: dict[str, str] = field(default_factory=dict)


@dataclass
class Task:
    request: str
    observations: list[str] = field(default_factory=list)
    trace: list[str] = field(default_factory=list)
    answer: str | None = None


def classify_issue(text: str) -> str:
    if type(text) is not str:
        raise TypeError("text must be a string")

    lowered = text.casefold()
    if any(
        re.search(pattern, lowered)
        for pattern in (
            r"\btimeouts?\b",
            r"\btime(?:d)?\s+out\b",
            r"\bnetworks?\b",
            r"\bapis?\b",
        )
    ):
        return "network"
    if any(
        re.search(pattern, lowered)
        for pattern in (
            r"\bimports?\b",
            r"\bpackages?\b",
            r"\bdependenc(?:y|ies)\b",
        )
    ):
        return "dependency"
    return "general"


def lookup_playbook(category: str) -> str:
    if type(category) is not str:
        raise TypeError("category must be a string")
    if category not in PLAYBOOKS:
        raise ValueError(f"Unknown category: {category}")
    return PLAYBOOKS[category]


TOOLS: dict[str, Callable[..., str]] = {
    "classify_issue": classify_issue,
    "lookup_playbook": lookup_playbook,
}

EXPECTED_ARGUMENTS = {
    "classify_issue": {"text": str},
    "lookup_playbook": {"category": str},
}

T = TypeVar("T")


class _CallTimedOut(Exception):
    """Internal signal used to distinguish call timeouts from user errors."""


def _validate_action(action: object) -> Action:
    if not isinstance(action, Action):
        raise TypeError("Planner must return an Action")
    if type(action.name) is not str:
        raise TypeError("Action name must be a string")
    if type(action.arguments) is not dict:
        raise TypeError("Action arguments must be a dictionary")
    if any(type(key) is not str for key in action.arguments):
        raise TypeError("Action argument names must be strings")
    return action


def _validate_arguments(
    action_name: str,
    arguments: dict[str, object],
    schema: dict[str, type],
) -> None:
    if set(arguments) != set(schema):
        raise ValueError(f"Invalid arguments for {action_name}")
    for name, expected_type in schema.items():
        if type(arguments[name]) is not expected_type:
            raise TypeError(
                f"Argument '{name}' for {action_name} must be "
                f"{expected_type.__name__}"
            )


def _validate_task(task: object) -> Task:
    if not isinstance(task, Task):
        raise TypeError("task must be a Task")
    if type(task.request) is not str:
        raise TypeError("Task request must be a string")
    if type(task.observations) is not list or any(
        type(item) is not str for item in task.observations
    ):
        raise TypeError("Task observations must be a list of strings")
    if type(task.trace) is not list or any(
        type(item) is not str for item in task.trace
    ):
        raise TypeError("Task trace must be a list of strings")
    if task.answer is not None and type(task.answer) is not str:
        raise TypeError("Task answer must be a string or None")
    return task


def _validate_timeout(name: str, value: float | None) -> float | None:
    if value is None:
        return None
    if type(value) not in (int, float) or not math.isfinite(value) or value <= 0:
        raise ValueError(f"{name} must be a positive finite number or None")
    return float(value)


def _call_with_timeout(
    callback: Callable[[], T], timeout_seconds: float | None
) -> T:
    """Return on timeout even if a provider or tool call remains blocked.

    Python threads cannot be forcibly stopped. A timed-out callback continues in
    a daemon thread, so real providers and tools still need their own transport
    timeouts to cancel network or subprocess work and prevent external effects.
    """
    if timeout_seconds is None:
        return callback()

    results: queue.Queue[tuple[bool, object]] = queue.Queue(maxsize=1)

    def invoke() -> None:
        try:
            results.put((True, callback()))
        except BaseException as error:
            results.put((False, error))

    worker = threading.Thread(target=invoke, daemon=True)
    worker.start()
    worker.join(timeout_seconds)
    if worker.is_alive():
        raise _CallTimedOut

    succeeded, result = results.get_nowait()
    if not succeeded:
        raise cast(BaseException, result)
    return cast(T, result)


def _bounded_call(
    callback: Callable[[], T],
    *,
    label: str,
    call_timeout_seconds: float | None,
    deadline_at: float | None,
    deadline_seconds: float | None,
) -> T:
    remaining = None if deadline_at is None else deadline_at - time.monotonic()
    if remaining is not None and remaining <= 0:
        raise TimeoutError(
            f"Agent exceeded the {deadline_seconds:g}-second wall-clock deadline"
        )

    effective_timeout = call_timeout_seconds
    deadline_is_limit = False
    if remaining is not None and (
        effective_timeout is None or remaining <= effective_timeout
    ):
        effective_timeout = remaining
        deadline_is_limit = True

    try:
        result = _call_with_timeout(callback, effective_timeout)
    except _CallTimedOut:
        if deadline_is_limit:
            raise TimeoutError(
                f"Agent exceeded the {deadline_seconds:g}-second wall-clock deadline"
            ) from None
        raise TimeoutError(
            f"{label} exceeded its {call_timeout_seconds:g}-second timeout"
        ) from None

    if deadline_at is not None and time.monotonic() >= deadline_at:
        raise TimeoutError(
            f"Agent exceeded the {deadline_seconds:g}-second wall-clock deadline"
        )
    return result


def decide(task: Task) -> Action:
    """Stand in for a model so every action stays free and deterministic."""
    if len(task.observations) == 0:
        return Action("classify_issue", {"text": task.request})
    if len(task.observations) == 1:
        return Action("lookup_playbook", {"category": task.observations[0]})
    return Action("finish", {"answer": task.observations[-1]})


def execute(action: Action) -> str:
    action = _validate_action(action)
    tool = TOOLS.get(action.name)
    if tool is None:
        raise ValueError(f"Unknown tool: {action.name}")
    _validate_arguments(
        action.name, action.arguments, EXPECTED_ARGUMENTS[action.name]
    )
    result = tool(**action.arguments)
    if type(result) is not str:
        raise TypeError(f"Tool {action.name} must return a string")
    return result


def run_agent(
    task: Task,
    planner: Callable[[Task], Action] = decide,
    max_steps: int = 3,
    *,
    deadline_seconds: float | None = 5.0,
    planner_timeout_seconds: float | None = 2.0,
    tool_timeout_seconds: float | None = 2.0,
) -> Task:
    """Run within separate step, wall-clock, planner, and tool budgets.

    ``max_steps`` limits decisions, not elapsed time. ``deadline_seconds`` caps
    the whole run, while the planner and tool timeouts cap individual calls.
    Passing ``None`` disables a time limit. Transport-level timeouts are still
    required to cancel the underlying provider, network, or subprocess work.
    """
    task = _validate_task(task)
    if not callable(planner):
        raise TypeError("planner must be callable")
    if type(max_steps) is not int or max_steps < 1:
        raise ValueError("max_steps must be at least 1")

    deadline_seconds = _validate_timeout("deadline_seconds", deadline_seconds)
    planner_timeout_seconds = _validate_timeout(
        "planner_timeout_seconds", planner_timeout_seconds
    )
    tool_timeout_seconds = _validate_timeout(
        "tool_timeout_seconds", tool_timeout_seconds
    )
    deadline_at = (
        None if deadline_seconds is None else time.monotonic() + deadline_seconds
    )

    for step in range(1, max_steps + 1):
        planner_task = Task(
            request=task.request,
            observations=list(task.observations),
            trace=list(task.trace),
            answer=task.answer,
        )
        action = _bounded_call(
            lambda: planner(planner_task),
            label="Planner",
            call_timeout_seconds=planner_timeout_seconds,
            deadline_at=deadline_at,
            deadline_seconds=deadline_seconds,
        )
        action = _validate_action(action)
        task.trace.append(f"{step}. {action.name}")

        if action.name == "finish":
            _validate_arguments("finish", action.arguments, {"answer": str})
            task.answer = action.arguments["answer"]
            return task

        observation = _bounded_call(
            lambda: execute(action),
            label=f"Tool {action.name}",
            call_timeout_seconds=tool_timeout_seconds,
            deadline_at=deadline_at,
            deadline_seconds=deadline_seconds,
        )
        task.observations.append(observation)

    raise RuntimeError(f"Agent exceeded the {max_steps}-step budget")


if __name__ == "__main__":
    result = run_agent(Task("My API request times out"))
    print("\n".join(result.trace))
    print(f"\nAnswer: {result.answer}")