"""SQLite 连接与建表。"""

from __future__ import annotations

import sqlite3
from contextlib import contextmanager
from typing import Iterator

from . import config

_SCHEMA = """
CREATE TABLE IF NOT EXISTS providers (
    id            TEXT PRIMARY KEY,
    name          TEXT NOT NULL,
    adapter       TEXT NOT NULL,
    base_url      TEXT NOT NULL DEFAULT '',
    api_key       TEXT NOT NULL DEFAULT '',
    default_model TEXT NOT NULL DEFAULT '',
    enabled       INTEGER NOT NULL DEFAULT 1,
    note          TEXT NOT NULL DEFAULT '',
    created_at    TEXT NOT NULL,
    updated_at    TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS models (
    id                   TEXT PRIMARY KEY,
    provider_id          TEXT NOT NULL,
    model_id             TEXT NOT NULL,
    display_name         TEXT NOT NULL DEFAULT '',
    enabled              INTEGER NOT NULL DEFAULT 0,
    kind                 TEXT NOT NULL DEFAULT 'image',
    size_hint            TEXT NOT NULL DEFAULT '',
    supports_image_input INTEGER NOT NULL DEFAULT 0,
    note                 TEXT NOT NULL DEFAULT '',
    discovered_at        TEXT NOT NULL DEFAULT '',
    created_at           TEXT NOT NULL,
    updated_at           TEXT NOT NULL,
    FOREIGN KEY (provider_id) REFERENCES providers(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS generations (
    id            TEXT PRIMARY KEY,
    provider_id   TEXT NOT NULL DEFAULT '',
    provider_name TEXT NOT NULL DEFAULT '',
    model         TEXT NOT NULL DEFAULT '',
    profile_name  TEXT NOT NULL DEFAULT 'user',
    prompt        TEXT NOT NULL DEFAULT '',
    size          TEXT NOT NULL DEFAULT '',
    ratio         TEXT NOT NULL DEFAULT '',
    n             INTEGER NOT NULL DEFAULT 1,
    status        TEXT NOT NULL DEFAULT 'success',
    error         TEXT NOT NULL DEFAULT '',
    files         TEXT NOT NULL DEFAULT '[]',
    duration_ms   INTEGER NOT NULL DEFAULT 0,
    created_at    TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_models_provider ON models(provider_id);
CREATE INDEX IF NOT EXISTS idx_generations_created ON generations(created_at DESC);
"""


def connect() -> sqlite3.Connection:
    config.ensure_dirs()
    conn = sqlite3.connect(config.DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


@contextmanager
def session() -> Iterator[sqlite3.Connection]:
    conn = connect()
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def init_db() -> None:
    with session() as conn:
        conn.executescript(_SCHEMA)