from __future__ import annotations

import hashlib
import json
import sqlite3
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from sqlalchemy import create_engine, select
from sqlalchemy.dialects.postgresql import insert as postgres_insert

from .schema import build_session_tables


TABLES = ("sessions", "turns", "blocks", "resumes", "resume_imports", "optimization_runs")


@dataclass(frozen=True, slots=True)
class MigrationReport:
    source_counts: dict[str, int]
    target_counts: dict[str, int]
    source_checksum: str
    target_checksum: str


def migrate_sqlite_to_postgres(
    source_path: str | Path,
    target_url: str,
    *,
    schema: str = "resume_agent",
    dry_run: bool = False,
    conflict_policy: str = "error",
) -> MigrationReport:
    """Copy SQLite rows into PostgreSQL without deleting target-only records."""
    if conflict_policy not in {"error", "update"}:
        raise ValueError("conflict_policy must be 'error' or 'update'")
    source_rows = _read_source(Path(source_path))
    source_counts = {name: len(rows) for name, rows in source_rows.items()}
    source_checksum = _checksum(source_rows)
    if dry_run:
        return MigrationReport(source_counts, {}, source_checksum, "")

    engine = create_engine(target_url)
    metadata, tables = build_session_tables(schema)
    try:
        with engine.begin() as connection:
            connection.exec_driver_sql(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')
            metadata.create_all(connection)
            for name in TABLES:
                rows = source_rows[name]
                if not rows:
                    continue
                table = tables[name]
                statement = postgres_insert(table).values(rows)
                if conflict_policy == "update":
                    statement = statement.on_conflict_do_update(
                        index_elements=[table.c.id],
                        set_={column.name: statement.excluded[column.name]
                              for column in table.columns if column.name != "id"},
                    )
                connection.execute(statement)
            target_rows = _read_target(connection, tables, source_rows)
            target_checksum = _checksum(target_rows)
            if target_checksum != source_checksum:
                raise RuntimeError("PostgreSQL verification checksum does not match SQLite source")
    finally:
        engine.dispose()
    target_counts = {name: len(rows) for name, rows in target_rows.items()}
    target_checksum = _checksum(target_rows)
    return MigrationReport(source_counts, target_counts, source_checksum, target_checksum)


def _read_source(path: Path) -> dict[str, list[dict[str, Any]]]:
    connection = sqlite3.connect(f"file:{path.resolve().as_posix()}?mode=ro", uri=True)
    connection.row_factory = sqlite3.Row
    try:
        rows: dict[str, list[dict[str, Any]]] = {}
        existing = {row[0] for row in connection.execute(
            "SELECT name FROM sqlite_master WHERE type = 'table'"
        )}
        for name in TABLES:
            if name not in existing:
                rows[name] = []
                continue
            result = [dict(row) for row in connection.execute(f"SELECT * FROM {name}")]
            rows[name] = [_normalize_source(name, row) for row in result]
        return rows
    finally:
        connection.close()


def _normalize_source(name: str, row: dict[str, Any]) -> dict[str, Any]:
    result = dict(row)
    if name == "sessions":
        result["profile"] = json.loads(result.pop("profile_json"))
        if result["profile"].get("job_type") == "other":
            result["profile"]["job_type"] = "internship"
    elif name == "blocks":
        result["data"] = json.loads(result.pop("data_json"))
    elif name == "resumes":
        result["content"] = json.loads(result.pop("content_json"))
    elif name == "resume_imports":
        raw_document = result.pop("document_json")
        result["document"] = json.loads(raw_document) if raw_document else None
        result["field_reviews"] = json.loads(result.pop("field_reviews_json"))
    elif name == "optimization_runs":
        result["state"] = json.loads(result.pop("state_json"))
        raw_proposal = result.pop("proposal_json")
        result["proposal"] = json.loads(raw_proposal) if raw_proposal else None

    return result

def _read_target(
    connection: Any, tables: dict[str, Any], source_rows: dict[str, list[dict[str, Any]]]
) -> dict[str, list[dict[str, Any]]]:
    result: dict[str, list[dict[str, Any]]] = {}
    for name in TABLES:
        table = tables[name]
        source_ids = [row["id"] for row in source_rows[name]]
        if not source_ids:
            result[name] = []
            continue
        rows = connection.execute(select(table).where(
            table.c.id.in_(source_ids)
        ).order_by(table.c.id)).mappings().all()
        result[name] = [dict(row) for row in rows]
    return result


def _checksum(rows_by_table: dict[str, list[dict[str, Any]]]) -> str:
    stable: dict[str, list[dict[str, Any]]] = {}
    for name in TABLES:
        rows = rows_by_table.get(name, [])
        stable[name] = [_checksum_row(name, row) for row in rows]
        stable[name].sort(key=lambda row: str(row["id"]))
    serialized = json.dumps(stable, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(serialized.encode("utf-8")).hexdigest()


def _checksum_row(name: str, row: dict[str, Any]) -> dict[str, Any]:
    keys = {
        "sessions": ("id", "stage", "revision", "profile", "draft_id", "resume_id"),
        "turns": ("id", "session_id", "sequence", "role", "content", "composer_mode"),
        "blocks": ("id", "session_id", "turn_id", "block_index", "type", "lifecycle", "data", "version"),
        "resumes": ("id", "session_id", "idempotency_key", "revision", "content"),
        "resume_imports": (
            "id", "session_id", "file_name", "mime_type", "size_bytes", "sha256",
            "object_key", "status", "document", "field_reviews", "error_code",
        ),
        "optimization_runs": (
            "id", "session_id", "entry_id", "mode", "status", "source_revision",
            "state", "proposal",
        ),
    }[name]
    return {key: row[key] for key in keys}