"""Persistent history and background execution for quality runs."""
from __future__ import annotations

import asyncio
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
import json
from pathlib import Path
import sqlite3
from typing import Any, Awaitable, Callable
from uuid import uuid4

from logos_adk.config import TaskQueueConfig
from logos_adk.errors import ConfigError, StateBackendError
from logos_adk.persistence import (
    get_postgresql_schema_version,
    get_sqlite_schema_version,
    set_postgresql_schema_version,
    set_sqlite_schema_version,
)
from logos_adk.postgres_utils import import_psycopg, resolve_database_url


QUALITY_HISTORY_SCHEMA_KIND = "quality_history"
QUALITY_HISTORY_SCHEMA_VERSION = 1


@dataclass
class QualityRunRecord:
    """Persisted record for a quality suite run."""

    id: str
    agent_name: str
    kind: str
    status: str
    created_at: str
    started_at: str | None = None
    finished_at: str | None = None
    summary: dict[str, Any] | None = None
    rows: list[dict[str, Any]] | None = None
    reports: dict[str, str] | None = None
    suite: dict[str, Any] | None = None
    error: str | None = None


class QualityHistoryStore:
    """Abstract persistent store for quality run history."""

    def create_run(self, record: QualityRunRecord) -> None:
        raise NotImplementedError

    def update_run(self, run_id: str, **fields: Any) -> None:
        raise NotImplementedError

    def get_run(self, run_id: str) -> QualityRunRecord | None:
        raise NotImplementedError

    def list_runs(
        self,
        *,
        agent_name: str | None = None,
        kind: str | None = None,
        limit: int = 50,
    ) -> list[QualityRunRecord]:
        raise NotImplementedError

    def compare_runs(self, run_id_a: str, run_id_b: str) -> dict[str, Any]:
        raise NotImplementedError


class SQLiteQualityHistoryStore(QualityHistoryStore):
    """SQLite-backed quality run history store."""

    def __init__(self, path: str = ".logos_adk/quality_runs.db", table_name: str = "quality_runs") -> None:
        self.path = path
        self.table_name = table_name

    def _connect(self) -> sqlite3.Connection:
        db_path = Path(self.path)
        if db_path.parent and not db_path.parent.exists():
            db_path.parent.mkdir(parents=True, exist_ok=True)
        connection = sqlite3.connect(self.path)
        connection.row_factory = sqlite3.Row
        return connection

    def _ensure_schema(self, connection: sqlite3.Connection) -> None:
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                kind TEXT NOT NULL,
                status TEXT NOT NULL,
                created_at TEXT NOT NULL,
                started_at TEXT,
                finished_at TEXT,
                summary_json TEXT,
                rows_json TEXT,
                reports_json TEXT,
                suite_json TEXT,
                error TEXT
            )
            """
        )
        connection.execute(
            f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_agent_created "
            f"ON {self.table_name} (agent_name, created_at DESC)"
        )
        recorded_version = get_sqlite_schema_version(connection, QUALITY_HISTORY_SCHEMA_KIND)
        if recorded_version != QUALITY_HISTORY_SCHEMA_VERSION:
            set_sqlite_schema_version(connection, QUALITY_HISTORY_SCHEMA_KIND, QUALITY_HISTORY_SCHEMA_VERSION)
        connection.commit()

    def create_run(self, record: QualityRunRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                f"""
                INSERT INTO {self.table_name}
                (id, agent_name, kind, status, created_at, started_at, finished_at,
                 summary_json, rows_json, reports_json, suite_json, error)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    record.id,
                    record.agent_name,
                    record.kind,
                    record.status,
                    record.created_at,
                    record.started_at,
                    record.finished_at,
                    json.dumps(record.summary) if record.summary is not None else None,
                    json.dumps(record.rows) if record.rows is not None else None,
                    json.dumps(record.reports) if record.reports is not None else None,
                    json.dumps(record.suite) if record.suite is not None else None,
                    record.error,
                ),
            )
            connection.commit()

    def update_run(self, run_id: str, **fields: Any) -> None:
        if not fields:
            return
        assignments = []
        values = []
        for key, value in fields.items():
            column = f"{key}_json" if key in {"summary", "rows", "reports", "suite"} else key
            assignments.append(f"{column} = ?")
            if key in {"summary", "rows", "reports", "suite"} and value is not None:
                values.append(json.dumps(value))
            else:
                values.append(value)
        values.append(run_id)
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                f"UPDATE {self.table_name} SET {', '.join(assignments)} WHERE id = ?",
                values,
            )
            connection.commit()

    def get_run(self, run_id: str) -> QualityRunRecord | None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            row = connection.execute(
                f"SELECT * FROM {self.table_name} WHERE id = ?",
                (run_id,),
            ).fetchone()
        return self._row_to_record(row) if row else None

    def list_runs(
        self,
        *,
        agent_name: str | None = None,
        kind: str | None = None,
        limit: int = 50,
    ) -> list[QualityRunRecord]:
        clauses = []
        params: list[Any] = []
        if agent_name is not None:
            clauses.append("agent_name = ?")
            params.append(agent_name)
        if kind is not None:
            clauses.append("kind = ?")
            params.append(kind)
        where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        params.append(limit)
        with self._connect() as connection:
            self._ensure_schema(connection)
            rows = connection.execute(
                f"""
                SELECT * FROM {self.table_name}
                {where_sql}
                ORDER BY created_at DESC
                LIMIT ?
                """,
                params,
            ).fetchall()
        return [self._row_to_record(row) for row in rows]

    def compare_runs(self, run_id_a: str, run_id_b: str) -> dict[str, Any]:
        run_a = self.get_run(run_id_a)
        run_b = self.get_run(run_id_b)
        if run_a is None or run_b is None:
            raise StateBackendError("Both quality runs must exist for comparison")
        keys = sorted(set((run_a.summary or {}).keys()) | set((run_b.summary or {}).keys()))
        summary_diff = {
            key: {
                "a": (run_a.summary or {}).get(key),
                "b": (run_b.summary or {}).get(key),
            }
            for key in keys
            if (run_a.summary or {}).get(key) != (run_b.summary or {}).get(key)
        }
        return {
            "run_a": asdict(run_a),
            "run_b": asdict(run_b),
            "summary_diff": summary_diff,
            "row_count_diff": {
                "a": len(run_a.rows or []),
                "b": len(run_b.rows or []),
            },
        }

    def _row_to_record(self, row: sqlite3.Row) -> QualityRunRecord:
        return QualityRunRecord(
            id=row["id"],
            agent_name=row["agent_name"],
            kind=row["kind"],
            status=row["status"],
            created_at=row["created_at"],
            started_at=row["started_at"],
            finished_at=row["finished_at"],
            summary=json.loads(row["summary_json"]) if row["summary_json"] else None,
            rows=json.loads(row["rows_json"]) if row["rows_json"] else None,
            reports=json.loads(row["reports_json"]) if row["reports_json"] else None,
            suite=json.loads(row["suite_json"]) if row["suite_json"] else None,
            error=row["error"],
        )


class PostgreSQLQualityHistoryStore(QualityHistoryStore):
    """PostgreSQL-backed quality run history store."""

    def __init__(self, dsn: str, table_name: str = "quality_runs") -> None:
        self.dsn = dsn
        self.table_name = table_name

    def _connect(self) -> Any:
        psycopg = import_psycopg()
        try:
            return psycopg.connect(self.dsn)
        except Exception as exc:
            raise StateBackendError(f"Failed to open PostgreSQL quality history store: {exc}") from exc

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    agent_name TEXT NOT NULL,
                    kind TEXT NOT NULL,
                    status TEXT NOT NULL,
                    created_at TIMESTAMPTZ NOT NULL,
                    started_at TIMESTAMPTZ,
                    finished_at TIMESTAMPTZ,
                    summary_json JSONB,
                    rows_json JSONB,
                    reports_json JSONB,
                    suite_json JSONB,
                    error TEXT
                )
                """
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_agent_created "
                f"ON {self.table_name} (agent_name, created_at DESC)"
            )
        recorded_version = get_postgresql_schema_version(connection, QUALITY_HISTORY_SCHEMA_KIND)
        if recorded_version != QUALITY_HISTORY_SCHEMA_VERSION:
            set_postgresql_schema_version(connection, QUALITY_HISTORY_SCHEMA_KIND, QUALITY_HISTORY_SCHEMA_VERSION)
        connection.commit()

    def create_run(self, record: QualityRunRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self.table_name}
                    (id, agent_name, kind, status, created_at, started_at, finished_at,
                     summary_json, rows_json, reports_json, suite_json, error)
                    VALUES (%s, %s, %s, %s, %s, %s, %s,
                            %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb, %s)
                    """,
                    (
                        record.id,
                        record.agent_name,
                        record.kind,
                        record.status,
                        record.created_at,
                        record.started_at,
                        record.finished_at,
                        json.dumps(record.summary) if record.summary is not None else None,
                        json.dumps(record.rows) if record.rows is not None else None,
                        json.dumps(record.reports) if record.reports is not None else None,
                        json.dumps(record.suite) if record.suite is not None else None,
                        record.error,
                    ),
                )
            connection.commit()

    def update_run(self, run_id: str, **fields: Any) -> None:
        if not fields:
            return
        assignments = []
        values: list[Any] = []
        for key, value in fields.items():
            if key in {"summary", "rows", "reports", "suite"}:
                assignments.append(f"{key}_json = %s::jsonb")
                values.append(json.dumps(value) if value is not None else None)
            else:
                assignments.append(f"{key} = %s")
                values.append(value)
        values.append(run_id)
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"UPDATE {self.table_name} SET {', '.join(assignments)} WHERE id = %s",
                    values,
                )
            connection.commit()

    def get_run(self, run_id: str) -> QualityRunRecord | None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(f"SELECT * FROM {self.table_name} WHERE id = %s", (run_id,))
                row = cursor.fetchone()
        return self._row_to_record(row) if row else None

    def list_runs(
        self,
        *,
        agent_name: str | None = None,
        kind: str | None = None,
        limit: int = 50,
    ) -> list[QualityRunRecord]:
        clauses = []
        params: list[Any] = []
        if agent_name is not None:
            clauses.append("agent_name = %s")
            params.append(agent_name)
        if kind is not None:
            clauses.append("kind = %s")
            params.append(kind)
        where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        params.append(limit)
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    SELECT * FROM {self.table_name}
                    {where_sql}
                    ORDER BY created_at DESC
                    LIMIT %s
                    """,
                    params,
                )
                rows = cursor.fetchall()
        return [self._row_to_record(row) for row in rows]

    def compare_runs(self, run_id_a: str, run_id_b: str) -> dict[str, Any]:
        run_a = self.get_run(run_id_a)
        run_b = self.get_run(run_id_b)
        if run_a is None or run_b is None:
            raise StateBackendError("Both quality runs must exist for comparison")
        keys = sorted(set((run_a.summary or {}).keys()) | set((run_b.summary or {}).keys()))
        summary_diff = {
            key: {
                "a": (run_a.summary or {}).get(key),
                "b": (run_b.summary or {}).get(key),
            }
            for key in keys
            if (run_a.summary or {}).get(key) != (run_b.summary or {}).get(key)
        }
        return {
            "run_a": asdict(run_a),
            "run_b": asdict(run_b),
            "summary_diff": summary_diff,
            "row_count_diff": {"a": len(run_a.rows or []), "b": len(run_b.rows or [])},
        }

    def _row_to_record(self, row: Any) -> QualityRunRecord:
        return QualityRunRecord(
            id=row[0],
            agent_name=row[1],
            kind=row[2],
            status=row[3],
            created_at=row[4].isoformat() if hasattr(row[4], "isoformat") else str(row[4]),
            started_at=row[5].isoformat() if row[5] is not None and hasattr(row[5], "isoformat") else row[5],
            finished_at=row[6].isoformat() if row[6] is not None and hasattr(row[6], "isoformat") else row[6],
            summary=row[7],
            rows=row[8],
            reports=row[9],
            suite=row[10],
            error=row[11],
        )


def build_quality_history_store(
    config: TaskQueueConfig | None = None,
) -> QualityHistoryStore:
    """Build a persistent quality history store."""
    database_url = resolve_database_url(
        config.dsn if config is not None else None,
        env_vars=("LOGOS_ADK_QUALITY_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"),
    )
    if config is not None and config.backend in {"postgres", "postgresql"}:
        if not database_url:
            raise ConfigError("PostgreSQL quality history requires a DSN or DATABASE_URL")
        return PostgreSQLQualityHistoryStore(database_url)
    if database_url:
        return PostgreSQLQualityHistoryStore(database_url)
    path = config.path if config is not None else ".logos_adk/quality_runs.db"
    return SQLiteQualityHistoryStore(path=path)


class QualityExecutionManager:
    """Background execution manager for quality suites."""

    def __init__(self, store: QualityHistoryStore, *, concurrency: int = 2) -> None:
        self.store = store
        self._semaphore = asyncio.Semaphore(max(1, concurrency))
        self._tasks: dict[str, asyncio.Task[None]] = {}

    def submit(
        self,
        *,
        agent_name: str,
        kind: str,
        suite: dict[str, Any],
        runner: Callable[[], Awaitable[dict[str, Any]]],
    ) -> QualityRunRecord:
        run_id = uuid4().hex
        record = QualityRunRecord(
            id=run_id,
            agent_name=agent_name,
            kind=kind,
            status="queued",
            created_at=datetime.now(UTC).isoformat(),
            suite=suite,
        )
        self.store.create_run(record)
        self._tasks[run_id] = asyncio.create_task(self._run(record, runner))
        return record

    async def _run(
        self,
        record: QualityRunRecord,
        runner: Callable[[], Awaitable[dict[str, Any]]],
    ) -> None:
        async with self._semaphore:
            started_at = datetime.now(UTC).isoformat()
            self.store.update_run(record.id, status="running", started_at=started_at)
            try:
                payload = await runner()
            except Exception as exc:
                self.store.update_run(
                    record.id,
                    status="failed",
                    started_at=started_at,
                    finished_at=datetime.now(UTC).isoformat(),
                    error=str(exc),
                )
                return
            self.store.update_run(
                record.id,
                status="completed",
                started_at=started_at,
                finished_at=datetime.now(UTC).isoformat(),
                summary=payload.get("summary"),
                rows=payload.get("rows"),
                reports=payload.get("reports"),
            )

    def get(self, run_id: str) -> QualityRunRecord | None:
        return self.store.get_run(run_id)

    def list_runs(self, *, agent_name: str | None = None, kind: str | None = None, limit: int = 50) -> list[QualityRunRecord]:
        return self.store.list_runs(agent_name=agent_name, kind=kind, limit=limit)

    def compare(self, run_id_a: str, run_id_b: str) -> dict[str, Any]:
        return self.store.compare_runs(run_id_a, run_id_b)

    def record_completed(
        self,
        *,
        agent_name: str,
        kind: str,
        suite: dict[str, Any],
        payload: dict[str, Any],
    ) -> QualityRunRecord:
        """Persist a synchronously completed quality run."""
        timestamp = datetime.now(UTC).isoformat()
        record = QualityRunRecord(
            id=uuid4().hex,
            agent_name=agent_name,
            kind=kind,
            status="completed",
            created_at=timestamp,
            started_at=timestamp,
            finished_at=timestamp,
            summary=payload.get("summary"),
            rows=payload.get("rows"),
            reports=payload.get("reports"),
            suite=suite,
        )
        self.store.create_run(record)
        return record


_quality_execution_manager: QualityExecutionManager | None = None


def get_quality_execution_manager() -> QualityExecutionManager:
    """Get or initialize the process-wide quality execution manager."""
    global _quality_execution_manager
    if _quality_execution_manager is None:
        _quality_execution_manager = QualityExecutionManager(build_quality_history_store())
    return _quality_execution_manager


def reset_quality_execution_manager() -> QualityExecutionManager | None:
    """Reset quality execution manager. For tests only."""
    global _quality_execution_manager
    manager = _quality_execution_manager
    _quality_execution_manager = None
    return manager


__all__ = [
    "get_quality_execution_manager",
    "PostgreSQLQualityHistoryStore",
    "QualityExecutionManager",
    "QualityHistoryStore",
    "QualityRunRecord",
    "SQLiteQualityHistoryStore",
    "build_quality_history_store",
    "reset_quality_execution_manager",
]
