"""Worker-facing endpoints — registration, polling, heartbeat, result, fail."""

from __future__ import annotations

import json
from typing import Annotated, Any

from fastapi import APIRouter, Depends, Query, Request
from fastapi.responses import JSONResponse, Response

from hive_compute.coordinator.auth import authorize_worker_registration, require_worker_token
from hive_compute.coordinator.billing import create_billing_item
from hive_compute.coordinator.canaries import inject_canary_shard, should_inject_canary, validate_canary_result
from hive_compute.coordinator.database import get_db_ctx, record_event
from hive_compute.coordinator.fallback import check_and_trigger_fallback
from hive_compute.coordinator.manager_auth import require_manager_admin
from hive_compute.coordinator.merger import merge_job_shards
from hive_compute.coordinator.metrics import (
    increment_job_completed,
    increment_job_failed,
    increment_shard_claimed,
    observe_job_processing_duration,
    observe_job_turnaround_duration,
)
from hive_compute.coordinator.reputation import get_all_reputations, get_reputation
from hive_compute.coordinator.scheduler import claim_next_job, complete_shard, fail_shard, release_claim
from hive_compute.coordinator.webhooks import fire_webhook
from hive_compute.coordinator.worker_api_models import WorkerListResponse
from hive_compute.coordinator.worker_credentials import mint_worker_token
from hive_compute.coordinator.worker_observability import (
    list_workers_with_observability,
    record_worker_completion,
    record_worker_failure,
    record_worker_poll,
)
from hive_compute.models import (
    BenchmarkReportRequest,
    FailRequest,
    PollResponse,
    ResultSubmitRequest,
    WorkerRegisterRequest,
    WorkerRegisterResponse,
)

router = APIRouter(tags=["workers"])
ManagerAdmin = Annotated[dict[str, Any], Depends(require_manager_admin)]
BENCHMARK_HISTORY_LIMIT_PER_MODEL = 10


@router.get("/workers", response_model=WorkerListResponse)
def list_workers(_admin: ManagerAdmin):
    """List all registered workers and their status."""
    workers = list_workers_with_observability()
    return WorkerListResponse(workers=workers, total=len(workers))


@router.get("/workers/reputation")
def list_reputations(_admin: ManagerAdmin):
    """Get reputation scores for all workers."""
    return {"workers": get_all_reputations()}


@router.get("/workers/me/reputation")
def get_my_reputation(worker_id: str = Depends(require_worker_token)):
    """Get reputation info for the authenticated worker."""
    rep = get_reputation(worker_id)
    if not rep:
        return JSONResponse(status_code=404, content={"error": "Worker not found or no reputation data"})
    return rep


@router.get("/workers/{worker_id}/reputation")
def get_worker_reputation(worker_id: str, _admin: ManagerAdmin):
    """Get reputation info for a specific worker."""
    rep = get_reputation(worker_id)
    if not rep:
        return JSONResponse(status_code=404, content={"error": "Worker not found or no reputation data"})
    return rep


@router.post("/workers/register", response_model=WorkerRegisterResponse)
def register_worker(req: WorkerRegisterRequest, request: Request):
    """Register a new worker or rotate credentials for an existing one."""
    worker_token, worker_token_hash, worker_token_last_four = mint_worker_token()

    with get_db_ctx() as db:
        existing = db.execute(
            "SELECT worker_id, worker_token, worker_token_revoked_at FROM workers WHERE worker_id = ?",
            (req.worker_id,),
        ).fetchone()

        if existing:
            enrollment_token = request.headers.get("X-Enrollment-Token")
            used_enrollment_token = False
            if enrollment_token:
                auth_info = authorize_worker_registration(request, req.worker_id)
                used_enrollment_token = auth_info.get("authType") == "enrollment_token"

            # Re-registration: enrollment tokens bypass identity check (admin trust).
            # Otherwise the existing worker token is required to prevent takeover.
            if not used_enrollment_token:
                reauth_err = "Worker already registered. Provide current token to rotate."
                auth_header = request.headers.get("Authorization", "")
                if not auth_header.startswith("Bearer "):
                    return JSONResponse(
                        status_code=403, content={"error": reauth_err},
                    )
                from hive_compute.coordinator.worker_credentials import hash_worker_token
                token_match = hash_worker_token(auth_header[7:]) == existing["worker_token"]
                token_revoked = existing["worker_token_revoked_at"] is not None
                if not token_match or token_revoked:
                    return JSONResponse(
                        status_code=403, content={"error": reauth_err},
                    )
            db.execute(
                """
                UPDATE workers
                SET chip = ?, memory_gb = ?, gpu_cores = ?,
                    available_handlers = ?, available_models = ?, model_capabilities = ?,
                    status = 'idle', last_heartbeat = datetime('now'),
                    capabilities_advertised_at = datetime('now'),
                    capabilities_confirmed_at = datetime('now'),
                    worker_token = ?, worker_token_last_four = ?, worker_token_issued_at = datetime('now'),
                    last_auth_failure_at = NULL, last_auth_failure_detail = NULL,
                    worker_token_revoked_at = NULL
                WHERE worker_id = ?
                """,
                (
                    req.chip,
                    req.memory_gb,
                    req.gpu_cores,
                    json.dumps(req.available_handlers),
                    json.dumps(req.available_models),
                    json.dumps(req.model_capabilities),
                    worker_token_hash,
                    worker_token_last_four,
                    req.worker_id,
                ),
            )
            return WorkerRegisterResponse(
                worker_id=req.worker_id,
                worker_token=worker_token,
                message="Worker credential rotated",
            )

        authorize_worker_registration(request, req.worker_id)
        db.execute(
            """
            INSERT INTO workers (
                worker_id, worker_token, worker_token_last_four, worker_token_issued_at,
                chip, memory_gb, gpu_cores, available_handlers, available_models, model_capabilities,
                capabilities_advertised_at, capabilities_confirmed_at, status
            )
            VALUES (?, ?, ?, datetime('now'), ?, ?, ?, ?, ?, ?, datetime('now'), datetime('now'), 'idle')
            """,
            (
                req.worker_id,
                worker_token_hash,
                worker_token_last_four,
                req.chip,
                req.memory_gb,
                req.gpu_cores,
                json.dumps(req.available_handlers),
                json.dumps(req.available_models),
                json.dumps(req.model_capabilities),
            ),
        )

    return WorkerRegisterResponse(
        worker_id=req.worker_id,
        worker_token=worker_token,
        message="Worker registered",
    )


@router.put("/workers/me")
def update_worker(req: WorkerRegisterRequest, worker_id: str = Depends(require_worker_token)):
    """Update worker info on daemon startup (uses worker token, not enrollment key)."""
    with get_db_ctx() as db:
        db.execute(
            """
            UPDATE workers SET chip = ?, memory_gb = ?, gpu_cores = ?,
                available_handlers = ?, available_models = ?, model_capabilities = ?,
                status = 'idle', last_heartbeat = datetime('now'),
                capabilities_advertised_at = datetime('now'),
                capabilities_confirmed_at = datetime('now')
            WHERE worker_id = ?
            """,
            (
                req.chip,
                req.memory_gb,
                req.gpu_cores,
                json.dumps(req.available_handlers),
                json.dumps(req.available_models),
                json.dumps(req.model_capabilities),
                worker_id,
            ),
        )
    return {"worker_id": worker_id, "message": "Worker updated"}


@router.get("/workers/me/payouts")
def get_my_payouts(worker_id: str = Depends(require_worker_token)):
    """Return this worker's payout summary and manual payout history."""
    from hive_compute.coordinator.payouts import get_worker_earnings, list_payout_disbursements

    earnings = get_worker_earnings(worker_id)
    held_cents = earnings.get("held", 0)
    released_cents = earnings.get("released", 0)
    total_cents = earnings.get("total", held_cents + released_cents)
    return {
        "pending_usd": round(held_cents / 100, 2),
        "released_usd": round(released_cents / 100, 2),
        "total_usd": round(total_cents / 100, 2),
        "item_count": earnings.get("items", 0),
        "history": list_payout_disbursements(worker_id),
    }


@router.post("/workers/me/benchmarks")
def submit_my_benchmarks(
    req: BenchmarkReportRequest,
    worker_id: str = Depends(require_worker_token),
):
    """Store worker-reported local benchmark results."""
    with get_db_ctx() as db:
        models_seen: set[str] = set()
        for result in req.results:
            models_seen.add(result.model)
            db.execute(
                """
                INSERT INTO worker_benchmarks (
                    worker_id, benchmark_created_at, audio_path, machine_json, environment_json,
                    model, backend, language, audio_seconds, elapsed_seconds,
                    realtime_factor, memory_rss_mb, time_to_first_token_seconds,
                    word_count, transcript_preview, error
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    worker_id,
                    req.created_at,
                    req.audio_path,
                    json.dumps(req.machine),
                    json.dumps(req.environment),
                    result.model,
                    result.backend,
                    result.language,
                    result.audio_seconds,
                    result.elapsed_seconds,
                    result.realtime_factor,
                    result.memory_rss_mb,
                    result.time_to_first_token_seconds,
                    result.word_count,
                    result.transcript_preview,
                    result.error,
                ),
            )
        for model in models_seen:
            db.execute(
                """
                DELETE FROM worker_benchmarks
                WHERE worker_id = ? AND model = ?
                  AND benchmark_id NOT IN (
                      SELECT benchmark_id
                      FROM worker_benchmarks
                      WHERE worker_id = ? AND model = ?
                      ORDER BY benchmark_created_at DESC, benchmark_id DESC
                      LIMIT ?
                  )
                """,
                (
                    worker_id,
                    model,
                    worker_id,
                    model,
                    BENCHMARK_HISTORY_LIMIT_PER_MODEL,
                ),
            )
    return {"worker_id": worker_id, "stored_results": len(req.results)}


@router.get("/workers/poll")
def poll_for_job(
    worker_id: str = Depends(require_worker_token),
    types: str = Query(default="transcribe"),
):
    """Poll for the next available job or shard. Returns 204 if none available."""
    job_types = [t.strip() for t in types.split(",")]
    with get_db_ctx() as db:
        record_worker_poll(db, worker_id)
    job = claim_next_job(worker_id, job_types)

    if not job:
        return Response(status_code=204)

    # Canary injection: occasionally give the worker a hidden known-answer test.
    # Swap both shard_id and audio_url so the worker submits against the canary
    # shard, and release the original shard back to pending for another worker.
    if job.get("shard_id") and should_inject_canary(worker_id):
        canary = inject_canary_shard(job["job_id"], worker_id)
        if canary:
            original_shard_id = job["shard_id"]
            with get_db_ctx() as db:
                db.execute(
                    "UPDATE shards SET status = 'pending', claimed_by = NULL, "
                    "claimed_at = NULL, attempts = attempts - 1 WHERE shard_id = ?",
                    (original_shard_id,),
                )
                # Remove only the most recent shard_claimed event for this shard
                # to avoid leaking the substitution, preserving earlier valid claims.
                # Use ROWID (monotonic in SQLite) instead of created_at to avoid
                # second-resolution ties.
                db.execute(
                    "DELETE FROM job_events WHERE event_id = ("
                    "  SELECT event_id FROM job_events"
                    "  WHERE job_id = ? AND event_type = 'shard_claimed' AND detail LIKE ?"
                    "  ORDER BY ROWID DESC LIMIT 1"
                    ")",
                    (job["job_id"], f"%{original_shard_id}%"),
                )
            job["shard_id"] = canary["shard_id"]
            job["audio_url"] = canary["audio_url"]
            job["start_seconds"] = 0.0
            job["end_seconds"] = canary["duration_seconds"]
            job["canary_id"] = canary["canary_id"]
            # Override language so worker transcribes canary in the correct language
            cfg = json.loads(job["input_config"]) if job["input_config"] else {}
            cfg["language"] = canary.get("language", "en")
            job["input_config"] = json.dumps(cfg)

    if job.get("shard_id") and not job.get("canary_id"):
        increment_shard_claimed()

    input_config = json.loads(job["input_config"]) if job["input_config"] else {}
    return PollResponse(
        job_id=job["job_id"],
        job_type=job["job_type"],
        audio_url=job["audio_url"],
        input_config=input_config,
        timeout_seconds=job["timeout_seconds"],
        shard_id=job.get("shard_id"),
        start_seconds=job.get("start_seconds"),
        end_seconds=job.get("end_seconds"),
    )


@router.post("/jobs/{job_id}/release")
def release_job(
    job_id: str,
    worker_id: str = Depends(require_worker_token),
    shard_id: str | None = Query(default=None),
    reason: str = Query(default="worker_release"),
):
    """Release a claimed job or shard that the worker no longer intends to run."""
    released = release_claim(job_id, worker_id, shard_id=shard_id, reason=reason)
    if not released:
        return JSONResponse(status_code=404, content={"error": "Job not found or not owned by this worker"})
    return {"ok": True}


@router.post("/jobs/{job_id}/heartbeat")
def heartbeat(
    job_id: str,
    worker_id: str = Depends(require_worker_token),
    shard_id: str | None = Query(default=None),
):
    """Update heartbeat for a claimed job or shard."""
    with get_db_ctx() as db:
        if shard_id:
            result = db.execute(
                """
                UPDATE shards SET last_heartbeat = datetime('now'), status = 'processing'
                WHERE shard_id = ? AND claimed_by = ? AND status IN ('claimed', 'processing')
                """,
                (shard_id, worker_id),
            )
            if result.rowcount == 0:
                return JSONResponse(status_code=404, content={"error": "Shard not found or not owned by this worker"})
        else:
            result = db.execute(
                """
                UPDATE jobs SET last_heartbeat = datetime('now'), status = 'processing'
                WHERE job_id = ? AND claimed_by = ? AND status IN ('claimed', 'processing')
                """,
                (job_id, worker_id),
            )
            if result.rowcount == 0:
                return JSONResponse(status_code=404, content={"error": "Job not found or not owned by this worker"})
        # Skip heartbeat events for canary shards to keep them hidden
        is_canary = False
        if shard_id:
            canary_check = db.execute(
                "SELECT canary_id FROM shards WHERE shard_id = ?", (shard_id,),
            ).fetchone()
            is_canary = bool(canary_check and canary_check["canary_id"])
        if not is_canary:
            record_event(db, job_id, "heartbeat", to_status="processing", actor=worker_id)

        # Also refresh the worker's own heartbeat so observability queries
        # correctly detect this worker as active and the capability snapshot
        # as recently confirmed by a live worker process.
        db.execute(
            """
            UPDATE workers
            SET last_heartbeat = datetime('now'),
                status = 'busy',
                capabilities_confirmed_at = datetime('now')
            WHERE worker_id = ?
            """,
            (worker_id,),
        )

    return {"ok": True}


@router.post("/jobs/{job_id}/result")
def submit_result(
    job_id: str,
    req: ResultSubmitRequest,
    worker_id: str = Depends(require_worker_token),
    shard_id: str | None = Query(default=None),
):
    """Submit completed transcription result for a job or shard."""
    result_data = json.dumps(req.model_dump())

    if shard_id:
        # Shard-level result — complete_shard verifies ownership first
        shard_result = complete_shard(shard_id, worker_id, result_data)
        if not shard_result:
            return JSONResponse(status_code=404, content={"error": "Shard not found or not owned by this worker"})

        # Validate canary after ownership is confirmed, preventing replay attacks
        validate_canary_result(shard_id, worker_id, req.transcript)

        # If all shards are done, trigger merge
        if shard_result["all_done"]:
            merge_job_shards(job_id)

        return {
            "job_id": job_id,
            "shard_id": shard_id,
            "status": "completed",
            "shards_completed": shard_result["completed_shards"],
            "shards_total": shard_result["total_shards"],
        }

    # Non-sharded job result (backward compatible)
    with get_db_ctx() as db:
        result = db.execute(
            """
            UPDATE jobs
            SET status = 'completed',
                result_data = ?,
                completed_at = datetime('now'),
                audio_duration = COALESCE(audio_duration, ?)
            WHERE job_id = ? AND claimed_by = ? AND status IN ('claimed', 'processing')
            RETURNING partner_id, audio_duration, created_at, completed_at
            """,
            (result_data, req.duration_seconds, job_id, worker_id),
        ).fetchone()
        if not result:
            return JSONResponse(status_code=404, content={"error": "Job not found or not owned by this worker"})
        record_event(db, job_id, "completed", from_status="processing", to_status="completed", actor=worker_id)

        record_worker_completion(db, worker_id)
        increment_job_completed()
        observe_job_turnaround_duration(result["created_at"], result["completed_at"])
        observe_job_processing_duration(result["created_at"], result["completed_at"])

    if result["partner_id"] and result["audio_duration"]:
        create_billing_item(job_id, result["partner_id"], result["audio_duration"])

    fire_webhook(job_id)
    return {"job_id": job_id, "status": "completed"}


@router.post("/jobs/{job_id}/fail")
def report_failure(
    job_id: str,
    req: FailRequest,
    worker_id: str = Depends(require_worker_token),
    shard_id: str | None = Query(default=None),
):
    """Report job or shard failure. Requeues if attempts remain."""
    if shard_id:
        # Shard-level failure
        shard_result = fail_shard(shard_id, worker_id, req.error)
        if not shard_result:
            return JSONResponse(status_code=404, content={"error": "Shard not found or not owned by this worker"})

        # If shard permanently failed, try fallback before failing the job.
        # Skip fallback for canary shards — they're validation-only, not real work.
        if shard_result["status"] == "failed" and not shard_result.get("is_canary"):
            fallback = check_and_trigger_fallback(shard_id)
            if not fallback:
                merge_job_shards(job_id)  # No fallback available — fail the job

        return {
            "job_id": job_id,
            "shard_id": shard_id,
            "status": shard_result["status"],
            "attempts": shard_result["attempts"],
            "max_attempts": shard_result["max_attempts"],
        }

    # Non-sharded job failure (backward compatible)
    with get_db_ctx() as db:
        row = db.execute(
            "SELECT attempts, max_attempts FROM jobs WHERE job_id = ? AND claimed_by = ?",
            (job_id, worker_id),
        ).fetchone()

        if not row:
            return JSONResponse(status_code=404, content={"error": "Job not found or not owned by this worker"})

        if row["attempts"] >= row["max_attempts"]:
            db.execute(
                "UPDATE jobs SET status = 'failed', error_message = ?, completed_at = datetime('now') WHERE job_id = ?",
                (req.error, job_id),
            )
            new_status = "failed"
            record_event(db, job_id, "failed_permanent", to_status="failed", actor=worker_id, detail=req.error)
        else:
            db.execute(
                "UPDATE jobs SET status = 'pending', claimed_by = NULL, claimed_at = NULL,"
                " error_message = ? WHERE job_id = ?",
                (req.error, job_id),
            )
            new_status = "pending"
            record_event(db, job_id, "failed_retry", to_status="pending", actor=worker_id, detail=req.error)

        record_worker_failure(db, worker_id, job_id, req.error)

    if new_status == "failed":
        increment_job_failed()
        with get_db_ctx() as db:
            failed_job = db.execute(
                "SELECT created_at, completed_at FROM jobs WHERE job_id = ?",
                (job_id,),
            ).fetchone()
        observe_job_turnaround_duration(failed_job["created_at"], failed_job["completed_at"])

    if new_status == "failed":
        fire_webhook(job_id)

    return {"job_id": job_id, "status": new_status, "attempts": row["attempts"], "max_attempts": row["max_attempts"]}
