"""Crash recovery — reconcile orphaned containers and stale GPU allocations on startup."""

from __future__ import annotations

import logging
import subprocess
from datetime import datetime, timezone

from radix_edge.db import get_db

logger = logging.getLogger("radix_edge.recovery")


def reconcile_on_startup() -> dict:
    """Reconcile system state after an agent restart.

    Returns a summary dict with counts of actions taken.
    """
    summary = {
        "orphaned_containers_killed": 0,
        "stale_jobs_failed": 0,
        "stale_gpu_allocations_released": 0,
    }

    # 1. Find running Docker containers managed by radix
    managed_containers = _list_managed_containers()
    managed_job_ids = {c["job_id"] for c in managed_containers}

    # 2. Find jobs in "running" state in the database
    with get_db() as db:
        running_jobs = db.execute(
            "SELECT job_id FROM jobs WHERE status = 'running'"
        ).fetchall()
        running_job_ids = {row["job_id"] for row in running_jobs}

    # 3. Kill orphaned containers (container exists but no matching running job)
    for container in managed_containers:
        if container["job_id"] not in running_job_ids:
            _kill_container(container["container_id"])
            summary["orphaned_containers_killed"] += 1

    # 4. Fail stale jobs (job is "running" but no container exists)
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
    with get_db() as db:
        for job_id in running_job_ids:
            if job_id not in managed_job_ids:
                db.execute(
                    """UPDATE jobs SET status = 'failed',
                       error_message = 'Agent restarted during execution',
                       completed_at = ? WHERE job_id = ?""",
                    (now, job_id),
                )
                summary["stale_jobs_failed"] += 1
                logger.info("Marked orphaned job %s as failed", job_id)

    # 5. Release stale GPU allocations
    with get_db() as db:
        stale_gpus = db.execute(
            """SELECT gpu_index, assigned_job_id FROM gpus
               WHERE assigned_job_id IS NOT NULL AND status = 'allocated'"""
        ).fetchall()

        for gpu in stale_gpus:
            assigned_job = gpu["assigned_job_id"]
            # Check if the job is still actually running
            job_row = db.execute(
                "SELECT status FROM jobs WHERE job_id = ?", (assigned_job,)
            ).fetchone()

            if not job_row or job_row["status"] in ("succeeded", "failed", "cancelled"):
                db.execute(
                    "UPDATE gpus SET assigned_job_id = NULL, status = 'idle' WHERE gpu_index = ?",
                    (gpu["gpu_index"],),
                )
                summary["stale_gpu_allocations_released"] += 1
                logger.info(
                    "Released stale GPU %d (was assigned to %s)",
                    gpu["gpu_index"], assigned_job,
                )

    logger.info(
        "Recovery complete: %d containers killed, %d jobs failed, %d GPUs released",
        summary["orphaned_containers_killed"],
        summary["stale_jobs_failed"],
        summary["stale_gpu_allocations_released"],
    )
    return summary


def _list_managed_containers() -> list[dict]:
    """List Docker containers managed by radix (with radix.managed=true label)."""
    try:
        result = subprocess.run(
            [
                "docker", "ps",
                "--filter", "label=radix.managed=true",
                "--format", "{{.ID}}\t{{.Label \"radix.job_id\"}}",
            ],
            capture_output=True, text=True, timeout=10,
        )
        if result.returncode != 0:
            logger.warning("docker ps failed: %s", result.stderr)
            return []

        containers = []
        for line in result.stdout.strip().split("\n"):
            if not line.strip():
                continue
            parts = line.split("\t")
            if len(parts) >= 2:
                containers.append({
                    "container_id": parts[0],
                    "job_id": parts[1],
                })
        return containers
    except (subprocess.TimeoutExpired, FileNotFoundError) as e:
        logger.warning("Failed to list Docker containers: %s", e)
        return []


def _kill_container(container_id: str) -> None:
    """Kill and remove a Docker container."""
    try:
        subprocess.run(
            ["docker", "rm", "-f", container_id],
            capture_output=True, timeout=30,
        )
        logger.info("Killed orphaned container %s", container_id)
    except (subprocess.TimeoutExpired, FileNotFoundError) as e:
        logger.warning("Failed to kill container %s: %s", container_id, e)
