"""Docker-based agent runner with budget tracking and eject policies."""

from __future__ import annotations

import json
import logging
import os
import select
import subprocess
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable

from ._runtime import resolve_public
from .budget import BudgetTracker
from .eject import EjectPolicy, default_eject_policies
from .events import EventLogger
from .models import CacheStats, RunResult
from .provision import get_usage as provision_get_usage
from .provision import provision_key as provision_key_impl

logger = logging.getLogger(__name__)
get_usage = provision_get_usage
provision_key = provision_key_impl

BUDGET_POLL_INTERVAL_S = 30
FINAL_USAGE_POLL_INTERVAL_S = 0.5
FINAL_USAGE_POLL_ATTEMPTS = 6


def image_exists(image_name: str) -> bool:
    """Return True if a Docker image with the given name exists locally."""
    r = subprocess.run(
        ["docker", "image", "inspect", image_name],
        capture_output=True, text=True,
    )
    return r.returncode == 0


def build_image(
    image_name: str,
    dockerfile: str | Path = "Dockerfile",
    context: str | Path = ".",
) -> None:
    """Build a Docker image.

    Args:
        image_name: Name/tag to give the built image.
        dockerfile: Path to the Dockerfile.
        context: Build context directory.

    Raises:
        RuntimeError: If the build fails.
    """
    print(f"Building Docker image {image_name!r}…")
    result = subprocess.run(
        ["docker", "build", "-t", image_name, "-f", str(dockerfile), str(context)],
    )
    if result.returncode != 0:
        raise RuntimeError(f"Docker build failed for {image_name!r}")
    print(f"Image {image_name!r} built.\n")


def remove_container(name: str) -> None:
    """Force-remove a container by name (no-op if it doesn't exist)."""
    subprocess.run(["docker", "rm", "-f", name], capture_output=True)


def _write_env_file(env_vars: dict[str, str]) -> str:
    """Write Docker env vars to a temporary env file and return its path."""
    with tempfile.NamedTemporaryFile(
        mode="w",
        prefix="agent-preflight-",
        suffix=".env",
        delete=False,
        encoding="utf-8",
    ) as fh:
        for key, value in env_vars.items():
            if any(ch in key for ch in "\r\n") or any(ch in value for ch in "\r\n"):
                raise ValueError(f"Environment variable {key!r} contains a newline")
            fh.write(f"{key}={value}\n")
        return fh.name


def _format_mount_arg(mount: tuple[str, str, str] | str) -> str:
    """Normalize supported mount formats to a Docker ``-v`` argument."""
    if isinstance(mount, str):
        return mount

    host_path, container_path, mode = mount
    return f"{host_path}:{container_path}:{mode}"


def _poll_final_usage(api_key: str) -> float | None:
    """Poll final usage with short retries instead of a fixed sleep."""
    get_usage_fn = resolve_public("get_usage", get_usage, original=provision_get_usage)
    for attempt in range(FINAL_USAGE_POLL_ATTEMPTS):
        usage = get_usage_fn(api_key)
        if usage is not None:
            return usage
        if attempt < FINAL_USAGE_POLL_ATTEMPTS - 1:
            time.sleep(FINAL_USAGE_POLL_INTERVAL_S)
    return None


def run_container(
    *,
    image_name: str,
    model_id: str,
    alias: str,
    admin_key: str,
    budget_usd: float,
    run_dir: Path,
    run_id: str,
    env: dict[str, str] | None = None,
    mounts: list[tuple[str, str, str] | str] | None = None,
    initial_prompt: str | None = None,
    container_args: list[str] | None = None,
    eject_policy: EjectPolicy | None = None,
    on_event: Callable[[dict], None] | None = None,
    budget_mode: str = "simple",
    workspace_path: str = "/workspace/output",
    idle_timeout_s: float = 300.0,
    log_fn: Callable[..., None] | None = None,
) -> RunResult:
    """Run a single agent in a Docker container with budget enforcement.

    The container receives:
    - ``OPENROUTER_API_KEY`` env var (provisioned sub-key with ``budget_usd`` cap)
    - Any additional ``env`` vars you specify
    - File mounts from ``mounts`` as tuples or Docker ``host:container:mode`` strings

    Budget is polled from OpenRouter every ``BUDGET_POLL_INTERVAL_S`` seconds and
    written to ``{workspace_path}/budget.txt`` inside the container.

    If ``initial_prompt`` is set, it is sent to the container's stdin as a JSON
    RPC message (compatible with pi-coding-agent's ``--mode rpc``).

    Args:
        image_name: Docker image to run.
        model_id: Full OpenRouter model ID (e.g. ``"meta-llama/llama-3.3-70b-instruct"``).
        alias: Short human-readable name for logs/directories.
        admin_key: OpenRouter admin key for provisioning the sub-key.
        budget_usd: USD budget cap for this run.
        run_dir: Directory to store per-run artifacts (logs, events.jsonl, summary.json).
        run_id: Unique run identifier used in key labels and paths.
        env: Additional environment variables to set inside the container.
        mounts: Volume mounts as ``[(host_path, container_path, mode), ...]`` or
                Docker-style ``["host_path:container_path:mode", ...]`` strings.
                Tuple mode is typically ``"ro"`` or ``"rw"``.
        initial_prompt: If set, sent as a JSON RPC prompt via stdin after container starts.
        container_args: Extra args appended to ``docker run`` (e.g. ``["--model", model_id]``).
        eject_policy: Policy controlling early termination. Defaults to
                      :func:`~agent_preflight.eject.default_eject_policies`.
        on_event: Optional callback called with each parsed JSON event.
        budget_mode: ``"simple"`` or ``"detailed"`` budget display format.
        workspace_path: Path inside the container where ``budget.txt`` is written.
        idle_timeout_s: Seconds of inactivity before the container is killed.
        log_fn: Custom log function (defaults to ``print``).

    Note:
        Output polling currently uses ``select.select`` on pipe handles, so this
        implementation is intended for Unix-like hosts.

    Returns:
        :class:`~agent_preflight.models.RunResult`.
    """
    log = log_fn or print
    model_dir = run_dir / alias
    model_dir.mkdir(parents=True, exist_ok=True)

    if eject_policy is None:
        eject_policy = default_eject_policies(idle_timeout_s=idle_timeout_s)

    # Provision a capped sub-key
    try:
        provision_key_fn = resolve_public("provision_key", provision_key, original=provision_key_impl)
        model_key = provision_key_fn(admin_key, f"agent-preflight:{run_id}:{alias}", budget_usd)
    except Exception as exc:
        log(f"[{alias}] ERROR: Key provisioning failed: {exc}")
        return RunResult(
            alias=alias,
            model_id=model_id,
            budget_usd=budget_usd,
            started_at=datetime.now(timezone.utc).isoformat(),
            finished_at=datetime.now(timezone.utc).isoformat(),
            error=f"provisioning_failed: {exc}",
        )

    log(f"[{alias}] Starting — model={model_id}, budget=${budget_usd:.2f}")

    container_name = f"ap-{alias}-{run_id[:8]}"
    remove_container(container_name)
    env_vars = dict(env or {})
    env_vars["OPENROUTER_API_KEY"] = model_key
    env_file_path = _write_env_file(env_vars)

    # Build docker run command
    docker_cmd = [
        "docker", "run", "--rm", "--init", "-i",
        "--name", container_name,
        "--add-host", "host.docker.internal:host-gateway",
        "--env-file", env_file_path,
    ]
    for mount in (mounts or []):
        docker_cmd += ["-v", _format_mount_arg(mount)]
    docker_cmd.append(image_name)
    if container_args:
        docker_cmd += container_args

    tracker = BudgetTracker(
        api_key=model_key,
        limit_usd=budget_usd,
        poll_interval_s=BUDGET_POLL_INTERVAL_S,
        mode=budget_mode,
    )

    events: list[dict] = []
    context: dict = {"tool_calls": 0}
    cache_stats = CacheStats()

    started_at = datetime.now(timezone.utc)
    result = RunResult(
        alias=alias,
        model_id=model_id,
        budget_usd=budget_usd,
        started_at=started_at.isoformat(),
    )

    ev_logger = EventLogger(model_dir / "events.jsonl", start_time=started_at)

    try:
        proc = subprocess.Popen(
            docker_cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        )

        # Send initial prompt via JSON RPC if requested
        if initial_prompt:
            proc.stdin.write(json.dumps({
                "id": "init", "type": "prompt", "message": initial_prompt,
            }) + "\n")
            proc.stdin.flush()
            log(f"[{alias}] Prompt sent.")

        last_budget_poll_ts = time.time()

        while True:
            # Non-blocking read with 1s timeout
            ready, _, _ = select.select([proc.stdout], [], [], 1.0)

            if not ready:
                if proc.poll() is not None:
                    break

                # Periodic budget poll + injection
                if time.time() - last_budget_poll_ts >= BUDGET_POLL_INTERVAL_S:
                    last_budget_poll_ts = time.time()
                    status = tracker.inject_to_docker(container_name, workspace_path)
                    ev_logger.log("budget_poll", used=status.used, remaining=status.remaining)
                    log(f"[{alias}] [budget] {status}")

                    # Eject check
                    decision = eject_policy.check(
                        elapsed_s=(datetime.now(timezone.utc) - started_at).total_seconds(),
                        budget_status=status,
                        events=events,
                        context=context,
                    )
                    if decision.should_eject:
                        log(f"[{alias}] EJECTING: {decision.reason}")
                        proc.kill()
                        result.error = f"ejected: {decision.reason}"
                        break
                continue

            line = proc.stdout.readline()
            if not line:
                if proc.poll() is not None:
                    break
                continue

            line = line.strip()
            if not line:
                continue

            # Parse event
            try:
                event = json.loads(line)
            except json.JSONDecodeError:
                event = {"type": "raw", "data": line}

            events.append(event)
            ev_logger.log(event.get("type", "raw"), **{
                k: v for k, v in event.items() if k != "type"
            })

            if on_event:
                on_event(event)

            # Track tool calls and notify budget tracker
            etype = event.get("type", "")
            if etype == "tool_execution_start":
                context["tool_calls"] = context.get("tool_calls", 0) + 1
                tracker.on_tool_start(event.get("toolName", ""))
            elif etype == "tool_execution_end":
                tracker.on_tool_end(event.get("toolName", ""))

            # Track prompt cache metrics from usage events
            usage = event.get("usage", {})
            if usage and isinstance(usage, dict):
                prompt_tokens = usage.get("prompt_tokens", 0)
                if prompt_tokens:
                    cache_stats.request_count += 1
                    cache_stats.total_prompt_tokens += prompt_tokens

                    # OpenRouter normalized: prompt_tokens_details
                    details = usage.get("prompt_tokens_details", {})
                    cached = (details.get("cached_tokens", 0) if details else 0)
                    creation = (details.get("cache_write_tokens", 0) if details else 0)

                    # Fallback: Anthropic native fields
                    cached += usage.get("cache_read_input_tokens", 0)
                    creation += usage.get("cache_creation_input_tokens", 0)

                    if cached > 0:
                        cache_stats.cached_tokens += cached
                        cache_stats.cache_hit_requests += 1
                    cache_stats.cache_creation_tokens += creation

            # Agent signalled completion
            if etype == "agent_end":
                log(f"[{alias}] Agent finished.")
                time.sleep(2)
                proc.terminate()
                break

        proc.wait(timeout=10)

    except Exception as exc:
        result.error = str(exc)
        log(f"[{alias}] ERROR: {exc}")

    finally:
        ev_logger.close()
        if os.path.exists(env_file_path):
            os.unlink(env_file_path)

    # Populate result
    finished_at = datetime.now(timezone.utc)
    result.finished_at = finished_at.isoformat()
    result.duration_s = int((finished_at - started_at).total_seconds())
    result.tool_calls = context.get("tool_calls", 0)
    result.events_count = len(events)
    result.cache = cache_stats

    # Final usage poll with short retries so fast runs do not pay a fixed delay.
    final_usage = _poll_final_usage(model_key)
    cache_pct = f" cache={cache_stats.hit_rate * 100:.0f}%" if cache_stats.request_count else ""
    if final_usage is not None:
        result.cost_usd = round(final_usage, 4)
        log(f"[{alias}] Done — cost=${final_usage:.4f} duration={result.duration_s}s "
            f"tools={result.tool_calls}{cache_pct} error={result.error or 'none'}")
    else:
        log(f"[{alias}] Done — cost=unknown duration={result.duration_s}s "
            f"tools={result.tool_calls}{cache_pct} error={result.error or 'none'}")

    # Log detailed cache stats if any requests were tracked
    if cache_stats.request_count > 0:
        log(f"[{alias}] Cache: {cache_stats.hit_rate * 100:.1f}% hit rate — "
            f"{cache_stats.cached_tokens:,} cached / {cache_stats.total_prompt_tokens:,} "
            f"prompt tokens | {cache_stats.cache_hit_requests}/{cache_stats.request_count} "
            f"requests hit cache"
            + (f" | {cache_stats.cache_creation_tokens:,} write tokens"
               if cache_stats.cache_creation_tokens else ""))

    # Write summary.json
    (model_dir / "summary.json").write_text(json.dumps(result.to_dict(), indent=2))
    return result
