"""Recovery protocol — reconstruct state after orchestrator death.

Classifies all claimed tasks into categories and applies recovery actions:
- marker_tasks: have a completion marker → auto-complete
- stale_tasks: no heartbeat for 3x heartbeat_interval → re-queue
- expired_tasks: claim_expires_at < now → re-queue
- active_tasks: recent heartbeat → leave alone
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone

import asyncpg
from redis.asyncio import Redis

from loom.config import LoomConfig
from loom.graph import cache, deps, store
from loom.graph.task import Task
from loom.markers import read_markers


@dataclass
class RecoveryPlan:
    """Classification of all claimed tasks for recovery."""

    marker_tasks: list[dict] = field(default_factory=list)
    stale_tasks: list[dict] = field(default_factory=list)
    expired_tasks: list[dict] = field(default_factory=list)
    active_tasks: list[dict] = field(default_factory=list)

    @property
    def total_recoverable(self) -> int:
        return len(self.marker_tasks) + len(self.stale_tasks) + len(self.expired_tasks)

    def to_dict(self) -> dict:
        return {
            "marker_tasks": self.marker_tasks,
            "stale_tasks": self.stale_tasks,
            "expired_tasks": self.expired_tasks,
            "active_tasks": self.active_tasks,
            "total_recoverable": self.total_recoverable,
        }


async def build_recovery_plan(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    config: LoomConfig,
    project_dir: str | None = None,
) -> RecoveryPlan:
    """Classify all claimed tasks into recovery categories."""
    plan = RecoveryPlan()
    now = datetime.now(timezone.utc)
    heartbeat_interval = config.orchestration.heartbeat_interval_seconds

    # Get all claimed tasks
    claimed_tasks = await store.get_tasks_by_status(pool, project_id, "claimed", limit=500)

    # Read completion markers from filesystem
    markers_by_id: dict[str, dict] = {}
    if project_dir:
        for marker in read_markers(project_dir):
            markers_by_id[marker["task_id"]] = marker

    for task in claimed_tasks:
        task_info = {
            "task_id": task.id,
            "title": task.title,
            "assignee": task.assignee,
        }

        # Check if a completion marker exists
        if task.id in markers_by_id:
            marker = markers_by_id[task.id]
            task_info["marker_status"] = marker.get("status", "done")
            task_info["marker_output"] = marker.get("output", {})
            task_info["marker_branch"] = marker.get("branch_name", "")
            task_info["marker_reason"] = marker.get("reason", "")
            plan.marker_tasks.append(task_info)
            continue

        # Check claim expiry
        if task.claim_expires_at:
            expires = task.claim_expires_at
            if expires.tzinfo is None:
                expires = expires.replace(tzinfo=timezone.utc)
            if expires < now:
                task_info["expired_at"] = expires.isoformat()
                plan.expired_tasks.append(task_info)
                continue

        # Check heartbeat staleness via Redis progress data
        progress = await cache.get_task_progress(redis, project_id, task.id)
        if progress and progress.get("updated_at"):
            try:
                last_heartbeat = datetime.fromisoformat(progress["updated_at"])
                if last_heartbeat.tzinfo is None:
                    last_heartbeat = last_heartbeat.replace(tzinfo=timezone.utc)
                stale_threshold = heartbeat_interval * 3
                if (now - last_heartbeat).total_seconds() > stale_threshold:
                    task_info["last_heartbeat"] = progress["updated_at"]
                    plan.stale_tasks.append(task_info)
                    continue
            except (ValueError, TypeError):
                pass

        # If no progress data at all, check claimed_at for staleness
        if not progress and task.claimed_at:
            claimed_at = task.claimed_at
            if claimed_at.tzinfo is None:
                claimed_at = claimed_at.replace(tzinfo=timezone.utc)
            stale_threshold = heartbeat_interval * 3
            if (now - claimed_at).total_seconds() > stale_threshold:
                task_info["claimed_at"] = claimed_at.isoformat()
                task_info["no_heartbeat"] = True
                plan.stale_tasks.append(task_info)
                continue

        # Task is active
        if progress:
            task_info["last_heartbeat"] = progress.get("updated_at", "")
            task_info["progress"] = progress.get("message", "")
            task_info["percent"] = progress.get("percent")
        plan.active_tasks.append(task_info)

    # Also check for markers of tasks that are NOT claimed (sweeper already released)
    all_task_ids = {t.id for t in claimed_tasks}
    for tid, marker in markers_by_id.items():
        if tid not in all_task_ids:
            task_info = {
                "task_id": tid,
                "title": "(unclaimed — marker only)",
                "assignee": marker.get("agent_id", ""),
                "marker_status": marker.get("status", "done"),
                "marker_output": marker.get("output", {}),
                "marker_branch": marker.get("branch_name", ""),
                "marker_reason": marker.get("reason", ""),
            }
            plan.marker_tasks.append(task_info)

    return plan


async def execute_recovery_plan(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    config: LoomConfig,
    plan: RecoveryPlan,
    project_dir: str | None = None,
) -> dict:
    """Apply the recovery plan: complete marker tasks, reset stale/expired."""
    from loom.bus.events import EventType
    from loom.bus.publisher import publish_event
    from loom.markers import archive_marker, read_markers

    results = {
        "completed_from_markers": [],
        "failed_from_markers": [],
        "reset_stale": [],
        "reset_expired": [],
        "errors": [],
    }

    # Build marker path lookup once (avoids re-reading for every task)
    marker_paths: dict[str, str] = {}
    if project_dir:
        for marker in read_markers(project_dir):
            if "_path" in marker:
                marker_paths[marker["task_id"]] = marker["_path"]

    # Process marker tasks
    for mt in plan.marker_tasks:
        tid = mt["task_id"]
        try:
            if mt.get("marker_status") == "failed":
                # Mark as failed via store
                task = await store.fail_task(pool, tid, mt.get("marker_reason", "recovered from marker"))
                await cache.sync_task(redis, task)
                await cache.remove_from_ready_queue(redis, project_id, tid)
                results["failed_from_markers"].append(tid)
            else:
                # Complete the task
                output = mt.get("marker_output", {})
                branch_name = mt.get("marker_branch") or None
                try:
                    result = await store.complete_task(pool, tid, output, branch_name=branch_name)
                    task = result[0] if isinstance(result, tuple) else result
                except Exception:
                    # If complete_task fails (e.g., already pending after sweeper),
                    # try batch_done which handles both claimed and pending
                    done_tasks, failures = await store.batch_done(pool, [tid], output)
                    if failures:
                        results["errors"].append({"task_id": tid, "error": failures[0]["error"]})
                        continue
                    task = done_tasks[0]
                await cache.sync_task(redis, task)
                await cache.remove_from_ready_queue(redis, project_id, tid)
                await publish_event(redis, project_id, EventType.TASK_DONE, tid)
                await deps.check_and_unblock(pool, redis, tid, project_id)
                results["completed_from_markers"].append(tid)
        except Exception as e:
            results["errors"].append({"task_id": tid, "error": str(e)})

        # Archive the marker file
        if tid in marker_paths:
            try:
                archive_marker(marker_paths[tid])
            except Exception:
                pass

    # Reset stale tasks
    for st in plan.stale_tasks:
        tid = st["task_id"]
        try:
            task = await store.reset_task(pool, tid)
            await cache.sync_task(redis, task)
            await cache.add_to_ready_queue(redis, task)
            await publish_event(redis, project_id, EventType.TASK_RESET, tid)
            results["reset_stale"].append(tid)
        except Exception as e:
            results["errors"].append({"task_id": tid, "error": str(e)})

    # Reset expired tasks
    for et in plan.expired_tasks:
        tid = et["task_id"]
        try:
            task = await store.reset_task(pool, tid)
            await cache.sync_task(redis, task)
            await cache.add_to_ready_queue(redis, task)
            await publish_event(redis, project_id, EventType.TASK_RESET, tid)
            results["reset_expired"].append(tid)
        except Exception as e:
            results["errors"].append({"task_id": tid, "error": str(e)})

    return results
