"""Job scheduler — atomic claim, shard-aware scheduling, and stale job reaper."""

from __future__ import annotations

import json
from typing import Any

import structlog

from hive_compute.coordinator.database import get_db_ctx, record_event
from hive_compute.coordinator.metrics import (
    increment_job_failed,
    increment_shard_completed,
    increment_shard_failed,
    observe_job_turnaround_duration,
    observe_shard_processing_duration,
)
from hive_compute.coordinator.reputation import record_shard_outcome
from hive_compute.coordinator.worker_observability import (
    record_worker_claim,
    record_worker_completion,
    record_worker_failure,
    sync_worker_busy_state,
)
from hive_compute.model_ids import DEFAULT_QUANTIZATION, canonical_model_identifier

logger = structlog.get_logger()


def describe_worker_benchmark_context(worker_id: str) -> dict[str, Any] | None:
    """Summarize pending transcription work against a worker's latest benchmarks."""
    with get_db_ctx() as db:
        worker = db.execute("SELECT worker_id FROM workers WHERE worker_id = ?", (worker_id,)).fetchone()
        if not worker:
            return None

        benchmark_rows = db.execute(
            """
            SELECT benchmark_id, model, backend, realtime_factor, created_at
            FROM worker_benchmarks
            WHERE worker_id = ?
            ORDER BY benchmark_id DESC
            """,
            (worker_id,),
        ).fetchall()
        pending_rows = db.execute(
            """
            WITH pending_units AS (
                SELECT json_extract(j.input_config, '$.model') AS model,
                       j.priority AS priority,
                       j.created_at AS created_at
                FROM jobs j
                WHERE j.status = 'pending'
                  AND j.job_type = 'transcribe'
                  AND j.shard_count IS NULL

                UNION ALL

                SELECT json_extract(j.input_config, '$.model') AS model,
                       j.priority AS priority,
                       j.created_at AS created_at
                FROM shards s
                JOIN jobs j ON s.job_id = j.job_id
                WHERE s.status = 'pending'
                  AND j.status = 'pending'
                  AND j.job_type = 'transcribe'
            )
            SELECT
                model,
                COUNT(*) AS pending_units,
                MIN(priority) AS best_priority,
                MIN(created_at) AS oldest_created_at
            FROM pending_units
            WHERE COALESCE(TRIM(model), '') != ''
            GROUP BY model
            ORDER BY best_priority ASC, oldest_created_at ASC, model ASC
            """
        ).fetchall()

    latest_benchmarks: dict[str, dict[str, Any]] = {}
    for row in benchmark_rows:
        model_key = _canonicalize_benchmark_model(row["model"])
        if model_key in latest_benchmarks:
            continue
        latest_benchmarks[model_key] = {
            "model": model_key,
            "recorded_model": row["model"],
            "backend": row["backend"],
            "realtime_factor": row["realtime_factor"],
            "created_at": row["created_at"],
        }

    top_pending_benchmark: dict[str, Any] | None = None
    for row in pending_rows:
        benchmark = latest_benchmarks.get(_canonicalize_benchmark_model(row["model"]))
        if benchmark and benchmark["realtime_factor"] is not None:
            if top_pending_benchmark is None or benchmark["realtime_factor"] > top_pending_benchmark["realtime_factor"]:
                top_pending_benchmark = benchmark

    pending_models: list[dict[str, Any]] = []
    missing_benchmarks: list[str] = []
    slower_than_top: list[str] = []

    for row in pending_rows:
        model_key = _canonicalize_benchmark_model(row["model"])
        benchmark = latest_benchmarks.get(model_key)

        if benchmark is None:
            note = "No benchmark recorded for this model; scheduler falls back to priority and age."
            missing_benchmarks.append(model_key)
        elif (
            top_pending_benchmark
            and benchmark["realtime_factor"] is not None
            and benchmark["realtime_factor"] < top_pending_benchmark["realtime_factor"]
        ):
            note = (
                f"Latest benchmark is slower than '{top_pending_benchmark['recorded_model']}' "
                f"for this worker."
            )
            slower_than_top.append(model_key)
        else:
            note = "Latest benchmark makes this one of the fastest pending models for this worker."

        pending_models.append(
            {
                "model": model_key,
                "recorded_model": row["model"],
                "pending_units": row["pending_units"],
                "best_priority": row["best_priority"],
                "oldest_created_at": row["oldest_created_at"],
                "benchmark": benchmark,
                "note": note,
            }
        )

    if top_pending_benchmark:
        preferred_model = top_pending_benchmark["model"]
        preferred_reason = (
            "Highest latest realtime factor among pending models: "
            f"{top_pending_benchmark['realtime_factor']}x on {top_pending_benchmark['recorded_model']}."
        )
    elif pending_models:
        preferred_model = pending_models[0]["model"]
        preferred_reason = "No pending model has benchmark coverage; scheduler falls back to priority and age."
    else:
        preferred_model = None
        preferred_reason = "No pending transcription work was found."

    return {
        "worker_id": worker_id,
        "preferred_model": preferred_model,
        "preferred_reason": preferred_reason,
        "missing_benchmarks": missing_benchmarks,
        "slower_than_top": slower_than_top,
        "benchmarked_models": list(latest_benchmarks.values()),
        "pending_models": pending_models,
    }


def claim_next_job(worker_id: str, job_types: list[str] | None = None) -> dict | None:
    """Atomically claim the next available work for a worker.

    For sharded jobs: claims a pending shard and returns job + shard info.
    For non-sharded jobs: claims the whole job directly (backward compatible).

    SQLite serializes writes, so this is inherently safe against double-claims.
    """
    if not job_types:
        job_types = ["transcribe"]

    # Try shards first (sharded jobs have shard_count > 0)
    shard = _claim_next_shard(worker_id, job_types)
    if shard:
        return shard

    # Fall back to non-sharded jobs
    return _claim_next_whole_job(worker_id, job_types)


def _claim_next_shard(worker_id: str, job_types: list[str]) -> dict | None:
    """Claim the next pending shard from a sharded job."""
    placeholders = ",".join("?" for _ in job_types)

    with get_db_ctx() as db:
        available_models = _load_worker_available_models(db, worker_id)
        if not available_models:
            return None

        row = _claim_matching_shard_candidate(
            db,
            worker_id,
            job_types,
            available_models,
            placeholders,
            relax_anti_correlation=False,
        )
        if not row:
            row = _claim_matching_shard_candidate(
                db,
                worker_id,
                job_types,
                available_models,
                placeholders,
                relax_anti_correlation=True,
            )

        if not row:
            return None

        # Get parent job info
        job = db.execute(
            "SELECT * FROM jobs WHERE job_id = ?",
            (row["job_id"],),
        ).fetchone()

        record_worker_claim(db, worker_id, row["job_id"])

        record_event(
            db, row["job_id"], "shard_claimed",
            from_status="pending", to_status="claimed",
            actor=worker_id, detail=f"shard {row['shard_index']} ({row['shard_id']})",
        )

        logger.info(
            "Shard claimed",
            job_id=row["job_id"],
            shard_id=row["shard_id"],
            shard_index=row["shard_index"],
            worker_id=worker_id,
        )

        # Return combined job + shard info for the worker
        result = dict(job)
        result["shard_id"] = row["shard_id"]
        result["shard_index"] = row["shard_index"]
        result["start_seconds"] = row["start_seconds"]
        result["end_seconds"] = row["end_seconds"]
        result["canary_id"] = row["canary_id"]
        return result


def _claim_next_whole_job(worker_id: str, job_types: list[str]) -> dict | None:
    """Claim a non-sharded job (shard_count IS NULL)."""
    placeholders = ",".join("?" for _ in job_types)

    with get_db_ctx() as db:
        available_models = _load_worker_available_models(db, worker_id)
        if not available_models:
            return None

        candidates = db.execute(
            f"""
            SELECT job_id, input_config
            FROM jobs
            WHERE status = 'pending'
              AND job_type IN ({placeholders})
              AND attempts < max_attempts
              AND shard_count IS NULL
            ORDER BY priority ASC,
                     COALESCE((
                         SELECT wb.realtime_factor
                         FROM worker_benchmarks wb
                         WHERE wb.worker_id = ?
                           AND CASE
                                   WHEN wb.model = 'large-v3-turbo' THEN 'turbo'
                                   ELSE wb.model
                               END = CASE
                                   WHEN json_extract(jobs.input_config, '$.model') = 'large-v3-turbo' THEN 'turbo'
                                   ELSE json_extract(jobs.input_config, '$.model')
                               END
                         ORDER BY wb.created_at DESC
                         LIMIT 1
                     ), 0) DESC,
                     created_at ASC
            LIMIT 25
            """,
            (*job_types, worker_id),
        ).fetchall()

        row = None
        for candidate in candidates:
            if _requested_model_identifier(candidate["input_config"]) not in available_models:
                continue

            row = db.execute(
                """
                UPDATE jobs
                SET status = 'claimed',
                    claimed_by = ?,
                    claimed_at = datetime('now'),
                    last_heartbeat = datetime('now'),
                    attempts = attempts + 1
                WHERE job_id = ? AND status = 'pending'
                RETURNING *
                """,
                (worker_id, candidate["job_id"]),
            ).fetchone()
            if row:
                break

        if row:
            record_worker_claim(db, worker_id, row["job_id"])
            record_event(db, row["job_id"], "claimed", from_status="pending", to_status="claimed", actor=worker_id)
            logger.info("Job claimed", job_id=row["job_id"], worker_id=worker_id)
            return dict(row)

    return None

def _claim_matching_shard_candidate(
    db,
    worker_id: str,
    job_types: list[str],
    available_models: set[str],
    placeholders: str,
    *,
    relax_anti_correlation: bool,
):
    anti_correlation_filter = ""
    params: list[object] = [*job_types]
    if not relax_anti_correlation:
        anti_correlation_filter = """
              AND NOT (
                  EXISTS (
                      SELECT 1 FROM shards adj
                      WHERE adj.job_id = s.job_id
                        AND adj.claimed_by = ?
                        AND ABS(adj.shard_index - s.shard_index) = 1
                  )
                  AND EXISTS (
                      SELECT 1 FROM shards alt
                      WHERE alt.job_id = s.job_id
                        AND alt.status = 'pending'
                        AND alt.attempts < alt.max_attempts
                        AND NOT EXISTS (
                            SELECT 1 FROM shards adj2
                            WHERE adj2.job_id = alt.job_id
                              AND adj2.claimed_by = ?
                              AND ABS(adj2.shard_index - alt.shard_index) = 1
                        )
                  )
              )
        """
        params.extend([worker_id, worker_id])

    params.extend([worker_id, worker_id])
    candidates = db.execute(
        f"""
        SELECT s.shard_id, j.input_config
        FROM shards s
        JOIN jobs j ON s.job_id = j.job_id
        WHERE s.status = 'pending'
          AND j.job_type IN ({placeholders})
          AND j.status = 'pending'
          AND s.attempts < s.max_attempts
          {anti_correlation_filter}
          AND (
              (SELECT COUNT(*) FROM shards cap
               WHERE cap.job_id = s.job_id AND cap.claimed_by = ?
                 AND cap.status IN ('claimed', 'processing')) * 10
              < (SELECT COUNT(*) FROM shards tot WHERE tot.job_id = s.job_id) * 3
              OR (SELECT COUNT(*) FROM shards cap2
                  WHERE cap2.job_id = s.job_id AND cap2.claimed_by = ?
                    AND cap2.status IN ('claimed', 'processing')) < 2
          )
        ORDER BY j.priority ASC,
                 COALESCE((
                     SELECT wb.realtime_factor
                     FROM worker_benchmarks wb
                     WHERE wb.worker_id = ?
                       AND CASE
                               WHEN wb.model = 'large-v3-turbo' THEN 'turbo'
                               ELSE wb.model
                           END = CASE
                               WHEN json_extract(j.input_config, '$.model') = 'large-v3-turbo' THEN 'turbo'
                               ELSE json_extract(j.input_config, '$.model')
                           END
                     ORDER BY wb.created_at DESC
                     LIMIT 1
                 ), 0) DESC,
                 (
                     SELECT COUNT(*)
                     FROM shards completed
                     WHERE completed.job_id = s.job_id
                       AND completed.status = 'completed'
                 ) DESC,
                 j.created_at ASC,
                 s.shard_index ASC
        LIMIT 25
        """,
        [*params, worker_id],
    ).fetchall()

    for candidate in candidates:
        if _requested_model_identifier(candidate["input_config"]) not in available_models:
            continue

        row = db.execute(
            """
            UPDATE shards
            SET status = 'claimed',
                claimed_by = ?,
                claimed_at = datetime('now'),
                last_heartbeat = datetime('now'),
                attempts = attempts + 1
            WHERE shard_id = ? AND status = 'pending'
            RETURNING *
            """,
            (worker_id, candidate["shard_id"]),
        ).fetchone()
        if row:
            return row

    return None


def _load_worker_available_models(db, worker_id: str) -> set[str]:
    row = db.execute(
        "SELECT available_models FROM workers WHERE worker_id = ?",
        (worker_id,),
    ).fetchone()
    if not row or not row["available_models"]:
        return set()
    return {
        canonical_model_identifier(model, default_quantization="none")
        for model in json.loads(row["available_models"])
        if model
    }


def _requested_model_identifier(input_config_text: str | None) -> str:
    config = json.loads(input_config_text) if input_config_text else {}
    return canonical_model_identifier(
        config.get("model"),
        config.get("quantization"),
        default_quantization=DEFAULT_QUANTIZATION,
    )


def _canonicalize_benchmark_model(model: str | None) -> str:
    normalized = str(model or "").strip().lower()
    if normalized == "large-v3-turbo":
        return "turbo"
    return normalized


def release_claim(job_id: str, worker_id: str, *, shard_id: str | None = None, reason: str = "worker_release") -> bool:
    """Release a claimed job or shard back to the pending queue."""
    with get_db_ctx() as db:
        if shard_id:
            row = db.execute(
                "SELECT * FROM shards WHERE shard_id = ? AND claimed_by = ? AND status IN ('claimed', 'processing')",
                (shard_id, worker_id),
            ).fetchone()
            if not row:
                return False

            db.execute(
                """
                UPDATE shards
                SET status = 'pending',
                    claimed_by = NULL,
                    claimed_at = NULL,
                    last_heartbeat = NULL,
                    attempts = CASE WHEN attempts > 0 THEN attempts - 1 ELSE 0 END
                WHERE shard_id = ?
                """,
                (shard_id,),
            )
            if not row["canary_id"]:
                record_event(
                    db,
                    job_id,
                    "shard_released",
                    from_status=row["status"],
                    to_status="pending",
                    actor=worker_id,
                    detail=f"shard {row['shard_index']} ({shard_id}): {reason}",
                )
        else:
            row = db.execute(
                "SELECT * FROM jobs WHERE job_id = ? AND claimed_by = ? AND status IN ('claimed', 'processing')",
                (job_id, worker_id),
            ).fetchone()
            if not row:
                return False

            db.execute(
                """
                UPDATE jobs
                SET status = 'pending',
                    claimed_by = NULL,
                    claimed_at = NULL,
                    last_heartbeat = NULL,
                    attempts = CASE WHEN attempts > 0 THEN attempts - 1 ELSE 0 END
                WHERE job_id = ?
                """,
                (job_id,),
            )
            record_event(
                db,
                job_id,
                "released",
                from_status=row["status"],
                to_status="pending",
                actor=worker_id,
                detail=reason,
            )

        sync_worker_busy_state(db, worker_id, touch_last_poll=True, touch_last_heartbeat=True)

    return True


def complete_shard(shard_id: str, worker_id: str, result_data: str) -> dict | None:
    """Mark a shard as completed and check if the whole job is done."""
    with get_db_ctx() as db:
        result = db.execute(
            """
            UPDATE shards SET status = 'completed', result_data = ?, completed_at = datetime('now')
            WHERE shard_id = ? AND claimed_by = ? AND status IN ('claimed', 'processing')
            RETURNING *
            """,
            (result_data, shard_id, worker_id),
        ).fetchone()

        if not result:
            return None

        job_id = result["job_id"]

        # Skip audit events for canary shards to keep them hidden from clients
        if not result["canary_id"]:
            record_event(
                db, job_id, "shard_completed",
                from_status="processing", to_status="completed",
                actor=worker_id, detail=f"shard {result['shard_index']} ({shard_id})",
            )
            increment_shard_completed()
            observe_shard_processing_duration(result["claimed_at"], result["completed_at"])

        record_worker_completion(db, worker_id)

    # Record reputation outside the transaction — skip for canary shards
    # because validate_canary_result already handles canary reputation.
    if not result["canary_id"]:
        record_shard_outcome(worker_id, success=True)

    with get_db_ctx() as db:
        # Check if all non-canary shards are done (canary shards don't count)
        counts = db.execute(
            """
            SELECT
                COUNT(*) as total,
                SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) as completed,
                SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) as failed
            FROM shards WHERE job_id = ? AND canary_id IS NULL
            """,
            (job_id,),
        ).fetchone()

        return {
            "job_id": job_id,
            "shard_id": shard_id,
            "total_shards": counts["total"],
            "completed_shards": counts["completed"],
            "failed_shards": counts["failed"],
            "all_done": counts["completed"] + counts["failed"] == counts["total"],
        }


def fail_shard(shard_id: str, worker_id: str, error: str) -> dict | None:
    """Mark a shard as failed, requeue if attempts remain."""
    with get_db_ctx() as db:
        row = db.execute(
            "SELECT * FROM shards WHERE shard_id = ? AND claimed_by = ?",
            (shard_id, worker_id),
        ).fetchone()

        if not row:
            return None

        job_id = row["job_id"]

        if row["attempts"] >= row["max_attempts"]:
            db.execute(
                "UPDATE shards SET status = 'failed', error_message = ? WHERE shard_id = ?",
                (error, shard_id),
            )
            new_status = "failed"
            if not row["canary_id"]:
                increment_shard_failed()
            if not row["canary_id"]:
                record_event(
                    db, job_id, "shard_failed_permanent",
                    to_status="failed", actor=worker_id,
                    detail=f"shard {row['shard_index']} ({shard_id}): {error}",
                )
        else:
            db.execute(
                "UPDATE shards SET status = 'pending', claimed_by = NULL, claimed_at = NULL,"
                " error_message = ? WHERE shard_id = ?",
                (error, shard_id),
            )
            new_status = "pending"
            if not row["canary_id"]:
                record_event(
                    db, job_id, "shard_failed_retry",
                    to_status="pending", actor=worker_id,
                    detail=f"shard {row['shard_index']} ({shard_id}): {error}",
                )

        record_worker_failure(db, worker_id, job_id, error)

    # Skip reputation recording for canary shards — canary validation handles it
    if not row["canary_id"]:
        record_shard_outcome(worker_id, success=False)

    return {
        "shard_id": shard_id,
        "status": new_status,
        "attempts": row["attempts"],
        "max_attempts": row["max_attempts"],
        "is_canary": bool(row["canary_id"]),
    }


def reap_stale_jobs(stale_minutes: int = 5) -> int:
    """Requeue jobs and shards with expired heartbeats. Returns count of reaped items."""
    reaped = 0

    with get_db_ctx() as db:
        # Reap stale non-sharded jobs
        rows = db.execute(
            """
            SELECT job_id, attempts, max_attempts, claimed_by, created_at
            FROM jobs
            WHERE status IN ('claimed', 'processing')
              AND shard_count IS NULL
              AND last_heartbeat < datetime('now', ?)
            """,
            (f"-{stale_minutes} minutes",),
        ).fetchall()

        for row in rows:
            if row["attempts"] >= row["max_attempts"]:
                db.execute(
                    "UPDATE jobs SET status = 'failed',"
                    " error_message = 'Max attempts exceeded (stale)', completed_at = datetime('now') WHERE job_id = ?",
                    (row["job_id"],),
                )
                increment_job_failed()
                now_row = db.execute("SELECT completed_at FROM jobs WHERE job_id = ?", (row["job_id"],)).fetchone()
                observe_job_turnaround_duration(row["created_at"], now_row["completed_at"])
                record_event(
                    db, row["job_id"], "reaped_failed", to_status="failed",
                    actor="reaper", detail="stale heartbeat, max attempts exceeded",
                )
            else:
                db.execute(
                    "UPDATE jobs SET status = 'pending', claimed_by = NULL, claimed_at = NULL WHERE job_id = ?",
                    (row["job_id"],),
                )
                record_event(
                    db, row["job_id"], "reaped_requeued", to_status="pending",
                    actor="reaper", detail="stale heartbeat, requeued for retry",
                )
            reaped += 1
            if row["claimed_by"]:
                sync_worker_busy_state(db, row["claimed_by"])

        # Reap stale shards
        shard_rows = db.execute(
            """
            SELECT shard_id, job_id, shard_index, attempts, max_attempts, claimed_by, canary_id
            FROM shards
            WHERE status IN ('claimed', 'processing', 'fallback')
              AND last_heartbeat < datetime('now', ?)
            """,
            (f"-{stale_minutes} minutes",),
        ).fetchall()

        for row in shard_rows:
            if row["attempts"] >= row["max_attempts"]:
                db.execute(
                    "UPDATE shards SET status = 'failed',"
                    " error_message = 'Max attempts exceeded (stale)' WHERE shard_id = ?",
                    (row["shard_id"],),
                )
                if not row["canary_id"]:
                    increment_shard_failed()
                    record_event(
                        db, row["job_id"], "shard_reaped_failed", to_status="failed",
                        actor="reaper", detail=f"shard {row['shard_index']} stale, max attempts",
                    )
            else:
                db.execute(
                    "UPDATE shards SET status = 'pending', claimed_by = NULL, claimed_at = NULL WHERE shard_id = ?",
                    (row["shard_id"],),
                )
                if not row["canary_id"]:
                    record_event(
                        db, row["job_id"], "shard_reaped_requeued", to_status="pending",
                        actor="reaper", detail=f"shard {row['shard_index']} stale, requeued",
                    )
            reaped += 1
            if row["claimed_by"]:
                sync_worker_busy_state(db, row["claimed_by"])

    if reaped:
        logger.info("Reaped stale items", count=reaped)
    return reaped
