#!/usr/bin/env python3
"""Generate a PostgreSQL restore script from a Teleo SQLite backup.

This is a shadow-restore converter for restore drills while the canonical
runtime database is still SQLite. It preserves tables and rows, infers safe
Postgres column types from actual SQLite storage classes, and avoids blind
string rewrites of `.dump` output.
"""

from __future__ import annotations

import argparse
import gzip
import json
import shutil
import sqlite3
import tempfile
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

PG_NUMERIC_TYPES = {"bigint", "double precision"}


def pg_quote_identifier(value: str) -> str:
    return '"' + value.replace('"', '""') + '"'


def sqlite_quote_identifier(value: str) -> str:
    return '"' + value.replace('"', '""') + '"'


def copy_escape(value: Any, pg_type: str, stats: dict[str, int]) -> str:
    if value is None:
        return r"\N"
    if isinstance(value, bytes):
        text = value.hex()
        stats["blob_values_hex_encoded"] += 1
    elif pg_type in PG_NUMERIC_TYPES:
        text = str(value)
    else:
        text = str(value)

    if "\x00" in text:
        stats["nul_chars_replaced"] += text.count("\x00")
        text = text.replace("\x00", r"\u0000")

    return (
        text.replace("\\", "\\\\")
        .replace("\t", r"\t")
        .replace("\n", r"\n")
        .replace("\r", r"\r")
    )


def infer_pg_type(declared_type: str | None, storage_counts: dict[str, int]) -> tuple[str, str]:
    storage_types = {name for name, count in storage_counts.items() if count and name != "null"}
    declared = (declared_type or "").upper()

    if not storage_types:
        if "INT" in declared:
            return "bigint", "all_null_declared_integer"
        if any(token in declared for token in ("REAL", "FLOA", "DOUB")):
            return "double precision", "all_null_declared_real"
        return "text", "all_null_text_default"

    if storage_types == {"integer"}:
        return "bigint", "sqlite_integer_storage"
    if storage_types <= {"integer", "real"}:
        return "double precision", "sqlite_numeric_storage"
    return "text", "sqlite_text_or_mixed_storage"


def read_storage_counts(conn: sqlite3.Connection, table: str, column: str) -> dict[str, int]:
    sql = (
        f"select typeof({sqlite_quote_identifier(column)}) as storage_type, count(*) "
        f"from {sqlite_quote_identifier(table)} group by storage_type"
    )
    return {str(row[0]): int(row[1]) for row in conn.execute(sql)}


def user_tables(conn: sqlite3.Connection) -> list[dict[str, str]]:
    rows = conn.execute(
        """
        select name, sql
        from sqlite_master
        where type = 'table'
          and name not like 'sqlite_%'
        order by name
        """
    ).fetchall()
    return [{"name": str(row[0]), "sqlite_sql": row[1] or ""} for row in rows]


@contextmanager
def sqlite_database_path(path: Path):
    if path.suffix != ".gz":
        yield path
        return

    tmpdir = Path(tempfile.mkdtemp(prefix="teleo-sqlite-restore-"))
    try:
        db_path = tmpdir / "pipeline.db"
        with gzip.open(path, "rb") as source, db_path.open("wb") as target:
            shutil.copyfileobj(source, target)
        yield db_path
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def build_table_metadata(conn: sqlite3.Connection, table: str) -> dict[str, Any]:
    columns: list[dict[str, Any]] = []
    for cid, name, declared_type, notnull, default_value, primary_key in conn.execute(
        f"pragma table_info({sqlite_quote_identifier(table)})"
    ):
        storage_counts = read_storage_counts(conn, table, name)
        pg_type, type_reason = infer_pg_type(declared_type, storage_counts)
        columns.append(
            {
                "cid": cid,
                "name": name,
                "sqlite_declared_type": declared_type or "",
                "sqlite_notnull": bool(notnull),
                "sqlite_default": default_value,
                "sqlite_primary_key_position": primary_key,
                "postgres_type": pg_type,
                "postgres_type_reason": type_reason,
                "sqlite_storage_counts": storage_counts,
            }
        )
    row_count = conn.execute(f"select count(*) from {sqlite_quote_identifier(table)}").fetchone()[0]
    return {"name": table, "row_count": int(row_count), "columns": columns}


def write_target_counts_sql(path: Path, schema: str, tables: list[dict[str, Any]]) -> None:
    selects = []
    for table in tables:
        table_name = table["name"]
        table_literal = table_name.replace("'", "''")
        table_ref = f"{pg_quote_identifier(schema)}.{pg_quote_identifier(table_name)}"
        selects.append(f"select '{table_literal}' as table_name, count(*)::bigint as row_count from {table_ref}")
    body = "\nunion all\n".join(selects) or "select '' as table_name, 0::bigint as row_count where false"
    path.write_text(
        "copy (\n"
        f"{body}\n"
        "order by table_name\n"
        ") to stdout with (format csv, header true);\n"
    )


def convert_sqlite_to_postgres(
    sqlite_path: Path,
    output_sql: Path,
    summary_json: Path,
    schema: str,
    target_counts_sql: Path | None = None,
) -> dict[str, Any]:
    output_sql.parent.mkdir(parents=True, exist_ok=True)
    summary_json.parent.mkdir(parents=True, exist_ok=True)
    if target_counts_sql:
        target_counts_sql.parent.mkdir(parents=True, exist_ok=True)

    with sqlite_database_path(sqlite_path) as db_path:
        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        conn.row_factory = sqlite3.Row
        integrity = conn.execute("pragma integrity_check").fetchone()[0]
        tables = []
        table_names = [table["name"] for table in user_tables(conn)]
        for table_name in table_names:
            tables.append(build_table_metadata(conn, table_name))

        stats = {"blob_values_hex_encoded": 0, "nul_chars_replaced": 0}
        with output_sql.open("w", encoding="utf-8", newline="\n") as out:
            out.write("\\set ON_ERROR_STOP on\n")
            out.write("begin;\n")
            out.write(f"drop schema if exists {pg_quote_identifier(schema)} cascade;\n")
            out.write(f"create schema {pg_quote_identifier(schema)};\n")
            out.write(f"set search_path to {pg_quote_identifier(schema)};\n")

            for table in tables:
                table_ref = f"{pg_quote_identifier(schema)}.{pg_quote_identifier(table['name'])}"
                column_defs = [
                    f"{pg_quote_identifier(column['name'])} {column['postgres_type']}"
                    for column in table["columns"]
                ]
                out.write(f"create table {table_ref} (\n  " + ",\n  ".join(column_defs) + "\n);\n")

                if not table["columns"]:
                    continue

                column_list = ", ".join(pg_quote_identifier(column["name"]) for column in table["columns"])
                out.write(f"copy {table_ref} ({column_list}) from stdin with (format text, delimiter E'\\t', null '\\N');\n")

                select_columns = ", ".join(sqlite_quote_identifier(column["name"]) for column in table["columns"])
                for row in conn.execute(f"select {select_columns} from {sqlite_quote_identifier(table['name'])}"):
                    fields = [
                        copy_escape(row[column["name"]], column["postgres_type"], stats)
                        for column in table["columns"]
                    ]
                    out.write("\t".join(fields))
                    out.write("\n")
                out.write("\\.\n")

            out.write("commit;\n")
            out.write("analyze;\n")

        if target_counts_sql:
            write_target_counts_sql(target_counts_sql, schema, tables)

        summary = {
            "artifact": "teleo_sqlite_to_postgres_dump",
            "generated_at_utc": datetime.now(UTC).isoformat(),
            "source_sqlite_path": str(sqlite_path),
            "source_integrity_check": integrity,
            "schema": schema,
            "table_count": len(tables),
            "source_total_rows": sum(table["row_count"] for table in tables),
            "tables": tables,
            "conversion_notes": {
                "mode": "shadow_restore",
                "constraints_and_indexes": "not_recreated",
                "blob_storage": "hex_encoded_text_when_needed",
                "postgres_types": "inferred_from_sqlite_storage_classes",
            },
            "conversion_stats": stats,
            "output_sql": str(output_sql),
            "target_counts_sql": str(target_counts_sql) if target_counts_sql else None,
        }
        summary_json.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
        return summary


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--sqlite-backup", required=True, type=Path, help="SQLite .db or .db.gz backup")
    parser.add_argument("--output-sql", required=True, type=Path, help="Postgres import SQL path")
    parser.add_argument("--summary-json", required=True, type=Path, help="Non-secret conversion summary JSON")
    parser.add_argument("--schema", default="teleo_restore", help="Postgres schema to recreate")
    parser.add_argument("--target-counts-sql", type=Path, help="Optional psql query file for target row counts")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    summary = convert_sqlite_to_postgres(
        sqlite_path=args.sqlite_backup,
        output_sql=args.output_sql,
        summary_json=args.summary_json,
        schema=args.schema,
        target_counts_sql=args.target_counts_sql,
    )
    print(
        json.dumps(
            {
                "artifact": summary["artifact"],
                "source_integrity_check": summary["source_integrity_check"],
                "schema": summary["schema"],
                "table_count": summary["table_count"],
                "source_total_rows": summary["source_total_rows"],
                "output_sql": summary["output_sql"],
                "target_counts_sql": summary["target_counts_sql"],
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0 if summary["source_integrity_check"] == "ok" else 1


if __name__ == "__main__":
    raise SystemExit(main())
