from __future__ import annotations
import copy
from typing import Any, TYPE_CHECKING

import numpy as np
from qgis.PyQt.QtCore import Qt
from qgis.PyQt.QtWidgets import QTableWidget, QTableWidgetItem

if TYPE_CHECKING:
    from omrat import OMRAT

_USER_ROLE = getattr(Qt, 'UserRole', None) or Qt.ItemDataRole.UserRole


def dict_ndarray_to_list(data: dict[str, dict[str, np.ndarray]]) -> dict[str, dict[str, list]]:
    """Convert all np.ndarray values in a nested dict to lists."""
    result: dict[str, dict[str, list]] = {}
    for key, inner in data.items():
        result[key] = {}
        for subkey, value in inner.items():
            if isinstance(value, np.ndarray):
                result[key][subkey] = value.tolist()
            else:
                result[key][subkey] = value
    return result


def dict_list_to_ndarray(data: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
    """Convert all list values in a nested dict to np.ndarray."""
    result: dict[str, dict[str, Any]] = {}
    for key, inner in data.items():
        result[key] = {}
        for subkey, value in inner.items():
            if isinstance(value, list):
                result[key][subkey] = np.array(value)
            else:
                result[key][subkey] = value
    return result


class GatherData:
    def __init__(self, parent: OMRAT) -> None:
        self.p = parent
        self.data = {}

    def get_segment_tbl(self):
        """Extends the segment_data in self.data, must be called after it is created."""
        for j, col in enumerate(['Segment_Id', 'Route_Id', 'Leg_name', 'Start_Point', 'End_Point', 'Width']):
            for i, key in enumerate(self.data['segment_data'].keys()):
                item = self.p.main_widget.twRouteList.item(i, j)
                if item is None:
                    return
                value = item.text()
                self.data["segment_data"][key][col] = copy.deepcopy(value)

    def get_all_for_save(self) -> dict[str, Any]:
        self.data['pc'] = copy.deepcopy(self.p.causation_f.data)
        self.data['drift'] = copy.deepcopy(self.p.drift_values)
        self.p.distributions.change_dist_segment(self.p.distributions.last_id)  # Saves the current settings on the leg
        self.data['traffic_data'] = copy.deepcopy(self.p.traffic_data)
        self.data['segment_data'] = copy.deepcopy(self.p.segment_data)
        for key, item in self.data['segment_data'].items():
            self.data['segment_data'][key]['dist1'] = list(item.get('dist1', np.array([])))
            self.data['segment_data'][key]['dist2'] = list(item.get('dist2', np.array([])))
        self.get_segment_tbl()

        self.data['depths'] = []
        self.data['objects'] = []
        self.copy_depths_and_objects()
        # Persist the ship category scheme (types and size intervals) used to build traffic tables
        self.data['ship_categories'] = self.get_ship_categories_for_save()
        # Persist the oil-spill consequence block.  Pulled from the
        # ``Consequence`` handler if the dialogs have been opened, otherwise
        # filled with category-aware defaults so the .omrat round-trips.
        self.data['consequence'] = self.get_consequence_for_save()
        # Persist the junction transition matrices.  Empty dict is fine
        # — the load path falls back to geometric defaults if missing.
        self.data['junctions'] = self.get_junctions_for_save()
        # Persist the project-level traffic-scaling state (global spinbox
        # + per-type "follow global" booleans).  The per-cell ``Scaling
        # (%)`` matrices ride inside ``traffic_data`` already.
        from omrat_utils.handle_traffic import _default_traffic_scaling
        scaling = getattr(self.p, 'traffic_scaling', None)
        if not isinstance(scaling, dict):
            scaling = _default_traffic_scaling()
        self.data['traffic_scaling'] = {
            'global_percent': float(scaling.get('global_percent', 100.0)),
            'follow_global': [bool(x) for x in scaling.get('follow_global', [])],
        }
        # Write-once snapshot of imported segment endpoints.  See
        # ``omrat_utils/audit_report.py`` for how this is consumed.
        self.data['segments_imported'] = self.get_segments_imported_for_save()
        return self.data

    def get_segments_imported_for_save(self) -> dict[str, dict[str, str]]:
        """Snapshot endpoints captured at import / first-digitize time.

        Seeds from the current ``segment_data`` for any leg that has no
        recorded import — so legacy projects (saved before the audit
        feature) get a baseline on first save instead of staying empty.
        Once stamped, an entry is never overwritten.
        """
        existing = dict(getattr(self.p, 'segments_imported', {}) or {})
        for sid, seg in (self.p.segment_data or {}).items():
            if sid in existing:
                continue
            sp = seg.get('Start_Point') if isinstance(seg, dict) else None
            ep = seg.get('End_Point') if isinstance(seg, dict) else None
            if sp and ep:
                existing[str(sid)] = {
                    'Start_Point': str(sp),
                    'End_Point': str(ep),
                }
        # Mirror back so subsequent saves see the seeded baseline.
        self.p.segments_imported = existing
        return existing

    def get_junctions_for_save(self) -> dict[str, Any]:
        """Snapshot the junction registry for the .omrat file."""
        handler = getattr(self.p, 'junctions', None)
        if handler is None:
            return {}
        try:
            return handler.to_dict()
        except Exception:  # nosec B110 B112
            return {}

    def get_consequence_for_save(self) -> dict[str, Any]:
        """Snapshot the four consequence matrices for the .omrat file.

        Reads the live state held by ``self.p.consequence`` (created in
        ``OMRAT.__init__``) and converts it to plain Python types.  The
        consequence handler keeps every matrix in shape with the current
        ship_categories, so we don't have to reshape here.
        """
        from omrat_utils.consequence_defaults import (
            default_oil_onboard,
            default_spill_probability,
            default_spill_fraction,
            default_catastrophe_levels,
        )

        ship_cats = self.data.get('ship_categories', {}) or {}
        types = list(ship_cats.get('types', []))
        intervals = list(ship_cats.get('length_intervals', []))

        consequence_handler = getattr(self.p, 'consequence', None)
        if consequence_handler is not None:
            return {
                'oil_onboard': consequence_handler.oil_onboard,
                'spill_probability': consequence_handler.spill_probability,
                'spill_fraction': consequence_handler.spill_fraction,
                'catastrophe_levels': [
                    dict(level) for level in consequence_handler.catastrophe_levels
                ],
            }
        # Handler not constructed (e.g. headless tests).  Fall back to
        # defaults so the .omrat still validates.
        return {
            'oil_onboard': default_oil_onboard(types, intervals),
            'spill_probability': default_spill_probability(),
            'spill_fraction': default_spill_fraction(),
            'catastrophe_levels': default_catastrophe_levels(),
        }

    def copy_depths_and_objects(self):
        # Read rows from QTableWidget instead of iterating the widget
        depth_rows = self.obtain_table_data(self.p.main_widget.twDepthList)
        object_rows = self.obtain_table_data(self.p.main_widget.twObjectList)
        for row in depth_rows:
            if len(row) < 3:
                # Skip incomplete rows
                continue
            id_ = row[0]
            depths = row[1]
            polygon = row[2]
            # Support values like "48-...": take the first part as depth
            depth_val = depths.split('-')[0] if isinstance(depths, str) else str(depths)
            # Store as list [id, depth, polygon] per schema
            self.data['depths'].append([str(id_), str(depth_val), str(polygon)])
        for row in object_rows:
            if len(row) < 3:
                continue
            id_ = row[0]
            height = row[1]
            polygon = row[2]
            # Store as list [id, height, polygon] per schema
            self.data['objects'].append([str(id_), str(height), str(polygon)])

    def get_ship_categories_for_save(self) -> dict[str, Any]:
        """Capture the ship types (rows) and size intervals (columns) used for traffic tables.

        Returns a dict with keys:
        - types: list[str] of ship type labels (row headers from cvTypes)
        - length_intervals: list[dict] with keys {'min': float, 'max': float, 'label': str}
        - selection_mode: optional str indicating UI mode (e.g., 'simple_ais' or 'manual')
        """
        result: dict[str, Any] = {
            'types': [],
            'length_intervals': [],
        }
        try:
            # Extract ship types from Ship Categories widget
            scw = self.p.ship_cat.scw  # ShipCategoriesWidget
            types: list[str] = []
            if hasattr(scw, 'cvTypes') and scw.cvTypes is not None:
                rows = scw.cvTypes.rowCount()
                for i in range(rows):
                    it = scw.cvTypes.item(i, 0)
                    text = it.text() if it is not None else ''
                    if text:
                        types.append(text)
            result['types'] = types

            # Extract size intervals from twLengths (Min/Max columns)
            intervals: list[dict[str, Any]] = []
            if hasattr(scw, 'twLengths') and scw.twLengths is not None:
                rows = scw.twLengths.rowCount()
                for i in range(rows):
                    it_min = scw.twLengths.item(i, 0)
                    it_max = scw.twLengths.item(i, 1)
                    smin = it_min.text() if it_min is not None else ''
                    smax = it_max.text() if it_max is not None else ''
                    if smin == '' and smax == '':
                        continue
                    try:
                        vmin = float(smin)
                    except Exception:  # nosec B110 B112
                        vmin = smin
                    try:
                        vmax = float(smax)
                    except Exception:  # nosec B110 B112
                        vmax = smax
                    label = f"{smin} - {smax}".strip()
                    intervals.append({'min': vmin, 'max': vmax, 'label': label})
            result['length_intervals'] = intervals

            # Selection mode (optional)
            mode = None
            if hasattr(scw, 'radioButton') and hasattr(scw, 'radioButton_2'):
                try:
                    if scw.radioButton.isChecked():
                        mode = 'simple_ais'
                    elif scw.radioButton_2.isChecked():
                        mode = 'manual'
                except Exception:  # nosec B110 B112
                    mode = None
            if mode is not None:
                result['selection_mode'] = mode
        except Exception:  # nosec B110 B112
            # If widget not available, leave defaults (empty) to avoid breaking save
            pass
        return result

    def obtain_table_data(self, tbl) -> list:
        """Obtain data from a table.

        Cells that store a truncated display (e.g. WKT polygon summaries) keep
        the full value in Qt.UserRole; prefer that over the display text so that
        the full WKT is saved rather than the short label.
        """
        tbl_data = []
        rows = tbl.rowCount()
        cols = tbl.columnCount()
        for row in range(rows):
            line = []
            for col in range(cols):
                value = tbl.item(row, col)
                if value is not None:
                    get_data = getattr(value, 'data', None)
                    user_data = get_data(_USER_ROLE) if get_data is not None else None
                    line.append(user_data if user_data is not None else value.text())
            tbl_data.append(line)
        return tbl_data

    def normalize_depths_for_ui(self, depths: list) -> list[list[str]]:
        """Undo copy_depths_and_objects for depths: dicts -> [id, depth, polygon].

        Accepts either a list of dicts (with keys id, depth, polygon) or
        an already UI-ready list of rows, and returns list of [id, depth, polygon].
        """
        rows: list[list[str]] = []
        if not isinstance(depths, list):
            return rows
        for item in depths:
            # Assume it's already a sequence like [id, depth, polygon]
            try:
                rows.append([str(item[0]), str(item[1]), str(item[2])])
            except Exception:  # nosec B110 B112
                rows.append([str(item), '', ''])
        return rows

    def normalize_objects_for_ui(self, objects: list) -> list[list[str]]:
        """Undo copy_depths_and_objects for objects: dicts -> [id, height, polygon].

        Accepts either a list of dicts (with keys id, height/heights, polygon) or
        an already UI-ready list of rows, and returns list of [id, height, polygon].
        """
        rows: list[list[str]] = []
        if not isinstance(objects, list):
            return rows
        for item in objects:
            try:
                rows.append([str(item[0]), str(item[1]), str(item[2])])
            except Exception:  # nosec B110 B112
                rows.append([str(item), '', ''])
        return rows

    def _load_traffic_scaling(self, data: dict) -> None:
        from omrat_utils.handle_traffic import _default_traffic_scaling
        scaling_block = data.get('traffic_scaling')
        if isinstance(scaling_block, dict):
            try:
                gp = float(scaling_block.get('global_percent', 100.0))
            except (TypeError, ValueError):
                gp = 100.0
            flags = [bool(x) for x in scaling_block.get('follow_global', [])]
            self.p.traffic_scaling = {'global_percent': gp, 'follow_global': flags}
        else:
            self.p.traffic_scaling = _default_traffic_scaling()
        sb = getattr(self.p.main_widget, 'dsbGlobalTrafficScaling', None)
        if sb is not None:
            try:
                sb.blockSignals(True)
                sb.setValue(float(self.p.traffic_scaling['global_percent']))
            finally:
                sb.blockSignals(False)
        traffic = getattr(self.p, 'traffic', None)
        if traffic is not None and hasattr(traffic, 'ensure_scaling_present'):
            traffic.ensure_scaling_present()
        if traffic is not None and hasattr(traffic, 'populate_scaling_checkboxes'):
            traffic.populate_scaling_checkboxes()

    def _populate_causation_and_traffic(self, data: dict) -> None:
        if 'pc' in data and isinstance(data['pc'], dict):
            self.p.causation_f.data.update(data['pc'])
        self.p.traffic_data = data['traffic_data']
        if hasattr(self.p, 'traffic') and self.p.traffic is not None:
            self.p.traffic.traffic_data = self.p.traffic_data
        self._load_traffic_scaling(data)

    def _populate_segment_data(self, data: dict) -> None:
        self.p.segment_data = data['segment_data']
        self.p.drift_values = data['drift']
        self.p.drift_settings.drift_values = data['drift']
        imported_raw = data.get('segments_imported') or {}
        self.p.segments_imported = {
            str(sid): {'Start_Point': str(rec.get('Start_Point', '')),
                       'End_Point': str(rec.get('End_Point', ''))}
            for sid, rec in imported_raw.items() if isinstance(rec, dict)
        }
        for key, item in data['segment_data'].items():
            self.p.segment_data[key]['dist1'] = np.array(item.get('dist1', []))
            self.p.segment_data[key]['dist2'] = np.array(item.get('dist2', []))
        self.populate_segment_tbl(data['segment_data'], self.p.main_widget.twRouteList)
        self.populate_cbTrafficSelectSeg()
        self.p.main_widget.leNormMean1_1.setText('')
        self.p.distributions.change_dist_segment(self.p.distributions.last_id)

    def _populate_objects_and_canvas(self, data: dict, depth_rows: list, object_rows: list) -> None:
        self.p.object.area_type = 'depth'
        for i, dep in enumerate(depth_rows):
            depth_value = dep[1] if len(dep) > 1 else dep[0]
            self.p.object.load_area(f'Depth - {depth_value}m', dep[2], row=i,
                                    value=depth_value, value_field='Depth', defer_style=True)
        self.p.object._apply_depth_graduated_style()
        self.p.object.area_type = 'object'
        for j, obj in enumerate(object_rows):
            height_value = obj[1] if len(obj) > 1 else obj[0]
            self.p.object.load_area(f'Structure - {height_value}m', obj[2], row=j,
                                    value=height_value, value_field='Height')
        self.p.load_lines(data)

    def _populate_consequence_and_junctions(self, data: dict) -> None:
        consequence_handler = getattr(self.p, 'consequence', None)
        if consequence_handler is not None:
            consequence_handler.load_from_dict(
                data.get('consequence') or {}, data.get('ship_categories') or {},
            )
        junctions_handler = getattr(self.p, 'junctions', None)
        if junctions_handler is not None:
            junctions_handler.load_from_dict(
                data.get('junctions') or {}, data.get('segment_data') or {},
            )

    def populate(self, data):
        if 'ship_categories' in data and data['ship_categories']:
            self.populate_ship_categories(data['ship_categories'])
        self._populate_causation_and_traffic(data)
        self._populate_segment_data(data)
        depth_rows = self.normalize_depths_for_ui(data['depths'])
        object_rows = self.normalize_objects_for_ui(data['objects'])
        self.populate_tbl(depth_rows, self.p.main_widget.twDepthList)
        self.populate_tbl(object_rows, self.p.main_widget.twObjectList)
        self._populate_objects_and_canvas(data, depth_rows, object_rows)
        self._populate_consequence_and_junctions(data)

    def populate_ship_categories(self, ship_categories: dict[str, Any]):
        """Populate ship types and length intervals into the ship categories widget.

        Args:
            ship_categories: Dict with 'types' (list of ship type names) and
                           'length_intervals' (list of dicts with 'min', 'max', 'label')
        """
        scw = self.p.ship_cat.scw  # ShipCategoriesWidget

        # Populate ship types into cvTypes table
        types = ship_categories.get('types', [])
        if types and hasattr(scw, 'cvTypes') and scw.cvTypes is not None:
            scw.cvTypes.setRowCount(len(types))
            scw.cvTypes.setColumnCount(1)
            for i, ship_type in enumerate(types):
                item = QTableWidgetItem(str(ship_type))
                scw.cvTypes.setItem(i, 0, item)

        # Populate length intervals into twLengths table
        intervals = ship_categories.get('length_intervals', [])
        if intervals and hasattr(scw, 'twLengths') and scw.twLengths is not None:
            scw.twLengths.setRowCount(len(intervals))
            scw.twLengths.setColumnCount(2)
            for i, interval in enumerate(intervals):
                min_val = interval.get('min', '')
                max_val = interval.get('max', '')
                item_min = QTableWidgetItem(str(min_val))
                item_max = QTableWidgetItem(str(max_val))
                scw.twLengths.setItem(i, 0, item_min)
                scw.twLengths.setItem(i, 1, item_max)

        # Rebuild traffic table with new dimensions
        self.p.traffic.set_table_headings()

    def populate_cbTrafficSelectSeg(self):
        """Sets the segment names in cbTrafficSelectSeg"""
        self.p.main_widget.cbTrafficSelectSeg.clear()
        for key, seg in self.p.segment_data.items():
            label = seg.get('Leg_name') or str(key)
            self.p.main_widget.cbTrafficSelectSeg.addItem(label, str(key))
        self.p.traffic.c_seg = self.p.main_widget.cbTrafficSelectSeg.currentData() or ''

    def populate_tbl(self, data: list, tbl: QTableWidget):
        from omrat_utils.handle_object import _wkt_table_item
        _WKT_PREFIXES = ('polygon', 'multipolygon', 'point', 'linestring',
                         'multilinestring', 'geometrycollection')
        tbl.setRowCount(len(data))
        for i, line in enumerate(data):
            for j, value in enumerate(line):
                if (isinstance(value, str)
                        and value.lower().split('(')[0].strip() in _WKT_PREFIXES
                        and len(value) > 60):
                    item = _wkt_table_item(value)
                else:
                    item = QTableWidgetItem(value)
                tbl.setItem(i, j, item)

    def populate_segment_tbl(self, data: dict[str, dict[str, Any]], tbl: QTableWidget):
        tbl.setRowCount(len(data))
        for j, col in enumerate(['Segment_Id', 'Route_Id', 'Leg_name', 'Start_Point', 'End_Point', 'Width']):
            for i, key in enumerate(data.keys()):
                item = QTableWidgetItem(str(data[key][col]))
                self.p.main_widget.twRouteList.setItem(i, j, item)
