from __future__ import annotations

import json
import os
import re
import tempfile
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import pymysql


TILE_PATTERN = re.compile(r"^(?P<z>\d+)_(?P<x>\d+)_(?P<y>\d+)\.json$")


@dataclass(frozen=True)
class DbConfig:
    host: str = os.getenv("NAVSEA_DB_HOST", "localhost")
    port: int = int(os.getenv("NAVSEA_DB_PORT", "3306"))
    user: str = os.getenv("NAVSEA_DB_USER", "root")
    password: str = os.getenv("NAVSEA_DB_PASSWORD", "2chi9ks2")
    database: str = os.getenv("NAVSEA_DB_NAME", "pbf_analysis")
    unix_socket: str | None = os.getenv("NAVSEA_DB_SOCKET", "/tmp/mysql.sock")


class GeoJsonStageLoader:
    def __init__(
        self,
        config: DbConfig,
        geojson_dir: Path,
        stage_table: str = "feature_geometry_stage",
    ) -> None:
        self.config = config
        self.geojson_dir = geojson_dir
        self.stage_table = stage_table

    def connect(self, *, streaming: bool = False, local_infile: bool = False):
        kwargs = {
            "host": self.config.host,
            "port": self.config.port,
            "user": self.config.user,
            "password": self.config.password,
            "database": self.config.database,
            "charset": "utf8mb4",
            "autocommit": False,
            "local_infile": local_infile,
            "cursorclass": pymysql.cursors.SSCursor if streaming else pymysql.cursors.Cursor,
        }
        if self.config.unix_socket and self.config.host in {"localhost", "127.0.0.1"}:
            kwargs["unix_socket"] = self.config.unix_socket
        return pymysql.connect(**kwargs)

    def load(self) -> None:
        tile_files = self.list_tile_files()
        if not tile_files:
            raise RuntimeError(f"no geojson files found in {self.geojson_dir}")

        lookup = {tile: path for tile, path in tile_files}
        ordered_tiles = [tile for tile, _ in tile_files]

        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            newline="",
            suffix=".tsv",
            delete=False,
            dir="/tmp",
        ) as temp_file:
            temp_path = Path(temp_file.name)
            self.build_lookup_table()
            staged_rows = self.stream_and_stage(lookup, ordered_tiles, temp_file)

        with self.connect(local_infile=True) as conn:
            with conn.cursor() as cur:
                self.create_stage_table(cur)
                cur.execute(
                    f"""
                    LOAD DATA LOCAL INFILE %s
                    INTO TABLE {self.stage_table}
                    CHARACTER SET utf8mb4
                    FIELDS TERMINATED BY '\t'
                    LINES TERMINATED BY '\n'
                    (feature_id, geometry_wkt)
                    """,
                    (str(temp_path),),
                )
                cur.execute(f"SELECT COUNT(*) FROM {self.stage_table}")
                loaded_rows = cur.fetchone()[0]
            conn.commit()

        temp_path.unlink(missing_ok=True)
        print(f"Loaded {loaded_rows} rows into {self.stage_table} from {len(tile_files)} geojson tiles")
        if loaded_rows != staged_rows:
            raise RuntimeError(f"stage row mismatch: staged={staged_rows}, loaded={loaded_rows}")

    def list_tile_files(self) -> list[tuple[tuple[int, int, int], Path]]:
        rows = []
        for path in self.geojson_dir.glob("*.json"):
            match = TILE_PATTERN.match(path.name)
            if not match:
                continue
            rows.append(
                (
                    (int(match.group("z")), int(match.group("x")), int(match.group("y"))),
                    path,
                )
            )
        rows.sort()
        return rows

    def stream_and_stage(self, lookup: dict[tuple[int, int, int], Path], ordered_tiles: list[tuple[int, int, int]], temp_file) -> int:
        staged_rows = 0
        tile_index = 0
        current_tile = None
        current_rows: list[tuple[int, str, str, str]] = []

        with self.connect(streaming=True) as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT
                        z,
                        x,
                        y,
                        feature_id,
                        fid,
                        vt_layer,
                        geom_type
                    FROM feature_geojson_lookup FORCE INDEX (idx_tile_order)
                    ORDER BY z, x, y, feature_id
                    """
                )

                for z, x, y, feature_id, fid, vt_layer, geom_type in cur:
                    tile = (int(z), int(x), int(y))
                    if current_tile is None:
                        current_tile = tile
                    if tile != current_tile:
                        staged_rows += self.flush_tile(
                            current_tile,
                            current_rows,
                            lookup,
                            temp_file,
                        )
                        tile_index += 1
                        if tile_index % 5000 == 0:
                            print(f"processed {tile_index}/{len(ordered_tiles)} tiles, staged {staged_rows} geometries")
                        current_tile = tile
                        current_rows = []

                    current_rows.append((int(feature_id), str(fid), str(vt_layer), str(geom_type)))

                if current_tile is not None:
                    staged_rows += self.flush_tile(current_tile, current_rows, lookup, temp_file)
                    tile_index += 1

        if tile_index != len(ordered_tiles):
            raise RuntimeError(f"tile count mismatch: db_stream={tile_index}, geojson={len(ordered_tiles)}")
        return staged_rows

    def build_lookup_table(self) -> None:
        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute("DROP TABLE IF EXISTS feature_geojson_lookup")
                cur.execute(
                    """
                    CREATE TABLE feature_geojson_lookup AS
                    SELECT
                        f.z,
                        f.x,
                        f.y,
                        f.id AS feature_id,
                        pfid.v AS fid,
                        pvt.v AS vt_layer,
                        f.geom_type
                    FROM features f FORCE INDEX (idx_tile_order)
                    LEFT JOIN properties pfid
                        ON pfid.feature_id = f.id AND pfid.k = 'fid'
                    LEFT JOIN properties pvt
                        ON pvt.feature_id = f.id AND pvt.k = 'vt_layer'
                    """
                )
                cur.execute("ALTER TABLE feature_geojson_lookup ADD PRIMARY KEY (feature_id)")
                cur.execute("ALTER TABLE feature_geojson_lookup ADD KEY idx_tile_order (z, x, y, feature_id)")
            conn.commit()

    def flush_tile(
        self,
        tile: tuple[int, int, int],
        db_rows: list[tuple[int, str, str, str]],
        lookup: dict[tuple[int, int, int], Path],
        temp_file,
    ) -> int:
        path = lookup.get(tile)
        if path is None:
            raise RuntimeError(f"missing geojson for tile {tile}")

        data = json.loads(path.read_text(encoding="utf-8"))
        features = data.get("features", [])

        geo_groups: dict[tuple[str, str, str], list[str]] = defaultdict(list)
        for feature in features:
            props = feature.get("properties", {})
            fid = props.get("fid")
            vt_layer = props.get("vt_layer")
            geometry = feature.get("geometry")
            if fid is None or vt_layer is None or geometry is None:
                continue
            key = (str(fid), str(vt_layer), str(geometry.get("type")))
            geo_groups[key].append(self.geometry_to_wkt(geometry))

        db_groups: dict[tuple[str, str, str], list[int]] = defaultdict(list)
        for feature_id, fid, vt_layer, geom_type in db_rows:
            key = (fid, vt_layer, geom_type)
            db_groups[key].append(feature_id)

        if set(geo_groups) != set(db_groups):
            missing_in_db = list(sorted(set(geo_groups) - set(db_groups)))[:10]
            missing_in_geo = list(sorted(set(db_groups) - set(geo_groups)))[:10]
            raise RuntimeError(
                f"tile key mismatch for {path.name}: "
                f"missing_in_db={missing_in_db}, missing_in_geo={missing_in_geo}"
            )

        inserted = 0
        for key in sorted(geo_groups):
            geo_wkts = geo_groups[key]
            db_feature_ids = db_groups[key]
            if len(geo_wkts) != len(db_feature_ids):
                raise RuntimeError(
                    f"tile multiplicity mismatch for {path.name}, key={key}: "
                    f"geojson={len(geo_wkts)}, db={len(db_feature_ids)}"
                )
            for feature_id, geometry_wkt in zip(db_feature_ids, geo_wkts):
                temp_file.write(f"{feature_id}\t{geometry_wkt}\n")
                inserted += 1

        return inserted

    @staticmethod
    def create_stage_table(cur) -> None:
        cur.execute("DROP TABLE IF EXISTS feature_geometry_stage")
        cur.execute(
            """
            CREATE TABLE feature_geometry_stage (
                feature_id BIGINT NOT NULL,
                geometry_wkt LONGTEXT NOT NULL,
                PRIMARY KEY (feature_id)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """
        )

    @staticmethod
    def geometry_to_wkt(geometry: dict) -> str:
        geom_type = geometry["type"]
        coordinates = geometry["coordinates"]

        def fmt_pair(pair):
            return f"{pair[0]} {pair[1]}"

        def fmt_ring(ring):
            return "(" + ", ".join(fmt_pair(pair) for pair in ring) + ")"

        if geom_type == "Point":
            return f"POINT ({fmt_pair(coordinates)})"
        if geom_type == "MultiPoint":
            return "MULTIPOINT (" + ", ".join(f"({fmt_pair(pair)})" for pair in coordinates) + ")"
        if geom_type == "LineString":
            return "LINESTRING (" + ", ".join(fmt_pair(pair) for pair in coordinates) + ")"
        if geom_type == "MultiLineString":
            return "MULTILINESTRING (" + ", ".join("(" + ", ".join(fmt_pair(pair) for pair in line) + ")" for line in coordinates) + ")"
        if geom_type == "Polygon":
            return "POLYGON (" + ", ".join(fmt_ring(ring) for ring in coordinates) + ")"
        if geom_type == "MultiPolygon":
            polygons = []
            for polygon in coordinates:
                polygons.append("(" + ", ".join(fmt_ring(ring) for ring in polygon) + ")")
            return "MULTIPOLYGON (" + ", ".join(polygons) + ")"
        raise RuntimeError(f"unsupported geometry type: {geom_type}")


def main() -> None:
    loader = GeoJsonStageLoader(
        DbConfig(),
        Path("/home/wwwroot/newpec/exported_auto/geojson"),
    )
    loader.load()


if __name__ == "__main__":
    main()
