from __future__ import annotations

import json
import sqlite3
from array import array
from collections import defaultdict
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

import numpy as np

from geogst.core.geology.profiles.geoprofiles import GeoProfiles
from geogst.core.geology.profiles.profiletraces import (
    ZTrace,
    PointTrace,
    PlaneTrace,
    FocalMechanismTrace,
)
from geogst.core.utils.arrays import CategorizedArrays
from geogst.core.geometries.shapes.points import Point  # usato per point_breaks in ZTrace


def _connect(db_path: str) -> sqlite3.Connection:
    con = sqlite3.connect(db_path)
    con.row_factory = sqlite3.Row
    return con


def _table_columns(con: sqlite3.Connection, table: str) -> List[str]:
    cur = con.execute(f"PRAGMA table_info({table});")
    return [r["name"] for r in cur.fetchall()]


def _safe_json_loads(s: Optional[str]) -> Dict[str, Any]:
    if not s:
        return {}
    try:
        return json.loads(s)
    except Exception:
        return {}


def _as_array_d(vals: Sequence[float]) -> array:
    return array("d", [float(v) for v in vals])


class GeoProfileReader:
    """
    Read exported geoprofiles from the GeoProfiler gpkg/sqlite schema.
    Returns a GeoProfiles object ready to be plotted / re-exported.
    """

    @staticmethod
    def from_db(
        db_path: str,
        *,
        result_set_id: Optional[int] = None,
        profile_ids: Optional[Sequence[int]] = None,
    ) -> GeoProfiles:

        con = _connect(db_path)
        try:
            # ---------------------------
            # 1) pick profiles to load
            # ---------------------------
            if profile_ids is not None and len(profile_ids) > 0:

                prof_rows = con.execute(
                    # nosec B608 - interpolati solo i placeholder '?'; i valori
                    # sono bind parameter passati come secondo argomento
                    f"SELECT * FROM gp_profiles WHERE profile_id IN ({','.join(['?']*len(profile_ids))}) ORDER BY profile_id",   # nosec B608
                    list(profile_ids),
                ).fetchall()

            else:

                if result_set_id is None:
                    # default: newest result set
                    rs = con.execute(
                        "SELECT result_set_id FROM gp_result_sets ORDER BY result_set_id DESC LIMIT 1"
                    ).fetchone()
                    if not rs:
                        raise Exception("No gp_result_sets found in DB.")
                    result_set_id = int(rs["result_set_id"])

                prof_rows = con.execute(
                    "SELECT * FROM gp_profiles WHERE result_set_id = ? ORDER BY profile_id",
                    (int(result_set_id),),
                ).fetchall()

            if not prof_rows:
                raise Exception("No profiles found for given selection.")

            # CRS: take first profile wkt_crs (fallback empty)
            wkt_crs = prof_rows[0]["wkt_crs"] or ""

            # ---------------------------
            # 2) topo profiles -> ZTrace
            # ---------------------------

            topo_profiles: List[ZTrace] = []
            profile_id_list: List[int] = []

            for pr in prof_rows:

                pid = int(pr["profile_id"])
                profile_id_list.append(pid)

                samples = con.execute(
                    "SELECT idx, s, z, x, y FROM gp_profile_samples WHERE profile_id = ? ORDER BY idx",
                    (pid,),
                ).fetchall()

                if not samples:
                    # allow empty topo: create trivial ZTrace
                    s_arr = _as_array_d([0.0, float(pr["s_max"] or 1.0)])
                    z_arr = _as_array_d([0.0, 0.0])
                    s_breaks = _as_array_d([s_arr[0], s_arr[-1]])
                    point_breaks = [Point(0.0, 0.0), Point(1.0, 0.0)]
                    topo_profiles.append(ZTrace(s_arr, z_arr, s_breaks, point_breaks, wkt_crs, rec_id=pid))
                    continue

                s_vals = [float(r["s"]) for r in samples]
                z_vals = [float(r["z"]) for r in samples]
                s_arr = _as_array_d(s_vals)
                z_arr = _as_array_d(z_vals)

                # breaks: simplest (start/end)
                s_breaks = _as_array_d([s_arr[0], s_arr[-1]])

                # points: use first/last XY if available
                x0 = samples[0]["x"]
                y0 = samples[0]["y"]
                x1 = samples[-1]["x"]
                y1 = samples[-1]["y"]
                if x0 is None or y0 is None or x1 is None or y1 is None:
                    point_breaks = [Point(0.0, 0.0), Point(1.0, 0.0)]
                else:
                    point_breaks = [Point(float(x0), float(y0)), Point(float(x1), float(y1))]

                topo_profiles.append(ZTrace(s_arr, z_arr, s_breaks, point_breaks, wkt_crs, rec_id=pid))

            gps = GeoProfiles(topo_profiles=topo_profiles, wkt_crs=wkt_crs)

            # Helper: map profile_id -> index inside gps
            pid_to_ndx = {pid: i for i, pid in enumerate(profile_id_list)}

            # ----------------------------------
            # 3) projected points (gp_projected_points)
            # ----------------------------------

            pts_by_profile: Dict[int, Dict[Any, PointTrace]] = defaultdict(dict)

            cols_pts = set(_table_columns(con, "gp_projected_points"))
            # we expect: rec_id, profile_id, label, s, z, dist_to_profile, src_fid ...
            rows_pts = con.execute(
                # nosec B608 - interpolati solo i placeholder '?'; i valori
                # sono bind parameter passati come secondo argomento
                f"SELECT * FROM gp_projected_points WHERE profile_id IN ({','.join(['?']*len(profile_id_list))})",   # nosec B608
                profile_id_list,
            ).fetchall()

            for r in rows_pts:
                pid = int(r["profile_id"])
                # robust key: prefer src_fid if present, else internal rec_id
                key = r["src_fid"] if ("src_fid" in cols_pts and r["src_fid"] is not None) else int(r["rec_id"])
                label = (r["label"] or "")  # keep label as string in key if you prefer
                # If you want label grouping, swap key=label. For robustness keep numeric key.
                pt = PointTrace(
                    s=float(r["s"]),
                    z=float(r["z"]),
                    dist=float(r["dist_to_profile"]) if r["dist_to_profile"] is not None else 0.0,
                )
                pts_by_profile[pid][key] = pt

            # store in gps._points_projections["points"] as list per profile index
            points_list_per_profile: List[Dict[Any, PointTrace]] = [dict() for _ in profile_id_list]
            for pid, d in pts_by_profile.items():
                points_list_per_profile[pid_to_ndx[pid]] = d
            gps._points_projections["points"] = points_list_per_profile

            # ----------------------------------
            # 4) projected attitudes (gp_projected_attitudes)
            #    supports both old schema and PlaneTrace schema
            # ----------------------------------

            atts_by_profile: Dict[int, Dict[Any, PlaneTrace]] = defaultdict(dict)
            cols_att = set(_table_columns(con, "gp_projected_attitudes"))

            rows_att = con.execute(
                # nosec B608 - interpolati solo i placeholder '?'; i valori
                # sono bind parameter passati come secondo argomento
                f"SELECT * FROM gp_projected_attitudes WHERE profile_id IN ({','.join(['?']*len(profile_id_list))})",   # nosec B608
                profile_id_list,
            ).fetchall()

            has_plane_schema = {"slope_degr", "down_sense", "src_dip_dir", "src_dip_ang"}.issubset(cols_att)

            for r in rows_att:
                pid = int(r["profile_id"])
                key = r["src_fid"] if ("src_fid" in cols_att and r["src_fid"] is not None) else int(r["rec_id"])

                if has_plane_schema:
                    pl = PlaneTrace(
                        s=float(r["s"]),
                        z=float(r["z"]),
                        slope_degr=float(r["slope_degr"]),
                        down_sense=str(r["down_sense"]),
                        dist=float(r["dist_to_profile"]) if r["dist_to_profile"] is not None else 0.0,
                        src_dip_dir=float(r["src_dip_dir"]),
                        src_dip_ang=float(r["src_dip_ang"]),
                    )
                else:
                    # fallback (old schema): map azimuth/dip into src_* and slope, down_sense unknown
                    az = float(r["azimuth"]) if r["azimuth"] is not None else 0.0
                    dip = float(r["dip"]) if r["dip"] is not None else 0.0
                    pl = PlaneTrace(
                        s=float(r["s"]),
                        z=float(r["z"]),
                        slope_degr=dip,
                        down_sense="right",
                        dist=float(r["dist_to_profile"]) if r["dist_to_profile"] is not None else 0.0,
                        src_dip_dir=az,
                        src_dip_ang=dip,
                    )

                atts_by_profile[pid][key] = pl

            atts_list_per_profile: List[Dict[Any, PlaneTrace]] = [dict() for _ in profile_id_list]
            for pid, d in atts_by_profile.items():
                atts_list_per_profile[pid_to_ndx[pid]] = d
            gps._points_projections["attitudes"] = atts_list_per_profile

            # ----------------------------------
            # 4b) projected focal mechanisms (gp_projected_focal_mechanisms)
            #     the table is newer than the rest of the schema, so databases
            #     written before it existed are read back without mechanisms
            #     rather than failing
            # ----------------------------------

            fmecs_list_per_profile: List[Dict[Any, FocalMechanismTrace]] = [dict() for _ in profile_id_list]

            if con.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='gp_projected_focal_mechanisms';"
            ).fetchone() is not None:

                fmecs_by_profile: Dict[int, Dict[Any, FocalMechanismTrace]] = defaultdict(dict)
                cols_fmec = set(_table_columns(con, "gp_projected_focal_mechanisms"))

                rows_fmec = con.execute(
                    # nosec B608 - interpolati solo i placeholder '?'; i valori
                    # sono bind parameter passati come secondo argomento
                    f"SELECT * FROM gp_projected_focal_mechanisms WHERE profile_id IN ({','.join(['?']*len(profile_id_list))})",   # nosec B608
                    profile_id_list,
                ).fetchall()

                for r in rows_fmec:
                    pid = int(r["profile_id"])
                    key = r["src_fid"] if ("src_fid" in cols_fmec and r["src_fid"] is not None) else int(r["rec_id"])

                    fmecs_by_profile[pid][key] = FocalMechanismTrace(
                        s=float(r["s"]),
                        z=float(r["z"]),
                        dist=float(r["dist_to_profile"]) if r["dist_to_profile"] is not None else 0.0,
                        strike=float(r["strike"]),
                        dip=float(r["dip"]),
                        rake=float(r["rake"]),
                        profile_azimuth=float(r["profile_azimuth"]),
                    )

                for pid, d in fmecs_by_profile.items():
                    fmecs_list_per_profile[pid_to_ndx[pid]] = d

            gps._points_projections["focal_mechanisms"] = fmecs_list_per_profile

            # ----------------------------------
            # 5) line intersections -> CategorizedArrays
            # ----------------------------------
            if "gp_intersected_lines" in (con.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='gp_intersected_lines';"
            ).fetchone() or {}):
                rows_li = con.execute(
                    # nosec B608 - interpolati solo i placeholder '?'; i valori
                    # sono bind parameter passati come secondo argomento
                    f"SELECT * FROM gp_intersected_lines WHERE profile_id IN ({','.join(['?']*len(profile_id_list))}) ORDER BY rec_id",   # nosec B608
                    profile_id_list,
                ).fetchall()

                li_group: Dict[int, Dict[str, List[array]]] = defaultdict(lambda: defaultdict(list))
                for r in rows_li:
                    pid = int(r["profile_id"])
                    cat = (r["feat_category"] or "").strip()
                    if not cat:
                        cat = "undefined"
                    # store as array('d',[s]) (consistent with your current in-memory format)
                    li_group[pid][cat].append(_as_array_d([float(r["s"])]))

                lines_per_profile: List[List[List[CategorizedArrays]]] = []
                for pid in profile_id_list:
                    cats = []
                    for cat, arrs in li_group[pid].items():
                        cats.append(CategorizedArrays(cat, arrs))
                    # 1 segment only (wrap twice): [[[...]]]
                    lines_per_profile.append([cats])
                gps._lines_intersections = lines_per_profile

            # ----------------------------------
            # 6) polygon intersections -> CategorizedArrays with [s_from, s_to]
            # ----------------------------------
            rows_pi = con.execute(
                # nosec B608 - interpolati solo i placeholder '?'; i valori
                # sono bind parameter passati come secondo argomento
                f"SELECT * FROM gp_intersected_polygons WHERE profile_id IN ({','.join(['?']*len(profile_id_list))}) ORDER BY rec_id",   # nosec B608
                profile_id_list,
            ).fetchall()

            pi_group: Dict[int, Dict[str, List[array]]] = defaultdict(lambda: defaultdict(list))
            for r in rows_pi:
                pid = int(r["profile_id"])
                unit = str(r["unit_name"])
                s_from = float(r["s_from"])
                s_to = float(r["s_to"])
                pi_group[pid][unit].append(_as_array_d([s_from, s_to]))

            polygons_per_profile: List[List[List[CategorizedArrays]]] = []
            for pid in profile_id_list:
                cats = []
                for unit, arrs in pi_group[pid].items():
                    cats.append(CategorizedArrays(unit, arrs))
                polygons_per_profile.append([cats])
            gps._polygons_intersections = polygons_per_profile

            return gps

        finally:
            con.close()
