from __future__ import annotations

import json
import sqlite3
import threading
from collections.abc import Iterable
from pathlib import Path

from ai_orchestrator.schemas.task_state import TaskState
from ai_orchestrator.schemas.task_status import TaskLifecycleStatus


class SQLiteTaskStore:
    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 tasks (
                    task_id TEXT PRIMARY KEY,
                    status TEXT NOT NULL,
                    payload TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                )
                """
            )
            connection.commit()

    def create(self, task: TaskState) -> None:
        payload = task.model_dump(mode="json")
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                INSERT INTO tasks (task_id, status, payload, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (
                    task.task_id,
                    task.status.value,
                    json.dumps(payload),
                    task.timestamps.created_at.isoformat(),
                    task.timestamps.updated_at.isoformat(),
                ),
            )
            connection.commit()

    def save(self, task: TaskState) -> None:
        task.touch()
        payload = task.model_dump(mode="json")
        with self._lock, self._connect() as connection:
            connection.execute(
                """
                UPDATE tasks
                SET status = ?, payload = ?, updated_at = ?
                WHERE task_id = ?
                """,
                (
                    task.status.value,
                    json.dumps(payload),
                    task.timestamps.updated_at.isoformat(),
                    task.task_id,
                ),
            )
            connection.commit()

    def get(self, task_id: str) -> TaskState | None:
        with self._lock, self._connect() as connection:
            row = connection.execute(
                "SELECT payload FROM tasks WHERE task_id = ?",
                (task_id,),
            ).fetchone()
        if row is None:
            return None
        return TaskState.model_validate_json(row["payload"])

    def list_by_statuses(self, statuses: Iterable[TaskLifecycleStatus]) -> list[TaskState]:
        status_values = [status.value for status in statuses]
        if not status_values:
            return []
        placeholders = ", ".join("?" for _ in status_values)
        query = f"SELECT payload FROM tasks WHERE status IN ({placeholders})"
        with self._lock, self._connect() as connection:
            rows = connection.execute(query, status_values).fetchall()
        return [TaskState.model_validate_json(row["payload"]) for row in rows]
