"""Export one private inspection batch while replacing VINs with keyed HMACs.

The output is private analytical staging data, not a public-dashboard dataset.
It intentionally omits plates, ZIPs, stations, raw JSON, and operational data.
"""

from __future__ import annotations

import argparse
import csv
import gzip
import hashlib
import hmac
import json
import os
import re
import stat
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, Optional, Sequence

import psycopg2


QUERY_VERSION = "inspection_batch_v4"
MAX_BATCH_DAYS = 32
SECURE_SSL_MODES = {"require", "verify-ca", "verify-full"}
PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_KEY_FILE = PROJECT_ROOT / ".secrets/vin_hmac.key"
PRIVATE_DATA_ROOT = (PROJECT_ROOT / "data/private").resolve()
VIN_PATTERN = re.compile(r"^[A-HJ-NPR-Z0-9]{17}$")
KEY_VERSION_PATTERN = re.compile(r"^[A-Za-z0-9._-]{1,32}$")

OUTCOME_ALIASES = {
    "PASS": "pass",
    "P": "pass",
    "FAIL": "fail",
    "F": "fail",
    "REJECT": "reject",
    "ABORT": "abort",
}


def _normalized_outcome(value: object) -> Optional[str]:
    if not isinstance(value, str):
        return None
    return OUTCOME_ALIASES.get(value.strip().upper())


def canonicalize_outcome(
    overall_result: object,
    obd_result: object,
    source_era: object,
    program_type: object = None,
    test_type: object = None,
) -> tuple[Optional[str], Optional[str]]:
    """Return the approved label and its provenance for one inspection."""
    overall = _normalized_outcome(overall_result)
    if overall is not None:
        return overall, "overall_result"

    proxy_eligible = (
        isinstance(source_era, str)
        and source_era.strip().lower() == "utah"
        and isinstance(program_type, str)
        and program_type.strip().lower() == "obd"
        and isinstance(test_type, str)
        and test_type.strip().upper() == "OBD"
    )
    if proxy_eligible:
        proxy = _normalized_outcome(obd_result)
        if proxy is not None:
            # This narrowly validated Utah OBD/OBD proxy is approved only for
            # binary pass/non-pass labels. Multiclass work must retain
            # overall_result labels exclusively.
            return proxy, "utah_obd_proxy"
    return None, None


def _outcome_case_sql(column: str) -> str:
    cases = " ".join(
        f"WHEN '{raw}' THEN '{canonical}'"
        for raw, canonical in OUTCOME_ALIASES.items()
    )
    return f"CASE upper(btrim({column})) {cases} ELSE NULL END"


_OVERALL_OUTCOME_SQL = _outcome_case_sql("s.overall_result")
_UTAH_OBD_OUTCOME_SQL = _outcome_case_sql("s.obd_result")
_UTAH_OBD_PROXY_PREDICATE_SQL = """lower(btrim(s.county)) = 'utah'
         AND lower(btrim(s.program_type)) = 'obd'
         AND upper(btrim(s.test_type)) = 'OBD'"""

# Keep label provenance in private staging. Only Utah rows normalized to both
# program_type=obd and test_type=OBD may use the proxy, and then only for binary
# pass/non-pass modeling. Multiclass labels must use overall_result.
OUTCOME_SELECT_SQL = f"""CASE
        WHEN {_OVERALL_OUTCOME_SQL} IS NOT NULL THEN {_OVERALL_OUTCOME_SQL}
        WHEN {_UTAH_OBD_PROXY_PREDICATE_SQL} THEN {_UTAH_OBD_OUTCOME_SQL}
        ELSE NULL
    END AS canonical_outcome,
    CASE
        WHEN {_OVERALL_OUTCOME_SQL} IS NOT NULL THEN 'overall_result'
        WHEN {_UTAH_OBD_PROXY_PREDICATE_SQL}
         AND {_UTAH_OBD_OUTCOME_SQL} IS NOT NULL THEN 'utah_obd_proxy'
        ELSE NULL
    END AS outcome_label_source"""


EXTRACT_SQL = f"""
SELECT
    s.id AS internal_event_id,
    s.vin,
    s.test_start AS event_ts,
    lower(btrim(s.county)) AS source_era,
    CASE
        WHEN lower(btrim(s.county)) IN ('slc', 'slco') THEN 'salt_lake'
        ELSE lower(btrim(s.county))
    END AS public_county,
    {OUTCOME_SELECT_SQL},
    lower(nullif(btrim(s.program_type), '')) AS program_type,
    upper(nullif(btrim(s.test_type), '')) AS test_type,
    upper(nullif(btrim(v.make), '')) AS observed_make,
    upper(nullif(btrim(v.model), '')) AS observed_model,
    v.year AS observed_model_year
FROM data.inspection_search AS s
JOIN data.inspection_vehicle AS v USING (id)
WHERE s.test_start >= %(start_ts)s
  AND s.test_start < %(end_ts)s
  AND s.vin IS NOT NULL
  AND length(btrim(s.vin)) = 17
ORDER BY s.test_start, s.id
"""

OUTPUT_COLUMNS = (
    "internal_event_id",
    "vehicle_token",
    "vehicle_bucket",
    "event_ts",
    "source_era",
    "public_county",
    "canonical_outcome",
    "outcome_label_source",
    "program_type",
    "test_type",
    "observed_make",
    "observed_model",
    "observed_model_year",
)


def parse_timestamp(value: str) -> datetime:
    try:
        return datetime.fromisoformat(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(
            f"Invalid ISO timestamp or date: {value}"
        ) from exc


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--start", required=True, type=parse_timestamp)
    parser.add_argument("--end", required=True, type=parse_timestamp)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--page-size", type=int, default=10_000)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument(
        "--allow-external-output",
        action="store_true",
        help="Allow private row-level output outside data/private (unsafe).",
    )
    return parser.parse_args()


def validate_args(args: argparse.Namespace) -> Path:
    if args.end <= args.start:
        raise ValueError("--end must be later than --start")
    if (args.end - args.start).total_seconds() > MAX_BATCH_DAYS * 86400:
        raise ValueError(
            f"A source batch may span at most {MAX_BATCH_DAYS} days; "
            "export longer periods as closed monthly batches."
        )
    if not 100 <= args.page_size <= 50_000:
        raise ValueError("--page-size must be between 100 and 50000")

    output = args.output
    if output is None:
        name = f"inspection_{args.start:%Y%m%d}_{args.end:%Y%m%d}.csv.gz"
        output = PRIVATE_DATA_ROOT / "inspection_batches" / name
    output = output.resolve()
    if not args.allow_external_output and not output.is_relative_to(PRIVATE_DATA_ROOT):
        raise ValueError(
            f"Private exports must stay under {PRIVATE_DATA_ROOT}; "
            "use --allow-external-output only for controlled validation."
        )
    manifest = output.with_name(output.name + ".manifest.json")
    if (output.exists() or manifest.exists()) and not args.overwrite:
        raise FileExistsError(
            f"Refusing to overwrite {output} or its manifest; use --overwrite"
        )
    return output


def get_hmac_configuration() -> tuple[bytes, str, str]:
    encoded_key = os.environ.get("VIN_HASH_KEY", "").strip()
    key_file = Path(os.environ.get("VIN_HASH_KEY_FILE", DEFAULT_KEY_FILE))
    if not encoded_key and key_file.exists():
        mode = stat.S_IMODE(key_file.stat().st_mode)
        if mode & 0o077:
            raise ValueError(f"VIN HMAC key file must be mode 0600: {key_file}")
        encoded_key = key_file.read_text(encoding="ascii").strip()

    version = os.environ.get("VIN_HASH_KEY_VERSION", "v1").strip() or "v1"
    if not KEY_VERSION_PATTERN.fullmatch(version):
        raise ValueError("VIN_HASH_KEY_VERSION contains unsupported characters")
    try:
        key = bytes.fromhex(encoded_key)
    except ValueError as exc:
        raise ValueError("VIN_HASH_KEY must be a hexadecimal value") from exc
    if len(key) != 32:
        raise ValueError(
            "VIN_HASH_KEY must encode exactly 32 random bytes (64 hex characters). "
            "See .env.example."
        )
    if len(set(key)) < 8:
        raise ValueError(
            "VIN_HASH_KEY has insufficient byte diversity; generate 32 random bytes."
        )
    fingerprint = hashlib.sha256(
        b"utah-vehicle-health/key-fingerprint/v1\0" + key
    ).hexdigest()[:16]
    return key, version, fingerprint


def normalize_vin(vin: str) -> Optional[str]:
    normalized = vin.strip().upper()
    if not VIN_PATTERN.fullmatch(normalized):
        return None
    if len(set(normalized)) < 4:
        return None
    if normalized in {"12345678901234567", "98765432109876543"}:
        return None
    return normalized


def vehicle_token(vin: str, key: bytes, key_version: str) -> str:
    message = (
        f"utah-vehicle-health/{key_version}/vin\0".encode("ascii")
        + vin.encode("ascii")
    )
    return hmac.new(key, message, hashlib.sha256).hexdigest()


def private_row(
    row: Sequence[object], key: bytes, key_version: str
) -> Optional[tuple[object, ...]]:
    internal_event_id, vin, *remaining = row
    if not isinstance(vin, str):
        raise ValueError("The source returned a non-text VIN")
    normalized = normalize_vin(vin)
    if normalized is None:
        return None
    token = vehicle_token(normalized, key, key_version)
    bucket = int(token[:8], 16) % 100
    return (internal_event_id, token, bucket, *remaining)


def batches(
    connection: "psycopg2.extensions.connection",
    start: datetime,
    end: datetime,
    page_size: int,
) -> Iterable[list[tuple[object, ...]]]:
    with connection.cursor(name="uvh_inspection_batch") as cursor:
        cursor.itersize = page_size
        cursor.execute(EXTRACT_SQL, {"start_ts": start, "end_ts": end})
        while True:
            rows = cursor.fetchmany(page_size)
            if not rows:
                return
            yield rows


def connect_read_only() -> "psycopg2.extensions.connection":
    ssl_mode = os.environ.get("PGSSLMODE", "").lower()
    if ssl_mode not in SECURE_SSL_MODES:
        raise ValueError("PGSSLMODE must require TLS before exporting data")

    connection = psycopg2.connect(
        application_name="utah_vehicle_health_batch_export",
        connect_timeout=10,
        options=(
            "-c default_transaction_read_only=on "
            "-c statement_timeout=120000 "
            "-c lock_timeout=3000 "
            "-c idle_in_transaction_session_timeout=60000"
        ),
    )
    connection.set_session(
        readonly=True,
        isolation_level="REPEATABLE READ",
        autocommit=False,
    )
    with connection.cursor() as cursor:
        cursor.execute(
            """
            SELECT
                current_setting('transaction_read_only'),
                ssl
            FROM pg_stat_ssl
            WHERE pid = pg_backend_pid()
            """
        )
        read_only, ssl = cursor.fetchone()
    if read_only != "on" or not ssl:
        connection.close()
        raise RuntimeError("The database session is not read-only over TLS")
    return connection


def main() -> int:
    args = parse_args()
    os.umask(0o077)
    try:
        output = validate_args(args)
        key, key_version, key_fingerprint = get_hmac_configuration()
    except (FileExistsError, ValueError) as exc:
        print(f"Configuration error: {exc}", file=sys.stderr)
        return 2

    output.parent.mkdir(parents=True, exist_ok=True)
    partial = output.with_name(output.name + ".partial")
    manifest = output.with_name(output.name + ".manifest.json")
    partial_manifest = manifest.with_name(manifest.name + ".partial")

    try:
        connection = connect_read_only()
    except (ValueError, RuntimeError, psycopg2.Error) as exc:
        print(f"Secure read-only connection failed: {type(exc).__name__}", file=sys.stderr)
        return 1

    source_rows_read = 0
    rows_written = 0
    invalid_vins_skipped = 0
    last_order_key: Optional[tuple[object, object]] = None
    try:
        with gzip.open(partial, "wt", encoding="utf-8", newline="") as handle:
            writer = csv.writer(handle)
            writer.writerow(OUTPUT_COLUMNS)
            for source_rows in batches(
                connection, args.start, args.end, args.page_size
            ):
                for source_row in source_rows:
                    source_rows_read += 1
                    order_key = (source_row[2], source_row[0])
                    if last_order_key is not None and order_key <= last_order_key:
                        raise RuntimeError(
                            "Source join did not return a unique increasing "
                            "event_ts/internal_event_id order"
                        )
                    last_order_key = order_key
                    transformed = private_row(source_row, key, key_version)
                    if transformed is None:
                        invalid_vins_skipped += 1
                        continue
                    writer.writerow(transformed)
                    rows_written += 1
        connection.rollback()
    except Exception:
        partial.unlink(missing_ok=True)
        partial_manifest.unlink(missing_ok=True)
        raise
    finally:
        connection.close()

    digest = hashlib.sha256()
    with partial.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)

    manifest_data = {
        "export_kind": "bounded_batch",
        "query_version": QUERY_VERSION,
        "query_sha256": hashlib.sha256(EXTRACT_SQL.encode("utf-8")).hexdigest(),
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "source_start_inclusive": args.start.isoformat(),
        "source_end_exclusive": args.end.isoformat(),
        "vehicle_token_key_version": key_version,
        "vehicle_token_key_fingerprint": key_fingerprint,
        "source_rows_read": source_rows_read,
        "rows": rows_written,
        "invalid_vins_skipped": invalid_vins_skipped,
        "columns": OUTPUT_COLUMNS,
        "compressed_file_sha256": digest.hexdigest(),
        "compressed_file_bytes": partial.stat().st_size,
        "classification": "private_pseudonymized_analytical_staging",
        "source_transaction": "read_only_repeatable_read",
    }
    partial_manifest.write_text(
        json.dumps(manifest_data, indent=2) + "\n", encoding="utf-8"
    )
    os.replace(partial, output)
    os.replace(partial_manifest, manifest)
    print(f"Exported {rows_written:,} private pseudonymized rows to {output}")
    if invalid_vins_skipped:
        print(f"Skipped {invalid_vins_skipped:,} invalid/placeholder VIN rows")
    print(f"Wrote manifest to {manifest}")
    return 0


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