"""Main orchestrator loop — autonomous supervisor for the Loom project.

Periodically sweeps for expired claims, retryable tasks, escalations,
and dispatches ready merge tasks.
"""

from __future__ import annotations

import asyncio
import json
import logging

import asyncpg
from redis.asyncio import Redis

from loom.bus import channels
from loom.bus.events import EventType
from loom.bus.publisher import publish_event
from loom.config import LoomConfig
from loom.graph import cache, store
from loom.observability.events import log_orchestrator_tick as _log_tick
from loom.orchestration.escalation import handle_escalation
from loom.orchestration.retry import sweep_retryable_tasks
from loom.orchestration.sweeper import sweep_expired_claims

log = logging.getLogger(__name__)

# TTL for the idempotency guard key (seconds).  Must be long enough to
# cover the full merge lifecycle, but short enough that a failed spawn
# doesn't permanently block the task.
_MERGE_SPAWN_TTL = 600  # 10 minutes


async def _check_project_complete(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
) -> bool:
    """Check if all tasks are done (no pending/claimed/blocked remaining).

    Returns True if the project is complete.
    """
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            SELECT
                COUNT(*) FILTER (WHERE status IN ('pending', 'claimed', 'blocked')) AS active,
                COUNT(*) FILTER (WHERE status = 'done') AS done,
                COUNT(*) AS total
            FROM tasks
            WHERE project_id = $1::uuid AND status NOT IN ('epic', 'idea')
            """,
            project_id,
        )

    if row["total"] == 0:
        return False

    if row["active"] == 0 and row["done"] > 0:
        await publish_event(
            redis, project_id, EventType.PROJECT_COMPLETE,
            payload={"done_count": row["done"], "total": row["total"]},
        )
        return True

    return False


async def _process_escalations(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    config: LoomConfig,
    project_dir: str | None,
) -> int:
    """Pop and handle escalations from the queue. Returns count handled."""
    escalation_enabled = getattr(config.orchestration, 'enable_escalation', True)
    handled = 0
    max_per_tick = getattr(config.orchestration, 'max_escalations_per_tick', 5)
    while True:
        # Non-blocking pop (timeout=1 second)
        result = await redis.brpop(
            channels.escalation_queue(project_id), timeout=1
        )
        if result is None:
            break

        _, data = result
        escalation = json.loads(data)
        task_id = escalation.get("task_id", "")
        message = escalation.get("message", "")

        # When escalation is disabled, drain the queue without handling
        if not escalation_enabled:
            log.debug("Escalation disabled — draining queue entry for task %s", task_id)
            continue

        try:
            await handle_escalation(
                pool, redis, project_id, task_id, message, config, project_dir,
            )
            handled += 1
            if handled >= max_per_tick:
                log.info("Escalation rate limit reached (%d/%d)", handled, max_per_tick)
                break
            # Send Slack notification if configured
            if config.integrations.slack_webhook_url:
                from loom.integrations.slack import send_slack_notification
                try:
                    await send_slack_notification(
                        webhook_url=config.integrations.slack_webhook_url,
                        message=f"Escalation handled for task {task_id}: {message}",
                        task_id=task_id,
                        severity="warning",
                    )
                except Exception:
                    log.warning("Failed to send Slack notification for %s", task_id)
        except Exception:
            log.exception("Failed to handle escalation for task %s", task_id)

    return handled


def _is_merge_task(task) -> bool:
    """Return True if *task* looks like a merge task.

    A task is considered a merge task when its title starts with
    ``"Merge branch "`` (as created by ``store._create_merge_task_in_txn``)
    or when its context contains a ``branch_name`` key.
    """
    if task.title.startswith("Merge branch "):
        return True
    if task.context and task.context.get("branch_name"):
        return True
    return False


async def _handle_merge_task_ready(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    project_dir: str | None,
) -> int:
    """Detect merge tasks in the ready queue, claim, and dispatch them.

    Uses a Redis SETNX idempotency guard so that a merge task is dispatched at
    most once even if multiple orchestrator ticks overlap.

    Returns the number of merge tasks dispatched.
    """
    ready_tasks = await cache.get_ready_tasks(redis, pool, project_id, limit=20)
    dispatched = 0

    for task in ready_tasks:
        if not _is_merge_task(task):
            continue

        spawn_key = channels.merge_spawned_key(project_id, task.id)

        # Idempotency guard: SETNX with TTL.  If the key already exists
        # this merge was already dispatched.
        already_spawned = not await redis.set(
            spawn_key, "1", nx=True, ex=_MERGE_SPAWN_TTL,
        )
        if already_spawned:
            log.debug("Merge task %s already dispatched (spawn guard exists)", task.id)
            continue

        try:
            # Claim the task atomically via Postgres.
            claimed = await store.claim_task(pool, task.id, "merge-agent")
            await cache.sync_task(redis, claimed)
            await cache.remove_from_ready_queue(redis, project_id, task.id)
            await publish_event(
                redis, project_id, EventType.TASK_CLAIMED,
                task_id=task.id, agent_id="merge-agent",
            )

            # Extract merge parameters from context.
            branch_name = claimed.context.get("branch_name", "")
            worktree_path = claimed.context.get(
                "worktree_path",
                f".claude/worktrees/{branch_name.removeprefix('worktree-')}",
            )

            # Dispatch the merge executor as a fire-and-forget background task.
            from loom.skills.merge import execute_merge_task

            asyncio.create_task(
                execute_merge_task(
                    pool=pool,
                    redis_client=redis,
                    project_id=project_id,
                    task_id=task.id,
                    branch_name=branch_name,
                    project_dir=project_dir or ".",
                    worktree_path=worktree_path,
                ),
                name=f"merge-{task.id}",
            )
            dispatched += 1
            log.info("Dispatched merge task %s for branch %s", task.id, branch_name)

        except Exception:
            log.exception("Failed to dispatch merge task %s", task.id)
            # Remove the idempotency guard so a retry can happen next tick.
            await redis.delete(spawn_key)
            # Best-effort: fail the task so it does not block.
            try:
                failed = await store.fail_task(
                    pool, task.id, "Failed to dispatch merge agent",
                )
                await cache.sync_task(redis, failed)
                await publish_event(
                    redis, project_id, EventType.TASK_FAILED,
                    task_id=task.id,
                    payload={"reason": "merge dispatch failed"},
                )
            except Exception:
                log.error("Could not fail task %s after dispatch error", task.id)

    return dispatched


async def orchestrator_tick(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    config: LoomConfig,
    project_dir: str | None = None,
) -> dict:
    """Run a single orchestrator sweep. Returns summary of actions taken."""
    summary: dict = {
        "expired_claims_released": 0,
        "stale_tasks": 0,
        "tasks_retried": 0,
        "escalations_handled": 0,
        "merge_tasks_dispatched": 0,
        "project_complete": False,
    }

    # 1. Sweep expired claims
    summary["expired_claims_released"] = await sweep_expired_claims(
        pool, redis, project_id,
    )

    # 1b. Detect stale tasks (claimed but no recent heartbeat)
    heartbeat_interval = config.orchestration.heartbeat_interval_seconds
    stale_threshold = heartbeat_interval * 2  # 2x interval = considered stale
    async with pool.acquire() as conn:
        stale_rows = await conn.fetch(
            """
            SELECT id FROM tasks
            WHERE project_id = $1::uuid
              AND status = 'claimed'
              AND claimed_at IS NOT NULL
              AND (
                  claim_expires_at IS NULL
                  OR claim_expires_at > NOW()
              )
              AND updated_at < NOW() - ($2 || ' seconds')::interval
            """,
            project_id, str(stale_threshold),
        )
    if stale_rows:
        summary["stale_tasks"] = len(stale_rows)
        log.warning(
            "Detected %d stale claimed tasks (no update in %ds): %s",
            len(stale_rows), stale_threshold,
            [r["id"] for r in stale_rows],
        )

    # 2. Sweep retryable tasks
    summary["tasks_retried"] = await sweep_retryable_tasks(
        pool, redis, project_id,
    )

    # 3. Handle escalations
    summary["escalations_handled"] = await _process_escalations(
        pool, redis, project_id, config, project_dir,
    )

    # 4. Dispatch ready merge tasks
    summary["merge_tasks_dispatched"] = await _handle_merge_task_ready(
        pool, redis, project_id, project_dir,
    )

    # 5. Sweep epics whose children are all done
    from loom.graph.deps import sweep_epic_completion
    closed_epics = await sweep_epic_completion(pool, redis, project_id)
    summary["epics_closed"] = len(closed_epics)

    # 6. Check project completion
    summary["project_complete"] = await _check_project_complete(
        pool, redis, project_id,
    )

    _log_tick(
        swept=summary["expired_claims_released"],
        retried=summary["tasks_retried"],
        escalated=summary["escalations_handled"],
    )

    return summary


async def orchestrator_loop(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    config: LoomConfig,
    project_dir: str | None = None,
) -> None:
    """Main orchestrator loop — runs until all tasks are done or cancelled.

    Each tick: sweep expired claims, sweep retryable tasks, handle
    escalations, check project completion.
    """
    interval = config.orchestration.sweep_interval_seconds
    log.info("Orchestrator started (interval=%ds)", interval)

    while True:
        try:
            summary = await orchestrator_tick(
                pool, redis, project_id, config, project_dir,
            )

            if any(v for k, v in summary.items() if k != "project_complete"):
                log.info(
                    "Orchestrator tick: released=%d retried=%d escalations=%d merges=%d",
                    summary["expired_claims_released"],
                    summary["tasks_retried"],
                    summary["escalations_handled"],
                    summary["merge_tasks_dispatched"],
                )

            if summary["project_complete"]:
                log.info("All tasks complete — orchestrator stopping.")
                return

        except asyncio.CancelledError:
            log.info("Orchestrator cancelled.")
            raise
        except Exception:
            log.exception("Orchestrator tick error (will retry next cycle)")

        await asyncio.sleep(interval)
