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

import argparse
import csv
import gzip
import json
import os
import re
from pathlib import Path
from typing import Any

try:
    import pymysql
except ImportError:  # pragma: no cover
    pymysql = None

from load_salesforce_snapshot_to_tidb import build_export_object_map, collect_manifest_objects, load_export_config


DEFAULT_MANIFEST_NAME = "manifest.json"
DEFAULT_EXPORT_CONFIG = "configs/daily_full_export.json"
DATETIME_PATTERNS = (
    "%Y-%m-%dT%H:%M:%S.000Z",
    "%Y-%m-%dT%H:%M:%S.000+0000",
    "%Y-%m-%d %H:%M:%S",
)
DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
INT_RE = re.compile(r"^-?\d+$")
DECIMAL_RE = re.compile(r"^-?\d+(\.\d+)?$")
TEXT_NAME_HINTS = ("description", "body", "headers", "tagcsv", "oldvalue", "newvalue")
VARCHAR_NAME_HINTS = ("accesslevel", "rowcause", "datatype", "field", "visibility", "sharetype")
NOT_NULL_NAME_HINTS = {"id", "createddate", "lastmodifieddate", "systemmodstamp"}
INDEX_HINTS = ("isdeleted", "createddate", "lastmodifieddate", "ownerid")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Create missing TiDB tables inferred from Salesforce snapshot CSV headers.")
    parser.add_argument("--snapshot-dir", required=True, help="Snapshot directory containing manifest and csv.gz files.")
    parser.add_argument("--manifest-name", default=DEFAULT_MANIFEST_NAME, help="Snapshot manifest file name.")
    parser.add_argument(
        "--export-config",
        default=str(Path(__file__).resolve().parent / DEFAULT_EXPORT_CONFIG),
        help="Export JSON config used to resolve object -> target_table metadata.",
    )
    parser.add_argument("--objects", help="Optional comma-separated object allowlist.")
    parser.add_argument("--db-host", default=os.environ.get("TIDB_HOST", ""))
    parser.add_argument("--db-port", type=int, default=int(os.environ.get("TIDB_PORT", "4000")))
    parser.add_argument("--db-user", default=os.environ.get("TIDB_USER", ""))
    parser.add_argument("--db-password", default=os.environ.get("TIDB_PASSWORD", ""))
    parser.add_argument("--db-database", default=os.environ.get("TIDB_DATABASE", ""))
    parser.add_argument("--ssl-mode", default=os.environ.get("TIDB_SSL_MODE", "REQUIRED"))
    parser.add_argument("--sample-rows", type=int, default=200)
    parser.add_argument("--apply", action="store_true", help="Execute CREATE TABLE for missing tables.")
    return parser.parse_args()


def parse_list(raw: str | None) -> set[str]:
    if not raw:
        return set()
    return {item.strip() for item in raw.split(",") if item.strip()}


def quote_ident(value: str) -> str:
    return f"`{value.replace('`', '``')}`"


def connect(args: argparse.Namespace):
    if pymysql is None:
        raise RuntimeError("pymysql is required")
    ssl_arg = {"ssl": {}} if str(args.ssl_mode or "").upper() != "DISABLED" else None
    return pymysql.connect(
        host=args.db_host,
        port=args.db_port,
        user=args.db_user,
        password=args.db_password,
        database=args.db_database,
        ssl=ssl_arg,
        charset="utf8mb4",
        autocommit=True,
    )


def discover_objects(snapshot_dir: Path, manifest_name: str, export_config: dict[str, Any]) -> list[dict[str, Any]]:
    export_object_map = build_export_object_map(export_config)
    return collect_manifest_objects(snapshot_dir, manifest_name, export_object_map)


def table_exists(conn: Any, table: str) -> bool:
    sql = (
        "SELECT 1 FROM information_schema.columns "
        "WHERE table_schema=%s AND lower(table_name)=lower(%s) LIMIT 1"
    )
    with conn.cursor() as cur:
        cur.execute(sql, (conn.db.decode() if isinstance(conn.db, bytes) else conn.db, table))
        return cur.fetchone() is not None


def read_header_and_samples(path: Path, sample_rows: int) -> tuple[list[str], list[dict[str, str]]]:
    with gzip.open(path, "rt", encoding="utf-8-sig", newline="") as src:
        reader = csv.DictReader(src)
        header = [str(name).strip() for name in (reader.fieldnames or []) if str(name).strip()]
        rows: list[dict[str, str]] = []
        for idx, row in enumerate(reader):
            if idx >= sample_rows:
                break
            rows.append({key: str(row.get(key) or "") for key in header})
    return header, rows


def is_bool_values(values: list[str]) -> bool:
    return bool(values) and all(value.lower() in {"true", "false"} for value in values)


def is_datetime_values(values: list[str]) -> bool:
    if not values:
        return False
    for value in values:
        matched = False
        for fmt in DATETIME_PATTERNS:
            try:
                from datetime import datetime
                datetime.strptime(value, fmt)
                matched = True
                break
            except ValueError:
                continue
        if not matched:
            return False
    return True


def is_date_values(values: list[str]) -> bool:
    return bool(values) and all(DATE_RE.match(value) for value in values)


def is_int_values(values: list[str]) -> bool:
    return bool(values) and all(INT_RE.match(value) for value in values)


def is_decimal_values(values: list[str]) -> bool:
    return bool(values) and all(DECIMAL_RE.match(value) for value in values)


def infer_type(column: str, rows: list[dict[str, str]]) -> str:
    values = [row[column].strip() for row in rows if row[column].strip()]
    lower_name = column.lower()
    max_len = max((len(value) for value in values), default=0)
    if lower_name == "id" or lower_name.endswith("id") or lower_name.endswith("_id") or lower_name.endswith("__pc"):
        return "varchar(255)"
    if lower_name.endswith("_type"):
        return "varchar(255)"
    if lower_name == "isdeleted" or lower_name.startswith("is"):
        return "tinyint(1)"
    if lower_name in {"createddate", "lastmodifieddate", "systemmodstamp"} or lower_name.endswith("datetime") or lower_name.endswith("timestamp"):
        return "datetime"
    if lower_name.endswith("date"):
        return "date"
    if is_bool_values(values):
        return "tinyint(1)"
    if is_datetime_values(values):
        return "datetime"
    if is_date_values(values):
        return "date"
    if is_int_values(values) or is_decimal_values(values):
        return "decimal(65,30)"
    if max_len > 255 or any(hint in lower_name for hint in TEXT_NAME_HINTS):
        return "text"
    if any(hint in lower_name for hint in VARCHAR_NAME_HINTS):
        return "varchar(255)"
    return "varchar(255)"


def infer_nullable(column: str, rows: list[dict[str, str]]) -> bool:
    lower_name = column.lower()
    if lower_name == "id":
        return False
    if lower_name == "isdeleted":
        return False
    return True


def column_sql(column: str, data_type: str, nullable: bool) -> str:
    if data_type == "tinyint(1)" and column.lower() == "isdeleted":
        return f"{quote_ident(column)} tinyint(1) NOT NULL DEFAULT '0'"
    null_sql = "DEFAULT NULL" if nullable else "NOT NULL"
    return f"{quote_ident(column)} {data_type} {null_sql}"


def create_table_sql(table: str, columns: list[str], rows: list[dict[str, str]]) -> str:
    definitions = []
    index_defs = []
    present = {column.lower(): column for column in columns}
    for column in columns:
        data_type = infer_type(column, rows)
        nullable = infer_nullable(column, rows)
        definitions.append(column_sql(column, data_type, nullable))
    if "id" not in present:
        raise ValueError(f"{table} snapshot missing Id column")
    definitions.append(f"PRIMARY KEY ({quote_ident(present['id'])}) /*T![clustered_index] CLUSTERED */")
    for key in INDEX_HINTS:
        if key in present:
            col = present[key]
            index_defs.append(f"KEY {quote_ident(f'idx_{table}_{key}')} ({quote_ident(col)})")
    all_defs = definitions + index_defs
    body = ",\n  ".join(all_defs)
    return (
        f"CREATE TABLE {quote_ident(table)} (\n"
        f"  {body}\n"
        f") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
    )


def main() -> int:
    args = parse_args()
    snapshot_dir = Path(args.snapshot_dir).expanduser().resolve()
    export_config = load_export_config(Path(args.export_config).expanduser().resolve())
    include = parse_list(args.objects)
    objects = discover_objects(snapshot_dir, args.manifest_name, export_config)

    conn = connect(args)
    results: list[dict[str, Any]] = []
    for item in objects:
        object_name = str(item.get("object", "")).strip()
        target_table = str(item.get("target_table", "")).strip()
        file_name = str(item.get("file", "")).strip()
        if not object_name or not target_table or not file_name:
            continue
        if include and object_name not in include:
            continue
        path = snapshot_dir / file_name
        if not path.exists():
            results.append({"object": object_name, "target_table": target_table, "status": "missing_file"})
            continue
        if table_exists(conn, target_table):
            results.append({"object": object_name, "target_table": target_table, "status": "exists"})
            continue
        header, rows = read_header_and_samples(path, args.sample_rows)
        sql = create_table_sql(target_table, header, rows)
        if args.apply:
            with conn.cursor() as cur:
                cur.execute(sql)
            status = "created"
        else:
            status = "planned"
        results.append(
            {
                "object": object_name,
                "target_table": target_table,
                "status": status,
                "column_count": len(header),
                "sql": sql,
            }
        )

    print(json.dumps({"snapshot_dir": str(snapshot_dir), "results": results}, indent=2))
    return 0


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