from __future__ import annotations

from collections.abc import Iterator
from typing import Literal

from datapipeline.config.metadata import (
    FEATURE_VECTORS_COUNT_KEY,
    TARGET_VECTORS_COUNT_KEY,
    VectorMetadata,
)
from datapipeline.domain.sample import Sample
from datapipeline.domain.vector import Vector
from datapipeline.services.artifacts import (
    ArtifactNotRegisteredError,
    VECTOR_METADATA_SPEC,
)

from ..common import (
    VectorContextMixin,
    replace_vector,
    select_vector,
    try_get_current_context,
)


class VectorDropVerticalTransform(VectorContextMixin):
    required_artifacts = {VECTOR_METADATA_SPEC.key}
    """Drop partitions/features when metadata coverage falls below configured thresholds.

    Requires the optional `metadata.json` artifact generated by the
    `metadata` build task. The transform evaluates coverage using the recorded
    `present_count`/`null_count` metrics and prunes the schema cache once so
    downstream coverage checks stop expecting bad partitions.
    """

    def __init__(
        self,
        *,
        payload: Literal["features", "targets"] = "features",
        threshold: float,
    ) -> None:
        super().__init__(payload=payload)
        if not 0.0 <= threshold <= 1.0:
            raise ValueError("threshold must be between 0 and 1.")
        self._threshold = threshold
        self._drop_ids: set[str] | None = None
        self._schema_pruned = False

    def __call__(self, stream: Iterator[Sample]) -> Iterator[Sample]:
        return self.apply(stream)

    def apply(self, stream: Iterator[Sample]) -> Iterator[Sample]:
        drop_ids = self._resolve_drop_ids()
        if not drop_ids:
            yield from stream
            return
        self._maybe_prune_schema(drop_ids)
        for sample in stream:
            if not self._schema_pruned:
                self._maybe_prune_schema(drop_ids)
            vector = select_vector(sample, self._payload)
            if vector is None or not vector.values:
                yield sample
                continue
            retained: dict[str, object] = {}
            changed = False
            for fid, value in vector.values.items():
                if fid in drop_ids:
                    changed = True
                    continue
                retained[fid] = value
            if not changed:
                yield sample
            else:
                yield replace_vector(sample, self._payload, Vector(values=retained))

    def _resolve_drop_ids(self) -> set[str]:
        if self._drop_ids is not None:
            return self._drop_ids
        context = self._context or try_get_current_context()
        if not context:
            raise RuntimeError("VectorDropVerticalTransform requires an active pipeline context.")
        try:
            raw = context.require_artifact(VECTOR_METADATA_SPEC)
        except ArtifactNotRegisteredError as exc:
            raise RuntimeError(
                "Vector metadata artifact missing. Enable the `metadata` build task "
                "and rerun `jerry build --project <project.yaml>`."
            ) from exc
        meta = VectorMetadata.model_validate(raw)
        section_key = "targets" if self._payload == "targets" else "features"
        counts_key = (
            TARGET_VECTORS_COUNT_KEY
            if self._payload == "targets"
            else FEATURE_VECTORS_COUNT_KEY
        )

        entries = getattr(meta, section_key) or []
        window_size = self._window_size(getattr(meta, "window", None))
        total = window_size if window_size is not None else meta.counts.get(counts_key)
        if not isinstance(total, (int, float)) or total <= 0:
            if self._payload == "targets":
                raise RuntimeError(
                    "Vector metadata artifact missing counts for targets; "
                    "ensure your dataset defines target streams and rebuild."
                )
            raise RuntimeError(
                "Vector metadata artifact missing counts for features; "
                "rerun `jerry build --project <project.yaml>` to refresh metadata."
            )
        expected_buckets = float(total)
        drop_ids: set[str] = set()
        for entry in entries:
            if not isinstance(entry, dict):
                continue
            fid = entry.get("id")
            if not isinstance(fid, str):
                continue
            coverage = self._coverage_for_entry(entry, expected_buckets)
            if coverage < self._threshold:
                drop_ids.add(fid)
        self._drop_ids = drop_ids
        return drop_ids

    @staticmethod
    def _window_size(window) -> float | None:
        if window is None:
            return None
        if isinstance(window, dict):
            return window.get("size")
        return getattr(window, "size", None)

    @staticmethod
    def _coverage_for_entry(entry: dict, expected_buckets: float) -> float:
        if expected_buckets <= 0:
            return 0.0
        present = float(entry.get("present_count") or 0.0)
        nulls = float(entry.get("null_count") or 0.0)
        cadence_doc = entry.get("cadence")
        cadence = cadence_doc.get("target") if isinstance(cadence_doc, dict) else None
        observed_elements = entry.get("observed_elements")

        if isinstance(observed_elements, (int, float)) and cadence:
            # Base expected elements on buckets where this feature actually appeared
            # to avoid over-crediting sparse sequences.
            expected_elements = float(max(present, 0.0)) * float(cadence)
            if expected_elements > 0:
                return max(
                    0.0,
                    min(1.0, float(observed_elements) / expected_elements),
                )

        coverage = (present - nulls) / expected_buckets
        return max(0.0, min(1.0, coverage))

    def _maybe_prune_schema(self, drop_ids: set[str]) -> None:
        if self._schema_pruned or not drop_ids:
            return
        context = self._context or try_get_current_context()
        if not context:
            self._schema_pruned = True
            return
        cache = getattr(context, "_cache", None)
        if cache is None:
            self._schema_pruned = True
            return
        schema_key = f"schema:{self._payload}"
        if schema_key not in cache:
            return
        entries = cache.get(schema_key)
        if not entries:
            self._schema_pruned = True
            return
        kept = [entry for entry in entries if entry.get("id") not in drop_ids]
        cache[schema_key] = kept
        ids_key = f"expected_ids:{self._payload}"
        cache[ids_key] = [
            entry.get("id")
            for entry in kept
            if isinstance(entry.get("id"), str)
        ]
        self._schema_pruned = True
