from __future__ import annotations
import json
import os
import copy
from qgis.PyQt.QtCore import QSettings
from qgis.PyQt.QtWidgets import QTableWidget, QTableWidgetItem
from typing import Any
import numpy as np

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from omrat import OMRAT


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'] = copy.deepcopy(item['dist1'].tolist())
            self.data['segment_data'][key]['dist2'] = copy.deepcopy(item['dist2'].tolist())
        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:
            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:
                        vmin = smin
                    try:
                        vmax = float(smax)
                    except Exception:
                        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:
                    mode = None
            if mode is not None:
                result['selection_mode'] = mode
        except Exception:
            # 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"""
        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:
                    line.append(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:
                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:
                rows.append([str(item), '', ''])
        return rows
    
    def populate(self, data):
        # Load ship categories first (needed for traffic table dimensions)
        if 'ship_categories' in data and data['ship_categories']:
            self.populate_ship_categories(data['ship_categories'])

        # Load causation factors -- merge over the defaults rather
        # than replacing the dict, because legacy ``.omrat`` files only
        # persisted ``p_pc`` / ``d_pc`` and would otherwise wipe the
        # collision keys (``headon``, ``overtaking``, ...) and break
        # later ``set_values`` lookups.
        if 'pc' in data and isinstance(data['pc'], dict):
            self.p.causation_f.data.update(data['pc'])

        self.p.traffic_data = data['traffic_data']
        # The Traffic instance captures its own reference to the dict at
        # __init__; update it too so cbTrafficDirectionSelect gets populated.
        if hasattr(self.p, 'traffic') and self.p.traffic is not None:
            self.p.traffic.traffic_data = self.p.traffic_data
        # Project-level traffic-scaling state (global spinbox + follow-global
        # flags).  Optional on legacy projects -- the normalizer + the
        # Traffic helpers default it to 100% / all-ticked.
        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()
        # Reflect the loaded global value in the UI spinbox (if it exists).
        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)
        # The two scaling helpers live on the real ``Traffic`` instance
        # only -- regression tests stub ``self.p.traffic`` with a
        # ``SimpleNamespace``, so guard with ``hasattr`` before calling.
        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()
        self.p.segment_data = data['segment_data']
        self.p.drift_values = data['drift']
        self.p.drift_settings.drift_values = data['drift']
        # Load the write-once segments_imported snapshot.  Absent on
        # legacy files — ``get_all_for_save`` seeds it on next save.
        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['dist1'])
            self.p.segment_data[key]['dist2'] = np.array(item['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)
        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)
        
        # Load data to canvas.  Order matters: layers added EARLIER end
        # up at the BOTTOM of the QGIS layer panel and therefore render
        # behind layers added later.  We want depth polygons to sit at
        # the back (they fill large areas of the canvas), structures in
        # the middle, and the leg lines on top so they stay visible.
        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):
            # obj = [id, height_value, polygon] - use height_value (obj[1]) for layer name
            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')
        # Legs go on last so they stay on top of depths + structures.
        self.p.load_lines(data)

        # Load oil-spill consequence inputs into the handler.  If the
        # handler hasn't been constructed yet (testing path) we silently
        # skip; the next save will just write defaults.
        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 {},
            )

        # Load junction transition matrices.  Missing block is treated
        # as "rebuild from scratch" — the handler will discover the
        # junctions implied by the loaded segment_data and seed them
        # with geometric defaults.
        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_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 in self.p.segment_data.keys():
            self.p.main_widget.cbTrafficSelectSeg.addItem(str(key))
        self.p.traffic.c_seg = self.p.main_widget.cbTrafficSelectSeg.currentText()

    def populate_tbl(self, data:list, tbl:QTableWidget):
        tbl.setRowCount(len(data))
        for i, line in enumerate(data):
            for j, value in enumerate(line):
                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)
