from __future__ import annotations

import argparse
import hashlib
import json
import sqlite3
from pathlib import Path
from typing import Any


def digest(value: Any) -> str:
    content = json.dumps(
        value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str
    ).encode("utf-8")
    return hashlib.sha256(content).hexdigest()


def schema(connection: sqlite3.Connection) -> list[dict[str, Any]]:
    rows = connection.execute(
        """
        SELECT type, name, tbl_name, sql
        FROM sqlite_master
        WHERE name NOT LIKE 'sqlite_%'
        ORDER BY type, name
        """
    ).fetchall()
    return [dict(row) for row in rows]


def table_rows(
    connection: sqlite3.Connection,
    table: str,
    excluded_columns: set[str] | None = None,
) -> list[dict[str, Any]]:
    quoted = '"' + table.replace('"', '""') + '"'
    excluded_columns = excluded_columns or set()
    rows = [
        {
            key: value
            for key, value in dict(row).items()
            if key not in excluded_columns
        }
        for row in connection.execute(f"SELECT * FROM {quoted}").fetchall()
    ]
    return sorted(rows, key=lambda row: json.dumps(row, ensure_ascii=False, sort_keys=True, default=str))


def main() -> None:
    parser = argparse.ArgumentParser(description="Compare preservation SQLite databases")
    parser.add_argument("--original", type=Path, required=True)
    parser.add_argument("--migrated", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--exclude-column",
        action="append",
        default=[],
        metavar="TABLE.COLUMN",
        help="exclude a nondeterministic column from one table comparison",
    )
    parser.add_argument("tables", nargs="+")
    args = parser.parse_args()

    excluded_by_table: dict[str, set[str]] = {}
    for item in args.exclude_column:
        table, separator, column = item.partition(".")
        if not separator or not table or not column:
            parser.error("--exclude-column must use TABLE.COLUMN")
        excluded_by_table.setdefault(table, set()).add(column)

    original = sqlite3.connect(args.original)
    migrated = sqlite3.connect(args.migrated)
    original.row_factory = sqlite3.Row
    migrated.row_factory = sqlite3.Row
    try:
        original_schema = schema(original)
        migrated_schema = schema(migrated)
        tables = []
        all_equal = original_schema == migrated_schema
        for table in args.tables:
            excluded_columns = excluded_by_table.get(table, set())
            original_rows = table_rows(original, table, excluded_columns)
            migrated_rows = table_rows(migrated, table, excluded_columns)
            equal = original_rows == migrated_rows
            all_equal = all_equal and equal
            tables.append(
                {
                    "table": table,
                    "original_count": len(original_rows),
                    "migrated_count": len(migrated_rows),
                    "original_sha256": digest(original_rows),
                    "migrated_sha256": digest(migrated_rows),
                    "equal": equal,
                    "excluded_columns": sorted(excluded_columns),
                }
            )
        result = {
            "all_equal": all_equal,
            "schema": {
                "object_count": len(original_schema),
                "original_sha256": digest(original_schema),
                "migrated_sha256": digest(migrated_schema),
                "equal": original_schema == migrated_schema,
            },
            "tables": tables,
        }
    finally:
        original.close()
        migrated.close()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if not all_equal:
        raise SystemExit(1)


if __name__ == "__main__":
    main()
