"""Application state management.

AppState is the central data store for the web UI. It holds all loaded tests,
manages channel transforms, and provides the data that callbacks need to
render plots, compute criteria, and generate reports.

AppState is NOT a Dash Store — it lives server-side in Python memory. The Dash
stores hold lightweight references (test IDs, selected channels, plot config)
that callbacks use to look up data from AppState.

This design keeps large numpy arrays out of the browser and avoids serialization.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

from impakt.channel.model import Channel, ChannelGroup, TestData
from impakt.io.reader import get_registry
from impakt.template.library import TemplateLibrary
from impakt.template.model import PlotDefinition, SessionState, TemplateSpec
from impakt.template.session import SessionManager
from impakt.transform.align import YAlign
from impakt.transform.cfc import CFCFilter

logger = logging.getLogger(__name__)


class LoadedTest:
    """A single loaded test with its metadata and channel access."""

    def __init__(self, test_data: TestData, color_offset: int = 0) -> None:
        self.data = test_data
        self.color_offset = color_offset

    @property
    def test_id(self) -> str:
        return self.data.test_id

    @property
    def label(self) -> str:
        """Short label for UI display."""
        meta = self.data.metadata
        parts = [meta.test_number]
        if meta.vehicle.make:
            parts.append(f"{meta.vehicle.make} {meta.vehicle.model}".strip())
        return " — ".join(parts) if len(parts) > 1 else parts[0]

    @property
    def channel_count(self) -> int:
        return len(self.data)

    def __len__(self) -> int:
        return len(self.data)


class AppState:
    """Server-side application state.

    Manages multiple loaded tests and provides channel resolution
    for the UI callbacks.
    """

    def __init__(self) -> None:
        self._tests: dict[str, LoadedTest] = {}
        self._test_order: list[str] = []
        self._color_counter: int = 0
        self._template_library = TemplateLibrary()
        self._active_template: TemplateSpec | None = None
        self._session_managers: dict[str, SessionManager] = {}
        # Per-channel transform overrides: {channel_key: {"cfc": "600", "y_align": True}}
        self.channel_overrides: dict[str, dict[str, Any]] = {}
        # Active corridors: list of {name, time, lower, upper, visible}
        self.corridors: list[dict[str, Any]] = []

    def load_test(self, path: str | Path) -> LoadedTest:
        """Load a test from a path and add it to the state.

        If a test with the same ID is already loaded, replaces it.

        Returns:
            The loaded test.

        Raises:
            ValueError: If the path is not a valid test directory.
        """
        path = Path(path).resolve()
        registry = get_registry()
        test_data = registry.read(path)

        loaded = LoadedTest(test_data, color_offset=self._color_counter)
        self._color_counter += len(test_data)

        test_id = test_data.test_id
        if test_id in self._tests:
            self._tests[test_id] = loaded
        else:
            self._tests[test_id] = loaded
            self._test_order.append(test_id)

        logger.info("Loaded test '%s' from %s (%d channels)", test_id, path, len(test_data))
        return loaded

    def remove_test(self, test_id: str) -> None:
        """Remove a loaded test."""
        if test_id in self._tests:
            del self._tests[test_id]
            self._test_order.remove(test_id)

    @property
    def tests(self) -> list[LoadedTest]:
        """All loaded tests in load order."""
        return [self._tests[tid] for tid in self._test_order if tid in self._tests]

    @property
    def test_ids(self) -> list[str]:
        return list(self._test_order)

    @property
    def primary_test(self) -> LoadedTest | None:
        """The first loaded test (primary for criteria, etc.)."""
        if self._test_order:
            return self._tests.get(self._test_order[0])
        return None

    def get_test(self, test_id: str) -> LoadedTest | None:
        return self._tests.get(test_id)

    def get_channel(self, test_id: str, channel_name: str) -> Channel | None:
        """Resolve a channel from a test ID and channel name."""
        test = self._tests.get(test_id)
        if test is None:
            return None
        try:
            return test.data.get(channel_name)
        except KeyError:
            return None

    def resolve_channel(
        self,
        channel_key: str,
        cfc_class: int | None = None,
        y_align: bool = False,
    ) -> Channel | None:
        """Resolve a channel key (test_id::channel_name) and apply transforms.

        Channel keys use the format 'test_id::channel_name'. If no '::'
        separator, assumes the primary test.
        """
        if "::" in channel_key:
            test_id, ch_name = channel_key.split("::", 1)
        else:
            if self.primary_test is None:
                return None
            test_id = self.primary_test.test_id
            ch_name = channel_key

        ch = self.get_channel(test_id, ch_name)
        if ch is None:
            return None

        # Apply transforms
        if cfc_class is not None:
            try:
                ch = CFCFilter(cfc_class=cfc_class).apply(ch)
            except (ValueError, Exception):
                pass

        if y_align:
            ch = YAlign().apply(ch)

        return ch

    def build_channel_tree(
        self,
    ) -> dict[str, dict[str, dict[str, dict[str, list[dict[str, str]]]]]]:
        """Build a hierarchical channel tree across all loaded tests.

        Returns:
            {test_id: {object_label: {location_label: {measurement_label: [{name, label, key}]}}}}
        """
        result: dict[str, dict[str, dict[str, dict[str, list[dict[str, str]]]]]] = {}

        for loaded in self.tests:
            tree = loaded.data.channel_tree()
            test_tree: dict[str, dict[str, dict[str, list[dict[str, str]]]]] = {}

            for obj, locations in sorted(tree.items()):
                obj_tree: dict[str, dict[str, list[dict[str, str]]]] = {}
                for loc, measurements in sorted(locations.items()):
                    loc_tree: dict[str, list[dict[str, str]]] = {}
                    for meas, channels in sorted(measurements.items()):
                        ch_list = []
                        for ch in channels:
                            label = ch.code.short_label if ch.code.is_valid else ch.name
                            ch_list.append(
                                {
                                    "name": ch.name,
                                    "label": label,
                                    "key": f"{loaded.test_id}::{ch.name}",
                                    "unit": ch.unit,
                                    "peak": f"{ch.peak:.2f}",
                                    "samples": str(ch.n_samples),
                                    "rate": f"{ch.sample_rate:.0f}",
                                }
                            )
                        loc_tree[meas] = ch_list
                    obj_tree[loc] = loc_tree
                test_tree[obj] = obj_tree
            result[loaded.test_id] = test_tree

        return result

    def flat_channel_list(self) -> list[dict[str, str]]:
        """Flat list of all channels across all tests, for dropdown/checklist options."""
        items: list[dict[str, str]] = []
        multi = len(self._tests) > 1

        for loaded in self.tests:
            prefix = f"[{loaded.test_id}] " if multi else ""
            for ch in sorted(loaded.data, key=lambda c: c.name):
                label = ch.code.short_label if ch.code.is_valid else ch.name
                items.append(
                    {
                        "label": f"{prefix}{label}",
                        "value": f"{loaded.test_id}::{ch.name}",
                    }
                )

        return items

    # ----- Template & Session -----

    @property
    def template_library(self) -> TemplateLibrary:
        return self._template_library

    @property
    def active_template(self) -> TemplateSpec | None:
        return self._active_template

    @property
    def template_names(self) -> list[str]:
        return self._template_library.list()

    def apply_template(
        self,
        name: str,
        selected_keys: list[str] | None = None,
    ) -> tuple[list[str], dict[str, str]]:
        """Apply a template by name.

        Resolves the template's channel patterns against the primary test,
        sets the active template, and returns the resolved channel keys and
        transform settings.

        Returns:
            (selected_channel_keys, transform_settings)
        """
        template = self._template_library.get(name)
        self._active_template = template

        # Persist to session if primary test has a path
        primary = self.primary_test
        if primary and primary.data.path:
            mgr = self._get_session_manager(primary.data.path)
            mgr.apply_template(template)

        # Resolve channel patterns from the template
        resolved_keys: list[str] = []
        if primary:
            for plot_def in template.plots:
                for pattern in plot_def.channel_patterns:
                    matches = primary.data.find(pattern)
                    for ch in matches:
                        key = f"{primary.test_id}::{ch.name}"
                        if key not in resolved_keys:
                            resolved_keys.append(key)

        # Transform settings
        transforms: dict[str, str] = {}
        if template.default_cfc is not None:
            transforms["cfc"] = str(template.default_cfc)
        else:
            transforms["cfc"] = "none"

        logger.info(
            "Applied template '%s' — resolved %d channels",
            name,
            len(resolved_keys),
        )
        return resolved_keys, transforms

    def save_as_template(
        self,
        name: str,
        description: str,
        selected_keys: list[str],
        cfc_value: str,
        x1: float | None,
        x2: float | None,
        protocol: str = "",
    ) -> TemplateSpec:
        """Capture the current UI state as a new template.

        Converts selected channels back to patterns and stores the
        current filter/cursor/protocol settings.
        """
        # Convert selected keys to channel patterns
        patterns: list[str] = []
        for key in selected_keys:
            if "::" in key:
                _, ch_name = key.split("::", 1)
            else:
                ch_name = key
            # Use the raw channel name as a pattern (exact match)
            if ch_name not in patterns:
                patterns.append(ch_name)

        # Build plot definition
        plot = PlotDefinition(
            title=name,
            channel_patterns=patterns,
            x_cursors=(x1, x2) if x1 is not None and x2 is not None else None,
        )

        # Build transforms list
        transforms: list[dict[str, Any]] = []
        if cfc_value and cfc_value != "none":
            transforms.append({"type": "cfc_filter", "cfc_class": int(cfc_value)})
        if transforms:
            plot.transforms = transforms

        template = TemplateSpec(
            name=name,
            description=description,
            plots=[plot],
            default_cfc=int(cfc_value) if cfc_value and cfc_value != "none" else None,
            protocol=protocol,
        )

        self._template_library.save(template)
        self._active_template = template

        logger.info("Saved template '%s' with %d channel patterns", name, len(patterns))
        return template

    def save_session(self, selected_keys: list[str], cfc_value: str, **overrides: Any) -> None:
        """Auto-save current state to the session for the primary test."""
        primary = self.primary_test
        if not primary or not primary.data.path:
            return

        mgr = self._get_session_manager(primary.data.path)
        mgr.state.overrides = {
            "selected_channels": selected_keys,
            "cfc": cfc_value,
            **overrides,
        }
        mgr.save()

    def load_session_state(self) -> dict[str, Any] | None:
        """Load saved session state for the primary test, if any."""
        primary = self.primary_test
        if not primary or not primary.data.path:
            return None

        mgr = self._get_session_manager(primary.data.path)
        if not mgr.has_session:
            return None

        return {
            "selected_channels": mgr.state.overrides.get("selected_channels", []),
            "cfc": mgr.state.overrides.get("cfc", "none"),
            "template": mgr.state.template_name,
        }

    def _get_session_manager(self, path: Path) -> SessionManager:
        key = str(path)
        if key not in self._session_managers:
            self._session_managers[key] = SessionManager(path)
        return self._session_managers[key]

    @property
    def is_empty(self) -> bool:
        return len(self._tests) == 0

    @property
    def total_channels(self) -> int:
        return sum(t.channel_count for t in self.tests)

    def __repr__(self) -> str:
        test_info = ", ".join(f"{t.test_id}({t.channel_count}ch)" for t in self.tests)
        tmpl = f", template={self._active_template.name}" if self._active_template else ""
        return f"AppState([{test_info}]{tmpl})"
