from __future__ import annotations

import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

from baudolo.restore.run import docker_exec

from .version import guard

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator

_SUPERUSER_ONLY_PREFIXES = (b"COMMENT ON EXTENSION", b"ALTER DEFAULT PRIVILEGES")
_EMPTY_PRECLEAN_SQL = Path(__file__).parent / "empty_preclean.sql"


def filter_superuser_only_lines(lines: Iterable[bytes]) -> Iterator[bytes]:
    """Drop superuser-only statements an app-level psql replay cannot run.

    Args:
        lines: dump lines including their trailing newlines.

    Yields:
        Every line except top-level statements starting with a superuser-only
        prefix. Lines inside COPY ... FROM stdin data blocks are passed
        through untouched: a data row may legally start with the same bytes,
        and dropping it would silently corrupt the restored table.
    """
    in_copy = False
    for line in lines:
        if in_copy:
            yield line
            if line.rstrip(b"\r\n") == b"\\.":
                in_copy = False
            continue
        if line.startswith(b"COPY ") and line.rstrip(b"\r\n").endswith(b"FROM stdin;"):
            in_copy = True
            yield line
            continue
        if line.startswith(_SUPERUSER_ONLY_PREFIXES):
            continue
        yield line


def restore_postgres_sql(
    *,
    container: str,
    db_name: str,
    user: str,
    password: str,
    sql_path: str,
    empty: bool,
    check_version: bool = True,
) -> None:
    if not Path(sql_path).is_file():
        raise FileNotFoundError(sql_path)

    if check_version:
        guard(
            sql_path=sql_path,
            engine="postgres",
            container=container,
            user=user,
            password=password,
        )

    docker_env = {"PGPASSWORD": password}

    if empty:
        with _EMPTY_PRECLEAN_SQL.open(encoding="utf-8") as preclean:
            drop_sql = preclean.read()
        docker_exec(
            container,
            ["psql", "-v", "ON_ERROR_STOP=1", "-U", user, "-d", db_name],
            stdin=drop_sql.encode(),
            docker_env=docker_env,
        )

    # Filter into a spooled temp file instead of building the whole dump in
    # memory: production dumps reach many GB and the previous read/splitlines/
    # join needed roughly three times the dump size in RSS.
    with Path(sql_path).open("rb") as src, tempfile.TemporaryFile() as filtered:
        for line in filter_superuser_only_lines(src):
            filtered.write(line)
        filtered.seek(0)
        docker_exec(
            container,
            [
                "psql",
                "--single-transaction",
                "-v",
                "ON_ERROR_STOP=1",
                "-U",
                user,
                "-d",
                db_name,
            ],
            stdin=filtered,
            docker_env=docker_env,
        )

    print(f"PostgreSQL restore complete for db '{db_name}'.")