#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations

import argparse
from pathlib import Path
from time import perf_counter
from typing import Any

import pymysql
from dotenv import load_dotenv

from cm93db.config import load_config


TABLES = {
    "POINT": ("chart_point_geometry", "spx_point_geom"),
    "LINE": ("chart_line_geometry", "spx_line_geom"),
    "AREA": ("chart_area_geometry", "spx_area_geom"),
}


def connect() -> pymysql.connections.Connection:
    db = load_config(require_database=True)
    return pymysql.connect(
        host=db.host,
        port=db.port,
        user=db.user,
        password=db.password,
        database=db.database,
        charset=db.charset,
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=False,
        connect_timeout=20,
        read_timeout=900,
        write_timeout=30,
    )


def envelope_wkt(bbox: tuple[float, float, float, float]) -> str:
    west, south, east, north = bbox
    return (
        "POLYGON(("
        f"{west:.17g} {south:.17g},"
        f"{east:.17g} {south:.17g},"
        f"{east:.17g} {north:.17g},"
        f"{west:.17g} {north:.17g},"
        f"{west:.17g} {south:.17g}"
        "))"
    )


def fetch_old(conn, dataset_id: int, scale: int, geometry_type: str, bbox):
    west, south, east, north = bbox
    table, _idx = TABLES[geometry_type]
    with conn.cursor() as cur:
        cur.execute(
            f"""
            SELECT f.id AS feature_id
            FROM chart_feature f
            JOIN {table} g ON g.feature_id=f.id
            WHERE f.dataset_id=%s
              AND f.native_scale=%s
              AND f.geometry_type=%s
              AND f.bbox_max_lon >= %s
              AND f.bbox_min_lon <= %s
              AND f.bbox_max_lat >= %s
              AND f.bbox_min_lat <= %s
            """,
            (dataset_id, scale, geometry_type, west, east, south, north),
        )
        return sorted(int(r["feature_id"]) for r in cur.fetchall())


def fetch_new(conn, dataset_id: int, scale: int, geometry_type: str, bbox):
    west, south, east, north = bbox
    table, idx = TABLES[geometry_type]
    wkt = envelope_wkt(bbox)
    with conn.cursor() as cur:
        cur.execute(
            f"""
            SELECT f.id AS feature_id
            FROM {table} g FORCE INDEX ({idx})
            STRAIGHT_JOIN chart_feature f ON f.id=g.feature_id
            WHERE MBRIntersects(g.geom, ST_GeomFromText(%s, 4326))
              AND f.dataset_id=%s
              AND f.native_scale=%s
              AND f.geometry_type=%s
              AND f.bbox_max_lon >= %s
              AND f.bbox_min_lon <= %s
              AND f.bbox_max_lat >= %s
              AND f.bbox_min_lat <= %s
            """,
            (wkt, dataset_id, scale, geometry_type, west, east, south, north),
        )
        return sorted(int(r["feature_id"]) for r in cur.fetchall())


def explain_new(conn, dataset_id: int, scale: int, geometry_type: str, bbox):
    west, south, east, north = bbox
    table, idx = TABLES[geometry_type]
    wkt = envelope_wkt(bbox)
    with conn.cursor() as cur:
        cur.execute(
            f"""
            EXPLAIN
            SELECT f.id AS feature_id
            FROM {table} g FORCE INDEX ({idx})
            STRAIGHT_JOIN chart_feature f ON f.id=g.feature_id
            WHERE MBRIntersects(g.geom, ST_GeomFromText(%s, 4326))
              AND f.dataset_id=%s
              AND f.native_scale=%s
              AND f.geometry_type=%s
              AND f.bbox_max_lon >= %s
              AND f.bbox_min_lon <= %s
              AND f.bbox_max_lat >= %s
              AND f.bbox_min_lat <= %s
            """,
            (wkt, dataset_id, scale, geometry_type, west, east, south, north),
        )
        return list(cur.fetchall())


def soundings_old(conn, dataset_id: int, scale: int, bbox):
    wkt = envelope_wkt(bbox)
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT s.id
            FROM chart_sounding_point s
            JOIN chart_feature f ON f.id=s.feature_id
            WHERE f.dataset_id=%s
              AND f.native_scale=%s
              AND MBRContains(ST_GeomFromText(%s, 4326), s.geom)
            """,
            (dataset_id, scale, wkt),
        )
        return sorted(int(r["id"]) for r in cur.fetchall())


def soundings_new(conn, dataset_id: int, scale: int, bbox):
    wkt = envelope_wkt(bbox)
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT s.id
            FROM chart_sounding_point s FORCE INDEX (spx_sounding_geom)
            STRAIGHT_JOIN chart_feature f ON f.id=s.feature_id
            WHERE MBRContains(ST_GeomFromText(%s, 4326), s.geom)
              AND f.dataset_id=%s
              AND f.native_scale=%s
            """,
            (wkt, dataset_id, scale),
        )
        return sorted(int(r["id"]) for r in cur.fetchall())



def soundings_explain_new(conn, dataset_id: int, scale: int, bbox):
    wkt = envelope_wkt(bbox)
    with conn.cursor() as cur:
        cur.execute(
            """
            EXPLAIN
            SELECT s.id
            FROM chart_sounding_point s FORCE INDEX (spx_sounding_geom)
            STRAIGHT_JOIN chart_feature f ON f.id=s.feature_id
            WHERE MBRContains(ST_GeomFromText(%s, 4326), s.geom)
              AND f.dataset_id=%s
              AND f.native_scale=%s
            """,
            (wkt, dataset_id, scale),
        )
        return list(cur.fetchall())


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-id", type=int, default=1)
    parser.add_argument("--scale", type=int, default=100000)
    parser.add_argument(
        "--bbox",
        nargs=4,
        type=float,
        metavar=("WEST", "SOUTH", "EAST", "NORTH"),
        default=(132.53356933593753, 34.15727269301868, 132.72033691406253, 34.31168124115256),
    )
    parser.add_argument("--skip-soundings", action="store_true")
    args = parser.parse_args()

    load_dotenv(dotenv_path=Path(".env").resolve())
    bbox = tuple(float(v) for v in args.bbox)

    print("CM95 v1.4.6 query equivalence / plan check")
    print("=" * 78)
    print(f"Dataset : {args.dataset_id}")
    print(f"Scale   : {args.scale}")
    print(f"BBox    : {bbox}")

    with connect() as conn:
        for geometry_type in ("POINT", "LINE", "AREA"):
            t0 = perf_counter()
            old_ids = fetch_old(conn, args.dataset_id, args.scale, geometry_type, bbox)
            old_seconds = perf_counter() - t0
            t1 = perf_counter()
            new_ids = fetch_new(conn, args.dataset_id, args.scale, geometry_type, bbox)
            new_seconds = perf_counter() - t1
            if old_ids != new_ids:
                old_set = set(old_ids)
                new_set = set(new_ids)
                print(f"{geometry_type:5} : FAIL old={len(old_ids)} new={len(new_ids)}")
                print(f"  old_only sample: {sorted(old_set-new_set)[:20]}")
                print(f"  new_only sample: {sorted(new_set-old_set)[:20]}")
                raise SystemExit(2)

            explain = explain_new(conn, args.dataset_id, args.scale, geometry_type, bbox)
            first = explain[0] if explain else {}
            expected_idx = TABLES[geometry_type][1]
            if first.get("table") != "g" or first.get("key") != expected_idx or first.get("type") != "range":
                print(f"{geometry_type:5} : FAIL spatial plan={first}")
                raise SystemExit(2)

            speedup = (old_seconds / new_seconds) if new_seconds > 0 else float("inf")
            print(
                f"{geometry_type:5} : PASS ids={len(old_ids)} "
                f"old={old_seconds:.3f}s new={new_seconds:.3f}s speedup={speedup:.1f}x "
                f"driver={first.get('table')} type={first.get('type')} "
                f"key={first.get('key')} rows={first.get('rows')}"
            )

        if not args.skip_soundings:
            old_s = soundings_old(conn, args.dataset_id, args.scale, bbox)
            new_s = soundings_new(conn, args.dataset_id, args.scale, bbox)
            if old_s != new_s:
                print(f"SOUND : FAIL old={len(old_s)} new={len(new_s)}")
                raise SystemExit(2)
            sx = soundings_explain_new(conn, args.dataset_id, args.scale, bbox)
            first = sx[0] if sx else {}
            if first.get("table") != "s" or first.get("key") != "spx_sounding_geom" or first.get("type") != "range":
                print(f"SOUND : FAIL spatial plan={first}")
                raise SystemExit(2)
            print(
                f"SOUND : PASS ids={len(old_s)} semantics=MBRContains unchanged "
                f"driver={first.get('table')} type={first.get('type')} "
                f"key={first.get('key')} rows={first.get('rows')}"
            )

        conn.rollback()

    print("OVERALL: PASS")


if __name__ == "__main__":
    main()