from __future__ import annotations

import os
from datetime import datetime, timezone
from typing import Optional

import aiosqlite
from pydantic import BaseModel


VALID_FRAME_COUNTS = {8 * n + 1 for n in range(1, 129)}
VALID_STATUSES = {"pending", "processing", "completed", "failed"}


class JobRecord(BaseModel):
    id: str
    status: str = "pending"
    prompt: Optional[str] = None
    params: dict | None = None
    error: Optional[str] = None
    created_at: str
    completed_at: Optional[str] = None


class JobDB:
    def __init__(self, db_path: str) -> None:
        self.db_path = db_path

    async def init(self) -> None:
        os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    id TEXT PRIMARY KEY,
                    status TEXT NOT NULL DEFAULT 'pending',
                    prompt TEXT,
                    params TEXT,
                    error TEXT,
                    created_at TEXT NOT NULL,
                    completed_at TEXT
                )
                """
            )
            await db.commit()

    async def create_job(self, job: JobRecord) -> None:
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                "INSERT INTO jobs (id, status, prompt, params, created_at) VALUES (?, ?, ?, ?, ?)",
                (
                    job.id,
                    job.status,
                    job.prompt,
                    _json_encode(job.params),
                    job.created_at,
                ),
            )
            await db.commit()

    async def get_job(self, job_id: str) -> JobRecord | None:
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)) as cursor:
                row = await cursor.fetchone()
                if row is None:
                    return None
                return JobRecord(
                    id=row["id"],
                    status=row["status"],
                    prompt=row["prompt"],
                    params=_json_decode(row["params"]),
                    error=row["error"],
                    created_at=row["created_at"],
                    completed_at=row["completed_at"],
                )

    async def update_status(
        self,
        job_id: str,
        status: str,
        error: Optional[str] = None,
    ) -> None:
        now = datetime.now(timezone.utc).isoformat()
        async with aiosqlite.connect(self.db_path) as db:
            if status in {"completed", "failed"}:
                await db.execute(
                    "UPDATE jobs SET status = ?, error = ?, completed_at = ? WHERE id = ?",
                    (status, error, now, job_id),
                )
            else:
                await db.execute(
                    "UPDATE jobs SET status = ?, error = ? WHERE id = ?",
                    (status, error, job_id),
                )
            await db.commit()

    async def delete_job(self, job_id: str) -> bool:
        async with aiosqlite.connect(self.db_path) as db:
            cur = await db.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
            await db.commit()
            return cur.rowcount > 0


def _json_encode(obj: dict | None) -> str | None:
    if obj is None:
        return None
    import json

    return json.dumps(obj)


def _json_decode(s: str | None) -> dict | None:
    if s is None:
        return None
    import json

    return json.loads(s)