"""Budget tracking with periodic polling and optional Docker injection."""

from __future__ import annotations

import logging
import subprocess
import time
from typing import Callable

from ._runtime import resolve_public
from .models import BudgetStatus
from .provision import get_usage as provision_get_usage

logger = logging.getLogger(__name__)
get_usage = provision_get_usage


def _get_usage(api_key: str) -> float | None:
    return resolve_public("get_usage", get_usage, original=provision_get_usage)(api_key)


class BudgetTracker:
    """Polls OpenRouter for actual spend and enforces a budget limit.

    Supports two display modes for the injected budget.txt:

    - ``"simple"``   (default) — one-line balance summary::

          Budget: $0.086 remaining of $0.15 (57% left)

    - ``"detailed"`` — adds per-action cost history and estimated remaining actions::

          Budget: $0.086 remaining of $0.15 (57% left)

          Recent costs:
            - $0.0012  (bash)
            - $0.0008  (read_file)
            ...

          Total actions: 14  |  Avg cost/action: $0.0010
          Estimated actions remaining: ~86

    Usage::

        tracker = BudgetTracker(api_key="sk-or-...", limit_usd=0.25)

        # Poll every 30s and check budget
        status = tracker.poll()
        if status.exceeded:
            # Kill the agent

        # Get a string for injecting into a system prompt
        msg = tracker.status_message()

        # Inject budget.txt into a running Docker container
        tracker.inject_to_docker("my-container-name")
    """

    def __init__(
        self,
        api_key: str,
        limit_usd: float,
        poll_interval_s: float = 30.0,
        mode: str = "simple",
        on_exceeded: Callable[[], None] | None = None,
    ) -> None:
        """
        Args:
            api_key: The provisioned sub-key to track (not the admin key).
            limit_usd: Budget cap in USD.
            poll_interval_s: Minimum seconds between OpenRouter polls.
            mode: ``"simple"`` or ``"detailed"``.
            on_exceeded: Optional callback invoked once when the budget is first exceeded.
        """
        self.api_key = api_key
        self.limit_usd = limit_usd
        self.poll_interval_s = poll_interval_s
        self.mode = mode
        self.on_exceeded = on_exceeded

        self._last_poll_ts: float = 0.0
        self._last_usage: float = 0.0
        self._exceeded_fired: bool = False
        self._action_history: list[dict] = []
        self._pending_tool: str | None = None

    # ------------------------------------------------------------------
    # Tool lifecycle hooks (for "detailed" mode cost attribution)
    # ------------------------------------------------------------------

    def on_tool_start(self, tool_name: str) -> None:
        """Call when a tool execution starts to track per-action cost."""
        if self.mode == "detailed":
            self._pending_tool = tool_name

    def on_tool_end(self, tool_name: str) -> None:  # noqa: ARG002
        """Call when a tool execution ends."""
        self._pending_tool = None

    # ------------------------------------------------------------------
    # Polling
    # ------------------------------------------------------------------

    def poll(self, force: bool = False) -> BudgetStatus:
        """Poll OpenRouter if the interval has elapsed.

        Args:
            force: Bypass the poll interval and always query the API.

        Returns:
            Current :class:`BudgetStatus`.
        """
        now = time.time()
        if not force and now - self._last_poll_ts < self.poll_interval_s:
            return BudgetStatus(used=self._last_usage, limit=self.limit_usd)

        self._last_poll_ts = now
        usage = _get_usage(self.api_key)
        if usage is not None and usage != self._last_usage:
            if self.mode == "detailed" and self._last_usage is not None:
                delta = usage - self._last_usage
                if delta > 0.0001:
                    self._action_history.append({
                        "ts": now,
                        "cost": round(delta, 5),
                        "label": self._pending_tool or "thinking",
                    })
            self._last_usage = usage

        status = BudgetStatus(used=self._last_usage, limit=self.limit_usd)
        if status.exceeded and not self._exceeded_fired:
            self._exceeded_fired = True
            if self.on_exceeded:
                self.on_exceeded()
        return status

    @property
    def exceeded(self) -> bool:
        """True if the last polled usage has reached the limit."""
        return self._last_usage >= self.limit_usd

    # ------------------------------------------------------------------
    # Injection
    # ------------------------------------------------------------------

    def status_message(self, force_poll: bool = False) -> str:
        """Return a budget status string for injection into a system prompt."""
        status = self.poll(force=force_poll)
        if self.mode == "detailed":
            return self._format_detailed(status)
        return status.format_simple()

    def inject_to_docker(
        self,
        container_name: str,
        workspace_path: str = "/workspace/output",
        timeout_s: float = 5.0,
    ) -> BudgetStatus:
        """Poll usage and write budget.txt into a running Docker container.

        The file is written to ``{workspace_path}/budget.txt``. Agents that
        read this file will see up-to-date cost information.

        Args:
            container_name: Docker container name or ID.
            workspace_path: Path inside the container to write budget.txt.
            timeout_s: Timeout for the docker exec command.

        Returns:
            Current :class:`BudgetStatus`.
        """
        status = self.poll()
        txt = self._format_detailed(status) if self.mode == "detailed" else status.format_simple()

        try:
            subprocess.run(
                [
                    "docker", "exec", container_name, "sh", "-c",
                    f"cat > {workspace_path}/budget.txt << 'BUDGET_EOF'\n{txt}\nBUDGET_EOF",
                ],
                capture_output=True,
                timeout=timeout_s,
            )
        except Exception as exc:
            logger.debug("budget injection failed for %s: %s", container_name, exc)

        return status

    # ------------------------------------------------------------------
    # Internal formatting
    # ------------------------------------------------------------------

    def _format_detailed(self, status: BudgetStatus) -> str:
        lines = [status.format_simple(), ""]
        if self._action_history:
            lines.append("Recent costs:")
            for a in self._action_history[-5:]:
                lines.append(f"  - ${a['cost']:.4f}  ({a['label']})")
            lines.append("")
            total = len(self._action_history)
            avg = self._last_usage / total if total > 0 else 0
            est = int(status.remaining / avg) if avg > 0.0001 else 0
            lines.append(f"Total actions: {total}  |  Avg cost/action: ${avg:.4f}")
            lines.append(f"Estimated actions remaining: ~{est}")
        else:
            lines.append("(No cost data yet)")
        return "\n".join(lines)
