from __future__ import annotations

import math
import re
import threading
import uuid
from collections.abc import Iterable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, field
from numbers import Integral, Real
from typing import Any, ClassVar

from models import LayerSummaryRecord, RunContext
from storage import LocalSessionStore, utc_now

try:
    import torch
    import torch.nn as nn
except Exception:  # pragma: no cover - exercised through runtime guard.
    torch = None
    nn = None

_DEFAULT_SATURATION_MODULE_TYPES = (
    "Sigmoid",
    "Tanh",
    "Hardsigmoid",
    "Hardtanh",
    "ReLU6",
)


def _normalize_string_values(values: Iterable[str], *, label: str) -> tuple[str, ...]:
    normalized_values: list[str] = []
    for raw_value in values:
        if not isinstance(raw_value, str):
            raise TypeError(f"{label} values must be strings.")
        candidate = raw_value.strip()
        if not candidate:
            raise ValueError(f"{label} values must not be empty.")
        if candidate not in normalized_values:
            normalized_values.append(candidate)
    return tuple(normalized_values)


def _normalize_step_values(values: Iterable[int]) -> tuple[int, ...]:
    normalized_values: list[int] = []
    for raw_value in values:
        if isinstance(raw_value, bool) or not isinstance(raw_value, Integral):
            raise TypeError("explicit_steps values must be integers.")
        candidate = int(raw_value)
        if candidate < 0:
            raise ValueError("explicit_steps values must not be negative.")
        if candidate not in normalized_values:
            normalized_values.append(candidate)
    return tuple(sorted(normalized_values))


@dataclass(frozen=True, slots=True)
class HookScheduleConfig:
    every_n_steps: int | None = 1
    start_step: int = 0
    stop_step: int | None = None
    explicit_steps: tuple[int, ...] = field(default_factory=tuple)

    def __post_init__(self) -> None:
        if self.every_n_steps is not None:
            if isinstance(self.every_n_steps, bool) or not isinstance(self.every_n_steps, Integral):
                raise TypeError("every_n_steps must be an integer or None.")
            if int(self.every_n_steps) <= 0:
                raise ValueError("every_n_steps must be a positive integer when provided.")
            object.__setattr__(self, "every_n_steps", int(self.every_n_steps))

        if isinstance(self.start_step, bool) or not isinstance(self.start_step, Integral):
            raise TypeError("start_step must be an integer.")
        if int(self.start_step) < 0:
            raise ValueError("start_step must not be negative.")
        object.__setattr__(self, "start_step", int(self.start_step))

        if self.stop_step is not None:
            if isinstance(self.stop_step, bool) or not isinstance(self.stop_step, Integral):
                raise TypeError("stop_step must be an integer or None.")
            if int(self.stop_step) < self.start_step:
                raise ValueError("stop_step must be greater than or equal to start_step.")
            object.__setattr__(self, "stop_step", int(self.stop_step))

        object.__setattr__(self, "explicit_steps", _normalize_step_values(self.explicit_steps))

    def should_probe(self, step: int) -> bool:
        if step in self.explicit_steps:
            return True
        if self.every_n_steps is None:
            return False
        if step < self.start_step:
            return False
        if self.stop_step is not None and step > self.stop_step:
            return False
        return (step - self.start_step) % self.every_n_steps == 0


@dataclass(frozen=True, slots=True)
class LayerFilterConfig:
    include_names: tuple[str, ...] = field(default_factory=tuple)
    exclude_names: tuple[str, ...] = field(default_factory=tuple)
    include_types: tuple[str, ...] = field(default_factory=tuple)
    exclude_types: tuple[str, ...] = field(default_factory=tuple)
    include_name_patterns: tuple[str, ...] = field(default_factory=tuple)
    exclude_name_patterns: tuple[str, ...] = field(default_factory=tuple)
    leaf_modules_only: bool = True

    def __post_init__(self) -> None:
        object.__setattr__(
            self,
            "include_names",
            _normalize_string_values(self.include_names, label="include_names"),
        )
        object.__setattr__(
            self,
            "exclude_names",
            _normalize_string_values(self.exclude_names, label="exclude_names"),
        )
        object.__setattr__(
            self,
            "include_types",
            _normalize_string_values(self.include_types, label="include_types"),
        )
        object.__setattr__(
            self,
            "exclude_types",
            _normalize_string_values(self.exclude_types, label="exclude_types"),
        )
        include_patterns = _normalize_string_values(
            self.include_name_patterns,
            label="include_name_patterns",
        )
        exclude_patterns = _normalize_string_values(
            self.exclude_name_patterns,
            label="exclude_name_patterns",
        )
        for pattern in include_patterns + exclude_patterns:
            re.compile(pattern)
        object.__setattr__(self, "include_name_patterns", include_patterns)
        object.__setattr__(self, "exclude_name_patterns", exclude_patterns)
        if not isinstance(self.leaf_modules_only, bool):
            raise TypeError("leaf_modules_only must be a boolean.")


@dataclass(frozen=True, slots=True)
class HookEngineConfig:
    schedule: HookScheduleConfig = field(default_factory=HookScheduleConfig)
    filters: LayerFilterConfig = field(default_factory=LayerFilterConfig)
    capture_forward: bool = True
    capture_backward: bool = True
    batch_size: int = 100
    raise_on_error: bool | None = None
    dead_unit_epsilon: float = 0.0
    saturation_threshold: float | None = 0.95
    saturation_module_types: tuple[str, ...] = field(
        default_factory=lambda: _DEFAULT_SATURATION_MODULE_TYPES
    )

    def __post_init__(self) -> None:
        if not isinstance(self.schedule, HookScheduleConfig):
            raise TypeError("schedule must be a HookScheduleConfig instance.")
        if not isinstance(self.filters, LayerFilterConfig):
            raise TypeError("filters must be a LayerFilterConfig instance.")
        if not isinstance(self.capture_forward, bool):
            raise TypeError("capture_forward must be a boolean.")
        if not isinstance(self.capture_backward, bool):
            raise TypeError("capture_backward must be a boolean.")
        if not self.capture_forward and not self.capture_backward:
            raise ValueError("HookEngineConfig must enable at least one hook direction.")
        if isinstance(self.batch_size, bool) or not isinstance(self.batch_size, Integral):
            raise TypeError("batch_size must be an integer.")
        if int(self.batch_size) <= 0:
            raise ValueError("batch_size must be a positive integer.")
        object.__setattr__(self, "batch_size", int(self.batch_size))

        if self.raise_on_error is not None and not isinstance(self.raise_on_error, bool):
            raise TypeError("raise_on_error must be a boolean or None.")

        if isinstance(self.dead_unit_epsilon, bool) or not isinstance(self.dead_unit_epsilon, Real):
            raise TypeError("dead_unit_epsilon must be a real number.")
        if float(self.dead_unit_epsilon) < 0.0:
            raise ValueError("dead_unit_epsilon must not be negative.")
        object.__setattr__(self, "dead_unit_epsilon", float(self.dead_unit_epsilon))

        if self.saturation_threshold is not None:
            if isinstance(self.saturation_threshold, bool) or not isinstance(
                self.saturation_threshold,
                Real,
            ):
                raise TypeError("saturation_threshold must be a real number or None.")
            if float(self.saturation_threshold) <= 0.0:
                raise ValueError("saturation_threshold must be positive when provided.")
            object.__setattr__(self, "saturation_threshold", float(self.saturation_threshold))

        object.__setattr__(
            self,
            "saturation_module_types",
            _normalize_string_values(
                self.saturation_module_types,
                label="saturation_module_types",
            ),
        )


@dataclass(slots=True)
class _ProbeState:
    step: int | None = None
    epoch: int | None = None
    active: bool = False
    call_counts: dict[tuple[str, str], int] = field(default_factory=dict)


@dataclass(frozen=True, slots=True)
class _ProbeReservation:
    step: int
    epoch: int | None
    call_index: int


@dataclass(frozen=True, slots=True)
class _TensorSummary:
    element_count: int
    tensor_shape: tuple[int, ...] | None
    mean: float
    std: float
    norm: float
    variance: float
    dead_unit_ratio: float | None
    saturation_ratio: float | None
    metadata: dict[str, Any] | None = None


class HookEngine:
    """PyTorch hook-based layer summary collector for a single training run."""

    _sequence_lock: ClassVar[threading.RLock] = threading.RLock()
    _next_sequence_by_run: ClassVar[dict[tuple[str, str], int]] = {}

    def __init__(
        self,
        session: Any,
        model: Any,
        optimizer: Any | None = None,
        *,
        config: HookEngineConfig | None = None,
    ) -> None:
        if torch is None or nn is None:
            raise ImportError("HookEngine requires PyTorch to be installed.")
        if not isinstance(model, nn.Module):
            raise TypeError("HookEngine expects a torch.nn.Module model instance.")

        self._context, self._store = self._resolve_session_dependencies(session)
        self._model = model
        self._optimizer = optimizer
        self._config = config or HookEngineConfig()
        self._include_name_patterns = tuple(
            re.compile(pattern) for pattern in self._config.filters.include_name_patterns
        )
        self._exclude_name_patterns = tuple(
            re.compile(pattern) for pattern in self._config.filters.exclude_name_patterns
        )
        self._raise_on_error = (
            self._config.raise_on_error
            if self._config.raise_on_error is not None
            else self._context.run.strict_mode
        )
        self._handles: list[Any] = []
        self._buffer: list[LayerSummaryRecord] = []
        self._lock = threading.RLock()
        self._closed = False
        self._probe_state = _ProbeState()
        self._instrumented_modules = tuple(self._iter_instrumented_modules())
        self._register_hooks()

    @property
    def config(self) -> HookEngineConfig:
        return self._config

    @property
    def run_id(self) -> str:
        return self._context.run_id

    @property
    def buffered_count(self) -> int:
        with self._lock:
            return len(self._buffer)

    @property
    def instrumented_module_names(self) -> tuple[str, ...]:
        return tuple(name for name, _module in self._instrumented_modules)

    def should_probe(self, step: int) -> bool:
        normalized_step = self._normalize_index("step", step)
        if normalized_step is None:
            raise ValueError("step must be provided for probe scheduling.")
        return self._config.schedule.should_probe(normalized_step)

    def begin_step(self, step: int, *, epoch: int | None = None) -> bool:
        """Arm the hook engine for the current training step."""

        self._ensure_open()
        normalized_step = self._normalize_index("step", step)
        if normalized_step is None:
            raise ValueError("step must be provided for begin_step.")
        normalized_epoch = self._normalize_index("epoch", epoch)
        active = self._config.schedule.should_probe(normalized_step)
        with self._lock:
            self._probe_state = _ProbeState(
                step=normalized_step,
                epoch=normalized_epoch,
                active=active,
            )
        return active

    def end_step(self) -> None:
        with self._lock:
            self._probe_state = _ProbeState()

    @contextmanager
    def step(self, step: int, *, epoch: int | None = None) -> Iterator[bool]:
        """Context manager that arms probes for a single forward/backward pass."""

        active = self.begin_step(step, epoch=epoch)
        try:
            yield active
        finally:
            self.end_step()

    def flush(self) -> tuple[LayerSummaryRecord, ...]:
        """Persist any buffered layer summaries."""

        self._ensure_open()
        with self._lock:
            if not self._buffer:
                return ()
            batch = tuple(self._buffer)
            self._buffer.clear()

        try:
            return self._store.write_layer_summaries(self.run_id, batch)
        except Exception:
            with self._lock:
                self._buffer = list(batch) + self._buffer
            raise

    def read(
        self,
        *,
        module_name: str | None = None,
        module_type: str | None = None,
        summary_type: str | None = None,
        hook_direction: str | None = None,
        step: int | None = None,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        limit: int | None = None,
    ) -> tuple[LayerSummaryRecord, ...]:
        """Flush pending summaries and return persisted summaries for the run."""

        if not self._closed:
            self.flush()
        return self._store.read_layer_summaries(
            self.run_id,
            module_name=module_name,
            module_type=module_type,
            summary_type=summary_type,
            hook_direction=hook_direction,
            step=step,
            step_min=step_min,
            step_max=step_max,
            epoch=epoch,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
            limit=limit,
        )

    def close(self) -> None:
        if self._closed:
            return

        flush_error: Exception | None = None
        try:
            self.flush()
        except Exception as exc:  # pragma: no cover - rare storage failure path.
            flush_error = exc
        finally:
            for handle in self._handles:
                handle.remove()
            self._handles.clear()
            self.end_step()
            self._closed = True

        if flush_error is not None:
            raise flush_error

    def __enter__(self) -> HookEngine:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        self.close()

    def _iter_instrumented_modules(self) -> Iterator[tuple[str, Any]]:
        for raw_name, module in self._model.named_modules():
            has_children = any(module.children())
            if raw_name == "" and has_children:
                continue
            if self._config.filters.leaf_modules_only and has_children:
                continue

            module_name = raw_name or "<root>"
            if not self._matches_filters(module_name, module):
                continue
            yield module_name, module

    def _matches_filters(self, module_name: str, module: Any) -> bool:
        type_names = self._module_type_names(module)
        filters = self._config.filters

        if filters.include_names and module_name not in filters.include_names:
            return False
        if module_name in filters.exclude_names:
            return False

        if filters.include_types and not any(type_name in filters.include_types for type_name in type_names):
            return False
        if any(type_name in filters.exclude_types for type_name in type_names):
            return False

        if self._include_name_patterns and not any(
            pattern.search(module_name) for pattern in self._include_name_patterns
        ):
            return False
        if any(pattern.search(module_name) for pattern in self._exclude_name_patterns):
            return False

        return True

    def _register_hooks(self) -> None:
        for module_name, module in self._instrumented_modules:
            if self._config.capture_forward:
                self._handles.append(
                    module.register_forward_hook(self._build_forward_hook(module_name))
                )
            if self._config.capture_backward:
                backward_hook = self._build_backward_hook(module_name)
                if hasattr(module, "register_full_backward_hook"):
                    handle = module.register_full_backward_hook(backward_hook)
                else:  # pragma: no cover - only for older PyTorch versions.
                    handle = module.register_backward_hook(backward_hook)
                self._handles.append(handle)

    def _build_forward_hook(self, module_name: str):
        def _hook(module: Any, inputs: tuple[Any, ...], output: Any) -> None:
            del inputs
            reservation = self._reserve_probe(module_name, "activation")
            if reservation is None:
                return
            try:
                self._record_summary(
                    module_name=module_name,
                    module=module,
                    value=output,
                    summary_type="activation",
                    hook_direction="forward",
                    reservation=reservation,
                )
            except Exception as exc:
                self._handle_hook_error(exc, module_name=module_name, hook_name="forward")

        return _hook

    def _build_backward_hook(self, module_name: str):
        def _hook(module: Any, grad_input: tuple[Any, ...], grad_output: tuple[Any, ...]) -> None:
            del grad_input
            reservation = self._reserve_probe(module_name, "gradient")
            if reservation is None:
                return
            try:
                self._record_summary(
                    module_name=module_name,
                    module=module,
                    value=grad_output,
                    summary_type="gradient",
                    hook_direction="backward",
                    reservation=reservation,
                )
            except Exception as exc:
                self._handle_hook_error(exc, module_name=module_name, hook_name="backward")

        return _hook

    def _record_summary(
        self,
        *,
        module_name: str,
        module: Any,
        value: Any,
        summary_type: str,
        hook_direction: str,
        reservation: _ProbeReservation,
    ) -> None:
        summary = self._summarize_value(value, summary_type=summary_type, module=module)
        if summary is None:
            return

        record = LayerSummaryRecord(
            id=uuid.uuid4().hex,
            run_id=self.run_id,
            sequence_id=self._allocate_sequence_id(),
            step=reservation.step,
            epoch=reservation.epoch,
            module_name=module_name,
            module_type=module.__class__.__name__,
            summary_type=summary_type,
            hook_direction=hook_direction,
            call_index=reservation.call_index,
            element_count=summary.element_count,
            timestamp=utc_now(),
            tensor_shape=summary.tensor_shape,
            mean=summary.mean,
            std=summary.std,
            norm=summary.norm,
            variance=summary.variance,
            dead_unit_ratio=summary.dead_unit_ratio,
            saturation_ratio=summary.saturation_ratio,
            metadata=summary.metadata,
        )
        self._enqueue(record)

    def _reserve_probe(
        self,
        module_name: str,
        summary_type: str,
    ) -> _ProbeReservation | None:
        with self._lock:
            if self._closed or not self._probe_state.active or self._probe_state.step is None:
                return None

            key = (module_name, summary_type)
            call_index = self._probe_state.call_counts.get(key, 0)
            self._probe_state.call_counts[key] = call_index + 1
            return _ProbeReservation(
                step=self._probe_state.step,
                epoch=self._probe_state.epoch,
                call_index=call_index,
            )

    def _enqueue(self, record: LayerSummaryRecord) -> None:
        should_flush = False
        with self._lock:
            self._buffer.append(record)
            should_flush = len(self._buffer) >= self._config.batch_size

        if should_flush:
            self.flush()

    def _summarize_value(
        self,
        value: Any,
        *,
        summary_type: str,
        module: Any,
    ) -> _TensorSummary | None:
        tensors = tuple(self._iter_tensors(value))
        if not tensors:
            return None

        element_count = 0
        tensor_shapes: list[tuple[int, ...]] = []
        summed_value = 0.0
        summed_square = 0.0
        zero_count = 0
        saturated_count = 0

        capture_dead_units = summary_type == "activation"
        capture_saturation = (
            summary_type == "activation"
            and self._config.saturation_threshold is not None
            and self._module_supports_saturation(module)
        )

        for tensor in tensors:
            prepared = self._prepare_tensor(tensor)
            if prepared is None:
                continue

            current_count = int(prepared.numel())
            if current_count == 0:
                continue

            element_count += current_count
            tensor_shapes.append(tuple(int(dimension) for dimension in tensor.shape))
            summed_value += float(prepared.sum().item())
            summed_square += float(prepared.square().sum().item())

            if capture_dead_units:
                if self._config.dead_unit_epsilon == 0.0:
                    zero_count += int(prepared.eq(0).sum().item())
                else:
                    zero_count += int(
                        prepared.abs().le(self._config.dead_unit_epsilon).sum().item()
                    )

            if capture_saturation:
                saturated_count += int(
                    prepared.abs().ge(self._config.saturation_threshold).sum().item()
                )

        if element_count == 0:
            return None

        mean = summed_value / element_count
        variance = max((summed_square / element_count) - (mean * mean), 0.0)
        metadata: dict[str, Any] | None = None
        tensor_shape = tensor_shapes[0] if len(tensor_shapes) == 1 else None
        if len(tensor_shapes) > 1:
            metadata = {"tensor_shapes": [list(shape) for shape in tensor_shapes]}

        return _TensorSummary(
            element_count=element_count,
            tensor_shape=tensor_shape,
            mean=mean,
            std=math.sqrt(variance),
            norm=math.sqrt(max(summed_square, 0.0)),
            variance=variance,
            dead_unit_ratio=(zero_count / element_count) if capture_dead_units else None,
            saturation_ratio=(saturated_count / element_count) if capture_saturation else None,
            metadata=metadata,
        )

    def _prepare_tensor(self, tensor: Any) -> Any | None:
        if torch is None or not isinstance(tensor, torch.Tensor):
            return None

        prepared = tensor.detach()
        if prepared.is_sparse:
            prepared = prepared.coalesce().values()
        if prepared.is_complex():
            prepared = prepared.abs()
        if prepared.numel() == 0:
            return None
        if not torch.is_floating_point(prepared) or prepared.dtype != torch.float32:
            prepared = prepared.float()
        return prepared

    def _iter_tensors(self, value: Any) -> Iterator[Any]:
        if torch is not None and isinstance(value, torch.Tensor):
            yield value
            return
        if isinstance(value, Mapping):
            for nested_value in value.values():
                yield from self._iter_tensors(nested_value)
            return
        if isinstance(value, (list, tuple)):
            for nested_value in value:
                yield from self._iter_tensors(nested_value)

    def _module_supports_saturation(self, module: Any) -> bool:
        type_names = self._module_type_names(module)
        return any(type_name in self._config.saturation_module_types for type_name in type_names)

    def _allocate_sequence_id(self) -> int:
        run_key = (self._context.db_path, self.run_id)
        with self._sequence_lock:
            next_sequence = self._next_sequence_by_run.get(run_key)
            if next_sequence is None:
                next_sequence = self._store.get_next_layer_summary_sequence(self.run_id)
            self._next_sequence_by_run[run_key] = next_sequence + 1
            return next_sequence

    def _handle_hook_error(
        self,
        error: Exception,
        *,
        module_name: str,
        hook_name: str,
    ) -> None:
        if self._raise_on_error:
            raise RuntimeError(
                f"HookEngine {hook_name} hook failed for module '{module_name}' in run '{self.run_id}'."
            ) from error

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("HookEngine is closed.")

    @staticmethod
    def _resolve_session_dependencies(session: Any) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(session, RunContext):
            return session, LocalSessionStore(session.storage_dir)

        context = getattr(session, "context", None)
        store = getattr(session, "_store", None)
        if isinstance(context, RunContext) and isinstance(store, LocalSessionStore):
            return context, store

        raise TypeError("HookEngine expects a TrainingSession or RunContext created by the Quark SDK.")

    @staticmethod
    def _normalize_index(label: str, value: int | None) -> int | None:
        if value is None:
            return None
        if isinstance(value, bool) or not isinstance(value, Integral):
            raise TypeError(f"{label} must be an integer when provided.")
        normalized_value = int(value)
        if normalized_value < 0:
            raise ValueError(f"{label} must not be negative.")
        return normalized_value

    @staticmethod
    def _module_type_names(module: Any) -> tuple[str, str]:
        module_type = module.__class__
        return (
            module_type.__name__,
            f"{module_type.__module__}.{module_type.__qualname__}",
        )
