"""Import the two Fengxian GIS deliveries into one project BizDB.

The source PRJ files declare EPSG:21481, while their coordinate values match the
project's Shanghai local CRS (SRID 900914). The importer records both facts and
assigns the validated local CRS without changing coordinate values.
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from dataclasses import dataclass
from datetime import date
from pathlib import Path

from psycopg import sql
from psycopg.conninfo import conninfo_to_dict

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from app.infra.db.project_routing import get_project_pgconn_string
from app.native.wndb.core.connection import project_connection, project_transaction


SCHEMA_SQL = PROJECT_ROOT / "resources/sql/010_source_gis_inventory.sql"
DEFAULT_APRIL_ROOT = Path(
    "/mnt/c/Users/admin/Desktop/data/奉贤/20200423全网除头桥数据导出"
)
DEFAULT_AUGUST_ROOT = Path(
    "/mnt/c/Users/admin/Desktop/data/奉贤/20200811奉贤100口径以上（含台账）(1)"
)


@dataclass(frozen=True)
class DatasetSpec:
    key: str
    title: str
    source_date: date


@dataclass(frozen=True)
class LayerSpec:
    key: str
    dataset: DatasetSpec
    display_name: str
    source_file: Path
    feature_kind: str
    asset_id_candidates: tuple[str, ...] = ()


APRIL_DATASET = DatasetSpec(
    key="fx_water_assets_20200423",
    title="奉贤全网除头桥 GIS 数据导出",
    source_date=date(2020, 4, 23),
)
AUGUST_DATASET = DatasetSpec(
    key="fx_pipe_inventory_20200811",
    title="奉贤 DN100 及以上管网台账",
    source_date=date(2020, 8, 11),
)


def layer_specs(april_root: Path, august_root: Path) -> tuple[LayerSpec, ...]:
    return (
        LayerSpec("fx_20200423_drain_valves", APRIL_DATASET, "排水阀", april_root / "排水阀.shp", "drain_valve", ("排水阀ID编", "排水阀ID编码")),
        LayerSpec("fx_20200423_service_connections", APRIL_DATASET, "接水点", april_root / "接水点.shp", "service_connection", ("接水点ID编", "接水点ID编码")),
        LayerSpec("fx_20200423_flow_meters", APRIL_DATASET, "流量仪", april_root / "流量仪.shp", "flow_meter", ("流量仪ID编", "流量仪ID编码")),
        LayerSpec("fx_20200423_hydrants", APRIL_DATASET, "消火栓", april_root / "消火栓.shp", "hydrant", ("消火栓ID编", "消火栓ID编码")),
        LayerSpec("fx_20200423_customer_meters", APRIL_DATASET, "用户水表", april_root / "用户水表.shp", "customer_meter", ("水表ID编码", "用户水表ID")),
        LayerSpec("fx_20200423_pipes", APRIL_DATASET, "管线", april_root / "管线.shp", "pipe", ("管线ID编码",)),
        LayerSpec("fx_20200423_valves", APRIL_DATASET, "阀门", april_root / "阀门.shp", "valve", ("阀门ID编码",)),
        LayerSpec("fx_20200423_non_residential_meters", APRIL_DATASET, "非居民用户水表", april_root / "非居民用户水表.shp", "non_residential_meter", ("水表ID编码", "用户水表ID")),
        LayerSpec(
            "fx_20200811_dn100_pipes",
            AUGUST_DATASET,
            "DN100 及以上管线台账",
            august_root / "20200811奉贤管网100及以上(含台账).shp",
            "pipe",
            ("PIPE_ID",),
        ),
    )


def _postgres_environment(database_name: str) -> dict[str, str]:
    parts = conninfo_to_dict(get_project_pgconn_string(db_name=database_name))
    env = os.environ.copy()
    for source, target in (
        ("host", "PGHOST"),
        ("port", "PGPORT"),
        ("user", "PGUSER"),
        ("password", "PGPASSWORD"),
        ("dbname", "PGDATABASE"),
    ):
        value = parts.get(source)
        if value:
            env[target] = value
    return env


def _load_staging_table(spec: LayerSpec, database_name: str) -> str:
    stage_name = f"_stage_{spec.key}"
    with project_connection(database_name) as conn:
        conn.execute(
            sql.SQL("DROP TABLE IF EXISTS {}.{} CASCADE").format(
                sql.Identifier("source_gis"), sql.Identifier(stage_name)
            )
        )
    command = [
        "ogr2ogr",
        "--config",
        "PG_USE_COPY",
        "YES",
        "--config",
        "OSR_USE_NON_DEPRECATED",
        "NO",
        "-f",
        "PostgreSQL",
        "PG:",
        str(spec.source_file),
        "-overwrite",
        "-preserve_fid",
        "-dim",
        "XY",
        "-mapFieldType",
        "Date=String",
        "-nln",
        stage_name,
        "-lco",
        "SCHEMA=source_gis",
        "-lco",
        "FID=source_fid",
        "-lco",
        "GEOMETRY_NAME=geom",
        "-lco",
        "LAUNDER=NO",
    ]
    subprocess.run(
        command,
        env=_postgres_environment(database_name),
        check=True,
        text=True,
    )
    return stage_name


def _first_column(columns: set[str], candidates: tuple[str, ...]) -> str | None:
    return next((candidate for candidate in candidates if candidate in columns), None)


def _column_as_text(column: str | None) -> sql.Composable:
    if column is None:
        return sql.SQL("NULL::text")
    return sql.SQL("NULLIF(BTRIM({}::text), '')").format(sql.Identifier(column))


def _column_as_float(column: str | None) -> sql.Composable:
    if column is None:
        return sql.SQL("NULL::double precision")
    return sql.SQL("{}::double precision").format(sql.Identifier(column))


def _store_layer(database_name: str, spec: LayerSpec, stage_name: str) -> int:
    with project_connection(database_name) as conn, conn.transaction(), conn.cursor() as cur:
        cur.execute(
            """
            SELECT a.attname
            FROM pg_attribute a
            WHERE a.attrelid = to_regclass(%s)
              AND a.attnum > 0
              AND NOT a.attisdropped
            """,
            (f"source_gis.{stage_name}",),
        )
        columns = {row["attname"] for row in cur.fetchall()}
        if not {"source_fid", "geom"}.issubset(columns):
            raise RuntimeError(f"Staging table for {spec.key} is incomplete")

        cur.execute(
            """
            INSERT INTO source_gis.datasets (
                dataset_key, title, source_date, declared_srid, effective_srid,
                coordinate_note, imported_at
            ) VALUES (%s, %s, %s, 21481, 900914, %s, NOW())
            ON CONFLICT (dataset_key) DO UPDATE SET
                title = EXCLUDED.title,
                source_date = EXCLUDED.source_date,
                declared_srid = EXCLUDED.declared_srid,
                effective_srid = EXCLUDED.effective_srid,
                coordinate_note = EXCLUDED.coordinate_note,
                imported_at = NOW()
            """,
            (
                spec.dataset.key,
                spec.dataset.title,
                spec.dataset.source_date,
                "源 PRJ 声明 EPSG:21481，但坐标值与项目 SRID 900914 一致；原坐标未重投影。",
            ),
        )
        cur.execute(
            sql.SQL("SELECT COUNT(*) AS count, GeometryType(geom) AS geometry_type FROM {}.{} GROUP BY GeometryType(geom)").format(
                sql.Identifier("source_gis"), sql.Identifier(stage_name)
            )
        )
        geometry_rows = cur.fetchall()
        count = sum(int(row["count"]) for row in geometry_rows)
        geometry_type = ",".join(sorted({row["geometry_type"] for row in geometry_rows}))
        cur.execute(
            """
            INSERT INTO source_gis.layers (
                layer_key, dataset_key, display_name, source_file, feature_kind,
                geometry_type, feature_count, imported_at
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, NOW())
            ON CONFLICT (layer_key) DO UPDATE SET
                dataset_key = EXCLUDED.dataset_key,
                display_name = EXCLUDED.display_name,
                source_file = EXCLUDED.source_file,
                feature_kind = EXCLUDED.feature_kind,
                geometry_type = EXCLUDED.geometry_type,
                feature_count = EXCLUDED.feature_count,
                imported_at = NOW()
            """,
            (
                spec.key,
                spec.dataset.key,
                spec.display_name,
                spec.source_file.name,
                spec.feature_kind,
                geometry_type,
                count,
            ),
        )
        cur.execute("DELETE FROM source_gis.features WHERE layer_key = %s", (spec.key,))

        object_id = _first_column(columns, ("OBJECTID", "ObjectID", "objectid"))
        keyword = _first_column(columns, ("GJZ", "关键字"))
        asset_id = _first_column(columns, spec.asset_id_candidates)
        diameter = _first_column(columns, ("PIPE_D", "DIAMETER", "直径"))
        from_node = _first_column(columns, ("FNODE", "起点编号", "起点ID"))
        to_node = _first_column(columns, ("TNODE", "终点编号", "终点ID"))
        insert_query = sql.SQL(
            """
            INSERT INTO source_gis.features (
                layer_key, source_fid, source_object_id, source_keyword,
                source_asset_id, source_diameter, from_node_id, to_node_id,
                attributes, geom
            )
            SELECT %s, source_fid, {}, {}, {}, {}, {}, {},
                   to_jsonb(stage) - 'source_fid' - 'geom',
                   ST_Force2D(ST_SetSRID(geom, 900914))
            FROM {}.{} AS stage
            """
        ).format(
            _column_as_text(object_id),
            _column_as_text(keyword),
            _column_as_text(asset_id),
            _column_as_float(diameter),
            _column_as_text(from_node),
            _column_as_text(to_node),
            sql.Identifier("source_gis"),
            sql.Identifier(stage_name),
        )
        cur.execute(insert_query, (spec.key,))
        cur.execute(
            sql.SQL("DROP TABLE {}.{}").format(
                sql.Identifier("source_gis"), sql.Identifier(stage_name)
            )
        )
    return count


def _rebuild_element_links(database_name: str) -> dict[str, int]:
    with project_connection(database_name) as conn, conn.transaction(), conn.cursor() as cur:
        cur.execute(
            """
            UPDATE source_gis.features AS april
            SET source_diameter = august.source_diameter
            FROM source_gis.features AS august
            WHERE april.layer_key = 'fx_20200423_pipes'
              AND august.layer_key = 'fx_20200811_dn100_pipes'
              AND april.source_object_id = august.source_object_id
              AND ST_Equals(april.geom, august.geom)
              AND april.source_diameter IS NULL
            """
        )
        cur.execute("DELETE FROM source_gis.element_links")
        cur.execute(
            """
            INSERT INTO source_gis.element_links (
                element_type, element_id, layer_key, source_fid,
                match_method, distance_m
            )
            SELECT 'pipes', pipe.id, layer.layer_key, match.source_fid,
                   'node_pair', 0
            FROM gis.pipes AS pipe
            CROSS JOIN source_gis.layers AS layer
            CROSS JOIN LATERAL (
                SELECT feature.source_fid
                FROM source_gis.features AS feature
                WHERE feature.layer_key = layer.layer_key
                  AND (
                    (feature.from_node_id = pipe.start_node_id AND feature.to_node_id = pipe.end_node_id)
                    OR
                    (feature.from_node_id = pipe.end_node_id AND feature.to_node_id = pipe.start_node_id)
                  )
                ORDER BY feature.source_fid
                LIMIT 1
            ) AS match
            WHERE layer.feature_kind = 'pipe'
            """
        )
        node_pair_count = cur.rowcount

        cur.execute(
            """
            INSERT INTO source_gis.element_links (
                element_type, element_id, layer_key, source_fid,
                match_method, distance_m
            )
            SELECT 'pipes', model.element_id, model.layer_key, match.source_fid,
                   'geometry', match.distance_m
            FROM (
                SELECT pipe.id AS element_id, layer.layer_key,
                       pipe.diameter,
                       ST_Transform(pipe.geom, 900914) AS geom
                FROM gis.pipes AS pipe
                CROSS JOIN source_gis.layers AS layer
                WHERE layer.feature_kind = 'pipe'
                  AND NOT EXISTS (
                      SELECT 1 FROM source_gis.element_links AS existing
                      WHERE existing.element_type = 'pipes'
                        AND existing.element_id = pipe.id
                        AND existing.layer_key = layer.layer_key
                  )
                  AND ST_Length(ST_Transform(pipe.geom, 900914)) >= 0.1
            ) AS model
            CROSS JOIN LATERAL (
                SELECT candidate.source_fid, candidate.distance_m
                FROM (
                    SELECT feature.source_fid,
                           ST_HausdorffDistance(feature.geom, model.geom) AS distance_m
                    FROM source_gis.features AS feature
                    WHERE feature.layer_key = model.layer_key
                      AND feature.geom && ST_Expand(model.geom, 1.0)
                      AND (
                          feature.source_diameter IS NULL
                          OR model.diameter IS NULL
                          OR ABS(feature.source_diameter - model.diameter)
                             <= GREATEST(1.0, model.diameter * 0.02)
                      )
                    ORDER BY feature.geom <-> model.geom
                    LIMIT 8
                ) AS candidate
                ORDER BY candidate.distance_m, candidate.source_fid
                LIMIT 1
            ) AS match
            WHERE match.distance_m <= 1.0
            ON CONFLICT (element_type, element_id, layer_key) DO NOTHING
            """
        )
        geometry_count = cur.rowcount

        cur.execute(
            """
            WITH point_matches AS (
                SELECT feature.layer_key, feature.source_fid,
                       nearest.node_id AS element_id,
                       nearest.distance_m,
                       ROW_NUMBER() OVER (
                           PARTITION BY nearest.node_id, feature.layer_key
                           ORDER BY nearest.distance_m, feature.source_fid
                       ) AS match_rank
                FROM source_gis.features AS feature
                JOIN source_gis.layers AS layer
                  ON layer.layer_key = feature.layer_key
                CROSS JOIN LATERAL (
                    SELECT node.node_id,
                           ST_Distance(node.geom, feature.geom) AS distance_m
                    FROM gis.node_geometries AS node
                    ORDER BY node.geom <-> feature.geom
                    LIMIT 1
                ) AS nearest
                JOIN gis.junctions AS junction
                  ON junction.id = nearest.node_id
                WHERE layer.feature_kind <> 'pipe'
                  AND nearest.distance_m <= 0.1
            )
            INSERT INTO source_gis.element_links (
                element_type, element_id, layer_key, source_fid,
                match_method, distance_m
            )
            SELECT 'junctions', element_id, layer_key, source_fid,
                   'geometry', distance_m
            FROM point_matches
            WHERE match_rank = 1
            """
        )
        junction_count = cur.rowcount

        cur.execute(
            """
            INSERT INTO source_gis.element_links (
                element_type, element_id, layer_key, source_fid,
                match_method, distance_m
            )
            SELECT 'valves', valve.id, layer.layer_key, match.source_fid,
                   'geometry', match.distance_m
            FROM gis.valves AS valve
            CROSS JOIN source_gis.layers AS layer
            CROSS JOIN LATERAL (
                SELECT feature.source_fid,
                       ST_Distance(feature.geom, ST_Transform(valve.geom, 900914)) AS distance_m
                FROM source_gis.features AS feature
                WHERE feature.layer_key = layer.layer_key
                ORDER BY feature.geom <-> ST_Transform(valve.geom, 900914)
                LIMIT 1
            ) AS match
            WHERE layer.feature_kind = 'valve'
              AND match.distance_m <= 2.0
            """
        )
        valve_count = cur.rowcount
    return {
        "node_pair": node_pair_count,
        "pipe_geometry": geometry_count,
        "junction_geometry": junction_count,
        "valve_geometry": valve_count,
    }


def import_sources(database_name: str, specs: tuple[LayerSpec, ...]) -> None:
    missing = [str(spec.source_file) for spec in specs if not spec.source_file.is_file()]
    if missing:
        raise FileNotFoundError("Missing GIS source files: " + ", ".join(missing))
    with project_connection(database_name) as conn:
        conn.execute(SCHEMA_SQL.read_text(encoding="utf-8"))

    staged_layers: list[tuple[LayerSpec, str]] = []
    for spec in specs:
        print(f"Loading {spec.display_name} ({spec.key})")
        staged_layers.append((spec, _load_staging_table(spec, database_name)))

    total = 0
    with project_transaction(database_name) as conn:
        conn.execute(
            "TRUNCATE TABLE source_gis.element_links, source_gis.features"
        )
        for spec, stage_name in staged_layers:
            count = _store_layer(database_name, spec, stage_name)
            total += count
            print(f"  stored {count:,} features")
        links = _rebuild_element_links(database_name)
    print(f"Imported {total:,} features into {database_name}")
    print(
        "Element links: "
        f"{links['node_pair']:,} node-pair, "
        f"{links['pipe_geometry']:,} pipe geometry, "
        f"{links['junction_geometry']:,} junction geometry, "
        f"{links['valve_geometry']:,} valve geometry"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--database", default="tjwater_v2")
    parser.add_argument("--april-root", type=Path, default=DEFAULT_APRIL_ROOT)
    parser.add_argument("--august-root", type=Path, default=DEFAULT_AUGUST_ROOT)
    args = parser.parse_args()
    import_sources(args.database, layer_specs(args.april_root, args.august_root))


if __name__ == "__main__":
    main()