#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import resource
import sys
import time
from dataclasses import dataclass, replace
from decimal import Decimal, InvalidOperation
from pathlib import Path
from typing import Any
from urllib.parse import unquote

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from csv_parser import iter_csv_rows  # noqa: E402
from csv_parser import new as new_csv_parser  # noqa: E402
from metadata import Column  # noqa: E402
from row_decoder import build_ordered_row_decoder_into  # noqa: E402
from type_mapping import (  # noqa: E402
    DECIMAL_BASE_TYPES,
    FLOAT_BASE_TYPES,
    INTEGER_BASE_TYPES,
    TEXT_BASE_TYPES,
    base_source_type,
    build_value_coercer,
    is_source_type_syncable,
)

SNAPSHOT_SCHEMA_SUFFIX = "-schema.sql"
SNAPSHOT_SCHEMA_CREATE_SUFFIX = "-schema-create.sql"


@dataclass(frozen=True)
class BenchTable:
    source_schema: str
    source_table: str
    snapshot_file_stem: str
    columns: tuple[Column, ...]

    @property
    def column_names(self) -> list[str]:
        return [column.name for column in self.columns]

    @property
    def label(self) -> str:
        return f"{self.source_schema}.{self.source_table}"


@dataclass
class BlackholeSink:
    upserts: int = 0

    def upsert(self, *, table: str, data: dict[str, Any]) -> None:
        self.upserts += 1


@dataclass
class BenchStats:
    files: int = 0
    bytes_processed: int = 0
    rows: int = 0
    duration_ms: float = 0.0

    def add(self, other: BenchStats) -> None:
        self.files += other.files
        self.bytes_processed += other.bytes_processed
        self.rows += other.rows
        self.duration_ms += other.duration_ms


def main() -> int:
    parser = argparse.ArgumentParser(description="Benchmark snapshot CSV parse/decode with a blackhole sink.")
    parser.add_argument("snapshot_dir", type=Path)
    parser.add_argument(
        "--mode",
        choices=("csv-only", "decode", "blackhole", "all"),
        default="blackhole",
        help="csv-only parses rows; decode also converts values; blackhole also calls an empty upsert.",
    )
    parser.add_argument("--table", action="append", default=[], help="Limit to source table name, e.g. lineitem.")
    parser.add_argument("--limit-files", type=int, default=0)
    parser.add_argument("--limit-rows", type=int, default=0)
    parser.add_argument("--null-marker", default=r"\N")
    parser.add_argument("--binary-encoding", default="utf8")
    parser.add_argument("--escape-backslash", action="store_true")
    parser.add_argument("--diagnostic-columns", action="store_true")
    parser.add_argument(
        "--decimal-mode",
        choices=("decimal", "regex", "raw"),
        default="decimal",
        help="decimal matches connector behavior; regex/raw are benchmark-only variants.",
    )
    parser.add_argument(
        "--integer-mode",
        choices=("int", "raw"),
        default="int",
        help="int matches connector behavior; raw is a benchmark-only upper-bound variant.",
    )
    args = parser.parse_args()

    snapshot_dir = args.snapshot_dir.resolve()
    if not snapshot_dir.is_dir():
        raise SystemExit(f"snapshot dir does not exist: {snapshot_dir}")

    tables = list(iter_tables(snapshot_dir))
    if args.table:
        requested = set(args.table)
        tables = [table for table in tables if table.source_table in requested or table.label in requested]
    if not tables:
        raise SystemExit("no snapshot tables matched")

    modes = ("csv-only", "decode", "blackhole") if args.mode == "all" else (args.mode,)
    print_json(
        "blackhole snapshot bench starting",
        snapshot_dir=str(snapshot_dir),
        modes=list(modes),
        table_count=len(tables),
        limit_files=args.limit_files,
        limit_rows=args.limit_rows,
        escape_backslash=args.escape_backslash,
        diagnostic_columns=args.diagnostic_columns,
        decimal_mode=args.decimal_mode,
        integer_mode=args.integer_mode,
    )

    for mode in modes:
        total = BenchStats()
        start = time.perf_counter()
        for table in tables:
            table_stats = bench_table(
                snapshot_dir,
                table,
                mode=mode,
                null_marker=args.null_marker,
                binary_encoding=args.binary_encoding,
                escape_backslash=args.escape_backslash,
                diagnostic_columns=args.diagnostic_columns,
                decimal_mode=args.decimal_mode,
                integer_mode=args.integer_mode,
                limit_files=args.limit_files,
                limit_rows=args.limit_rows,
            )
            total.add(table_stats)
            print_result("blackhole snapshot table completed", mode=mode, table=table.label, stats=table_stats)
        total.duration_ms = elapsed_ms(start)
        print_result(
            "blackhole snapshot bench completed",
            mode=mode,
            table="all",
            stats=total,
            peak_rss_mb=peak_rss_mb(),
        )

    return 0


def bench_table(
    snapshot_dir: Path,
    table: BenchTable,
    *,
    mode: str,
    null_marker: str,
    binary_encoding: str,
    escape_backslash: bool,
    diagnostic_columns: bool,
    decimal_mode: str,
    integer_mode: str,
    limit_files: int,
    limit_rows: int,
) -> BenchStats:
    stats = BenchStats()
    decode_row_into = build_bench_row_decoder_into(
        columns=table.columns,
        null_marker=null_marker,
        binary_encoding=binary_encoding,
        decimal_mode=decimal_mode,
        integer_mode=integer_mode,
    )
    data: dict[str, Any] = {}
    sink = BlackholeSink()
    files = snapshot_files(snapshot_dir, table)
    if limit_files > 0:
        files = files[:limit_files]

    start = time.perf_counter()
    for file_path in files:
        file_rows = 0
        stats.files += 1
        stats.bytes_processed += file_path.stat().st_size
        with file_path.open("r", encoding="utf-8", newline="") as handle:
            reader = snapshot_csv_rows(handle, escape_backslash=escape_backslash)
            first_row = next(reader, None)
            if first_row is None:
                continue
            rows = reader if is_header(first_row, table.column_names) else chain_first(first_row, reader)
            for raw_values in rows:
                if mode != "csv-only":
                    decode_row_into(raw_values, data)
                    if diagnostic_columns:
                        data["_incremental_commit_ts"] = 0
                        data["_incremental_source_schema"] = table.source_schema
                        data["_incremental_source_table"] = table.source_table
                    if mode == "blackhole":
                        sink.upsert(table=table.label, data=data)
                stats.rows += 1
                file_rows += 1
                if limit_rows > 0 and stats.rows >= limit_rows:
                    stats.duration_ms = elapsed_ms(start)
                    return stats
        if file_rows == 0:
            continue
    stats.duration_ms = elapsed_ms(start)
    return stats


def iter_tables(snapshot_dir: Path) -> list[BenchTable]:
    tables: list[BenchTable] = []
    for schema_path in sorted(snapshot_dir.glob(f"*{SNAPSHOT_SCHEMA_SUFFIX}")):
        if schema_path.name.endswith(SNAPSHOT_SCHEMA_CREATE_SUFFIX):
            continue
        stem = schema_path.name[: -len(SNAPSHOT_SCHEMA_SUFFIX)]
        source_schema, source_table = parse_schema_stem(stem)
        columns = parse_create_table_columns(schema_path.read_text(encoding="utf-8"))
        tables.append(
            BenchTable(
                source_schema=source_schema,
                source_table=source_table,
                snapshot_file_stem=stem,
                columns=columns,
            )
        )
    return tables


def parse_schema_stem(stem: str) -> tuple[str, str]:
    if "." not in stem:
        raise ValueError(f"schema file name must be <schema>.<table>-schema.sql: {stem}")
    source_schema, source_table = stem.split(".", 1)
    return unquote(source_schema), unquote(source_table)


def parse_create_table_columns(sql: str) -> tuple[Column, ...]:
    body = create_table_body(sql)
    columns: list[Column] = []
    primary_key: list[str] = []

    for item in split_top_level_commas(body):
        text = item.strip()
        if not text:
            continue
        upper = text.upper()
        if upper.startswith("PRIMARY KEY"):
            primary_key = parse_identifier_list(text)
            continue
        if upper.startswith(("KEY ", "UNIQUE ", "CONSTRAINT ", "INDEX ")):
            continue
        name, rest = parse_column_identifier(text)
        source_type = parse_column_type(rest)
        precision, scale = column_precision_scale(source_type)
        columns.append(
            Column(
                name=name,
                source_type=source_type,
                nullable="NOT NULL" not in rest.upper(),
                precision=precision,
                scale=scale,
            )
        )

    primary_key_set = set(primary_key)
    return tuple(replace(column, primary_key=column.name in primary_key_set) for column in columns)


def create_table_body(sql: str) -> str:
    match = re.search(r"\bCREATE\s+TABLE\b", sql, flags=re.IGNORECASE)
    if match is None:
        raise ValueError("schema file does not contain CREATE TABLE")
    start = sql.find("(", match.end())
    if start < 0:
        raise ValueError("CREATE TABLE does not contain a column list")
    depth = 0
    in_backtick = False
    in_single = False
    in_double = False
    for index in range(start, len(sql)):
        char = sql[index]
        if in_backtick:
            if char == "`":
                if index + 1 < len(sql) and sql[index + 1] == "`":
                    continue
                in_backtick = False
            continue
        if in_single:
            if char == "'" and sql[index - 1] != "\\":
                in_single = False
            continue
        if in_double:
            if char == '"' and sql[index - 1] != "\\":
                in_double = False
            continue
        if char == "`":
            in_backtick = True
        elif char == "'":
            in_single = True
        elif char == '"':
            in_double = True
        elif char == "(":
            depth += 1
        elif char == ")":
            depth -= 1
            if depth == 0:
                return sql[start + 1 : index]
    raise ValueError("CREATE TABLE column list is not closed")


def split_top_level_commas(text: str) -> list[str]:
    parts: list[str] = []
    start = 0
    depth = 0
    in_backtick = False
    in_single = False
    in_double = False
    index = 0
    while index < len(text):
        char = text[index]
        if in_backtick:
            if char == "`":
                if index + 1 < len(text) and text[index + 1] == "`":
                    index += 1
                else:
                    in_backtick = False
        elif in_single:
            if char == "'" and text[index - 1] != "\\":
                in_single = False
        elif in_double:
            if char == '"' and text[index - 1] != "\\":
                in_double = False
        elif char == "`":
            in_backtick = True
        elif char == "'":
            in_single = True
        elif char == '"':
            in_double = True
        elif char == "(":
            depth += 1
        elif char == ")":
            depth -= 1
        elif char == "," and depth == 0:
            parts.append(text[start:index])
            start = index + 1
        index += 1
    parts.append(text[start:])
    return parts


def parse_column_identifier(text: str) -> tuple[str, str]:
    text = text.lstrip()
    if text.startswith("`"):
        chars: list[str] = []
        index = 1
        while index < len(text):
            char = text[index]
            if char == "`":
                if index + 1 < len(text) and text[index + 1] == "`":
                    chars.append("`")
                    index += 2
                    continue
                return "".join(chars), text[index + 1 :].lstrip()
            chars.append(char)
            index += 1
        raise ValueError(f"unclosed column identifier: {text}")
    match = re.match(r"([A-Za-z_][0-9A-Za-z_$]*)\s*(.*)$", text, flags=re.DOTALL)
    if match is None:
        raise ValueError(f"failed to parse column identifier: {text}")
    return match.group(1), match.group(2).lstrip()


def parse_column_type(rest: str) -> str:
    match = re.match(r"([A-Za-z]+(?:\s*\([^)]*\))?(?:\s+UNSIGNED)?)\b", rest, flags=re.IGNORECASE)
    if match is None:
        raise ValueError(f"failed to parse column type: {rest}")
    return re.sub(r"\s+", " ", match.group(1).strip())


def parse_identifier_list(text: str) -> list[str]:
    start = text.find("(")
    end = text.find(")", start + 1)
    if start < 0 or end < 0:
        return []
    inner = text[start + 1 : end]
    result: list[str] = []
    for item in split_top_level_commas(inner):
        name, _rest = parse_column_identifier(item)
        result.append(name)
    return result


def column_precision_scale(source_type: str) -> tuple[int | None, int | None]:
    base = base_source_type(source_type)
    params = type_int_params(source_type)
    if base in {"BIT", "TINYINT"}:
        return (params[0] if params else None), None
    if base in {"DECIMAL", "NUMERIC", "NEWDECIMAL"}:
        precision = params[0] if params else 10
        scale = params[1] if len(params) > 1 else 0
        return precision, scale
    return None, None


def type_int_params(source_type: str) -> tuple[int, ...]:
    match = re.search(r"\(([^)]*)\)", source_type)
    if match is None:
        return ()
    return tuple(int(part.strip()) for part in match.group(1).split(",") if part.strip())


def build_bench_row_decoder_into(
    *,
    columns: tuple[Column, ...],
    null_marker: str,
    binary_encoding: str,
    decimal_mode: str,
    integer_mode: str,
) -> Any:
    if decimal_mode == "decimal" and integer_mode == "int":
        return build_ordered_row_decoder_into(
            columns,
            null_marker=null_marker,
            binary_encoding=binary_encoding,
        )

    column_decoders = tuple(
        (
            index,
            column.name,
            bench_value_decoder(
                column,
                null_marker=null_marker,
                binary_encoding=binary_encoding,
                decimal_mode=decimal_mode,
                integer_mode=integer_mode,
            ),
        )
        for index, column in enumerate(columns)
        if is_source_type_syncable(column.source_type, precision=column.precision)
    )

    def decode_row_into(raw_values: list[str], decoded: dict[str, Any]) -> None:
        for index, name, decode_value in column_decoders:
            decoded[name] = decode_value(raw_values[index])

    return decode_row_into


def bench_value_decoder(
    column: Column,
    *,
    null_marker: str,
    binary_encoding: str,
    decimal_mode: str,
    integer_mode: str,
) -> Any:
    base = base_source_type(column.source_type)
    if base in TEXT_BASE_TYPES - {"SET"}:
        return null_text_decoder(null_marker)
    if base in {"DATE", "DATETIME", "TIME", "VECTOR", "XML"}:
        return empty_null_text_decoder(null_marker)
    if base in DECIMAL_BASE_TYPES:
        if decimal_mode == "raw":
            return null_text_decoder(null_marker)
        if decimal_mode == "regex":
            return regex_decimal_decoder(null_marker, column.source_type)
        return decimal_validation_decoder(null_marker, column.source_type)
    if base in INTEGER_BASE_TYPES and integer_mode == "raw":
        return null_text_decoder(null_marker)
    if base in FLOAT_BASE_TYPES and integer_mode == "raw":
        return null_text_decoder(null_marker)
    return build_value_coercer(
        column.source_type,
        null_marker=null_marker,
        binary_encoding=binary_encoding,
        precision=column.precision,
        primary_key=column.primary_key,
    )


def null_text_decoder(null_marker: str) -> Any:
    def decode(raw: str | None) -> str | None:
        if raw is None or raw == null_marker:
            return None
        return raw

    return decode


def empty_null_text_decoder(null_marker: str) -> Any:
    def decode(raw: str | None) -> str | None:
        if raw is None or raw == null_marker or raw == "":
            return None
        return raw

    return decode


def decimal_validation_decoder(null_marker: str, source_type: str) -> Any:
    def decode(raw: str | None) -> str | None:
        if raw is None or raw == null_marker or raw == "":
            return None
        try:
            Decimal(raw)
        except (InvalidOperation, ValueError) as exc:
            raise ValueError(f"failed to coerce {raw!r} as {source_type}: {type(exc).__name__}: {exc}") from exc
        return raw

    return decode


_DECIMAL_RE = re.compile(r"^[+-]?(?:\d+(?:\.\d*)?|\.\d+)$")


def regex_decimal_decoder(null_marker: str, source_type: str) -> Any:
    def decode(raw: str | None) -> str | None:
        if raw is None or raw == null_marker or raw == "":
            return None
        if _DECIMAL_RE.fullmatch(raw) is None:
            raise ValueError(f"failed to coerce {raw!r} as {source_type}: invalid decimal literal")
        return raw

    return decode


def snapshot_files(snapshot_dir: Path, table: BenchTable) -> list[Path]:
    return sorted(snapshot_dir.glob(f"{table.snapshot_file_stem}.*.csv"))


def is_header(row: list[str], expected_header: list[str]) -> bool:
    if row == expected_header:
        return True
    if len(row) == len(expected_header) and set(row) == set(expected_header):
        raise ValueError("snapshot CSV header order does not match SQL schema column order")
    return False


def chain_first(first: list[str], rows: Any) -> Any:
    yield first
    yield from rows


def snapshot_csv_rows(handle: Any, *, escape_backslash: bool) -> Any:
    parser = new_csv_parser(error_type=ValueError, context="snapshot", strict=True)
    return iter_csv_rows(handle, parser=parser, escape_backslash=escape_backslash)


def elapsed_ms(start: float) -> float:
    return round((time.perf_counter() - start) * 1000, 3)


def peak_rss_mb() -> float:
    return round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024, 3)


def print_result(message: str, *, mode: str, table: str, stats: BenchStats, peak_rss_mb: float | None = None) -> None:
    duration_s = stats.duration_ms / 1000
    us_per_row = (stats.duration_ms * 1000 / stats.rows) if stats.rows else 0
    bytes_per_row = (stats.bytes_processed / stats.rows) if stats.rows else 0
    payload: dict[str, Any] = {
        "message": message,
        "mode": mode,
        "table": table,
        "duration_ms": round(stats.duration_ms, 3),
        "files": stats.files,
        "rows": stats.rows,
        "bytes": stats.bytes_processed,
        "mb": round(stats.bytes_processed / (1024 * 1024), 3),
        "rows_per_sec": round(stats.rows / duration_s, 3) if duration_s > 0 else 0,
        "mb_per_sec": round((stats.bytes_processed / (1024 * 1024)) / duration_s, 3) if duration_s > 0 else 0,
        "us_per_row": round(us_per_row, 3),
        "bytes_per_row": round(bytes_per_row, 1),
    }
    if peak_rss_mb is not None:
        payload["process_peak_rss_mb"] = peak_rss_mb
    print(json.dumps(payload, sort_keys=True), flush=True)


def print_json(message: str, **fields: Any) -> None:
    print(json.dumps({"message": message, **fields}, sort_keys=True), flush=True)


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