from __future__ import annotations

import json
import sqlite3
import threading
from collections.abc import Iterable
from datetime import timedelta
from pathlib import Path
from uuid import uuid4

from ai_orchestrator.schemas.queue_models import QueueJob, QueueJobStatus
from ai_orchestrator.schemas.task_state import utc_now


class SQLiteQueueStore:
    def __init__(self, database_path: Path) -> None:
        self.database_path = database_path
        self.database_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._initialize()

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.database_path)
        connection.row_factory = sqlite3.Row
        return connection

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    job_id TEXT PRIMARY KEY,
                    task_id TEXT NOT NULL,
                    engine_name TEXT NOT NULL,
                    status TEXT NOT NULL,
                    attempts INTEGER NOT NULL,
                    max_attempts INTEGER NOT NULL,
                    worker_id TEXT,
                    lease_until TEXT,
                    available_at TEXT NOT NULL,
                    last_heartbeat TEXT,
                    error_message TEXT,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                )
                """
            )
            connection.execute(
                "CREATE INDEX IF NOT EXISTS idx_jobs_status_available ON jobs (status, available_at)"
            )
            connection.commit()

    def enqueue(self, *, task_id: str, engine_name: str, max_attempts: int) -> QueueJob:
        now = utc_now()
        job = QueueJob(
            job_id=f"job_{uuid4().hex[:12]}",
            task_id=task_id,
            engine_name=engine_name,
            status=QueueJobStatus.queued,
            attempts=0,
            max_attempts=max_attempts,
            available_at=now,
            created_at=now,
            updated_at=now,
        )
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                INSERT INTO jobs (
                    job_id, task_id, engine_name, status, attempts, max_attempts,
                    worker_id, lease_until, available_at, last_heartbeat, error_message,
                    created_at, updated_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    job.job_id,
                    job.task_id,
                    job.engine_name,
                    job.status.value,
                    job.attempts,
                    job.max_attempts,
                    job.worker_id,
                    None,
                    job.available_at.isoformat(),
                    None,
                    None,
                    job.created_at.isoformat(),
                    job.updated_at.isoformat(),
                ),
            )
            connection.commit()
        return job

    def lease_next(self, *, worker_id: str, lease_seconds: int) -> QueueJob | None:
        now = utc_now()
        with self._lock, self._connect() as connection:
            row = connection.execute(
                """
                SELECT * FROM jobs
                WHERE (
                    status = ? AND available_at <= ?
                ) OR (
                    status = ? AND lease_until IS NOT NULL AND lease_until <= ?
                )
                ORDER BY created_at ASC
                LIMIT 1
                """,
                (
                    QueueJobStatus.queued.value,
                    now.isoformat(),
                    QueueJobStatus.leased.value,
                    now.isoformat(),
                ),
            ).fetchone()
            if row is None:
                return None

            attempts = int(row["attempts"]) + 1
            lease_until = now + timedelta(seconds=lease_seconds)
            connection.execute(
                """
                UPDATE jobs
                SET status = ?, worker_id = ?, attempts = ?, lease_until = ?, last_heartbeat = ?, updated_at = ?
                WHERE job_id = ?
                """,
                (
                    QueueJobStatus.leased.value,
                    worker_id,
                    attempts,
                    lease_until.isoformat(),
                    now.isoformat(),
                    now.isoformat(),
                    row["job_id"],
                ),
            )
            connection.commit()
            leased = dict(row)
            leased.update(
                {
                    "status": QueueJobStatus.leased.value,
                    "worker_id": worker_id,
                    "attempts": attempts,
                    "lease_until": lease_until.isoformat(),
                    "last_heartbeat": now.isoformat(),
                    "updated_at": now.isoformat(),
                }
            )
        return self._row_to_job(leased)

    def heartbeat(self, *, job_id: str, worker_id: str, lease_seconds: int) -> None:
        now = utc_now()
        lease_until = now + timedelta(seconds=lease_seconds)
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                UPDATE jobs
                SET lease_until = ?, last_heartbeat = ?, updated_at = ?
                WHERE job_id = ? AND worker_id = ? AND status = ?
                """,
                (
                    lease_until.isoformat(),
                    now.isoformat(),
                    now.isoformat(),
                    job_id,
                    worker_id,
                    QueueJobStatus.leased.value,
                ),
            )
            connection.commit()

    def mark_completed(self, job_id: str) -> None:
        self._mark_terminal(job_id, QueueJobStatus.completed, None)

    def mark_failed(self, job_id: str, error_message: str) -> None:
        self._mark_terminal(job_id, QueueJobStatus.failed, error_message)

    def mark_cancelled(self, job_id: str) -> None:
        self._mark_terminal(job_id, QueueJobStatus.cancelled, None)

    def cancel_for_task(self, task_id: str) -> None:
        now = utc_now().isoformat()
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                UPDATE jobs
                SET status = ?, updated_at = ?
                WHERE task_id = ? AND status IN (?, ?)
                """,
                (
                    QueueJobStatus.cancelled.value,
                    now,
                    task_id,
                    QueueJobStatus.queued.value,
                    QueueJobStatus.leased.value,
                ),
            )
            connection.commit()

    def requeue_expired(self) -> int:
        now = utc_now()
        with self._lock, self._connect() as connection:
            cursor = connection.execute(
                """
                UPDATE jobs
                SET status = ?, worker_id = NULL, lease_until = NULL, updated_at = ?
                WHERE status = ? AND lease_until IS NOT NULL AND lease_until <= ?
                """,
                (
                    QueueJobStatus.queued.value,
                    now.isoformat(),
                    QueueJobStatus.leased.value,
                    now.isoformat(),
                ),
            )
            connection.commit()
        return int(cursor.rowcount or 0)

    def get_by_task_id(self, task_id: str) -> list[QueueJob]:
        with self._lock, self._connect() as connection:
            rows = connection.execute(
                "SELECT * FROM jobs WHERE task_id = ? ORDER BY created_at ASC",
                (task_id,),
            ).fetchall()
        return [self._row_to_job(dict(row)) for row in rows]

    def list_by_statuses(self, statuses: Iterable[QueueJobStatus]) -> list[QueueJob]:
        values = [status.value for status in statuses]
        if not values:
            return []
        placeholders = ", ".join("?" for _ in values)
        query = f"SELECT * FROM jobs WHERE status IN ({placeholders}) ORDER BY created_at ASC"
        with self._lock, self._connect() as connection:
            rows = connection.execute(query, values).fetchall()
        return [self._row_to_job(dict(row)) for row in rows]

    def _mark_terminal(self, job_id: str, status: QueueJobStatus, error_message: str | None) -> None:
        now = utc_now().isoformat()
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                UPDATE jobs
                SET status = ?, error_message = ?, updated_at = ?, lease_until = NULL
                WHERE job_id = ?
                """,
                (
                    status.value,
                    error_message,
                    now,
                    job_id,
                ),
            )
            connection.commit()

    @staticmethod
    def _row_to_job(row: dict) -> QueueJob:
        normalized = dict(row)
        for key in ("lease_until", "available_at", "last_heartbeat", "created_at", "updated_at"):
            if normalized.get(key) is None:
                continue
        return QueueJob.model_validate(normalized)
