from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
import re
from typing import Dict, Generic, Iterable, List, Literal, TypeVar

from tqdm import tqdm

from .name_mapping import (
    _match_structured_expr,
    _rewrite_structured_expr,
    ResolvedMapping,
    match_expr_names,
    require_dest_missing,
    require_dest_present,
    resolve_name_mappings,
)
from .refs import TensorRef, format_tensor_ref, must_model, parse_model_expr, parse_slice
from .resolver import resolve_target_names as resolve_target_names_generic
from .types import StateDictProvider, TransformError
from .validation import ensure_mapping_payload, require_expr, validate_payload_keys


class TransformControl(Enum):
    CONTINUE = "continue"
    EXIT = "exit"


@dataclass(frozen=True)
class TransformResult:
    name: str
    count: int
    control: TransformControl = TransformControl.CONTINUE


@dataclass(frozen=True)
class CompiledTransform:
    transform: "BaseTransform"
    spec: object


class BaseTransform(ABC):
    name: str
    completion_requires_payload: bool = True

    @abstractmethod
    def compile(self, payload: dict, default_model: str | None) -> object:
        raise NotImplementedError

    @abstractmethod
    def apply(self, spec: object, provider: StateDictProvider) -> TransformResult:
        raise NotImplementedError

    @abstractmethod
    def infer_output_model(self, spec: object) -> str:
        raise NotImplementedError

    def contributes_output_model(self, spec: object) -> bool:
        del spec
        return True

    def completion_reference_keys(self) -> list[str]:
        return []

    def completion_payload_start_candidates(self, prefix_text: str) -> list[str] | None:
        del prefix_text
        return None

    def completion_key_candidates(self, before_cursor: str, prefix_text: str) -> list[str] | None:
        del before_cursor, prefix_text
        return None

    def completion_value_candidates(
        self,
        value_key: str | None,
        prefix_text: str,
        model_aliases: list[str],
    ) -> list[str] | None:
        del value_key, prefix_text, model_aliases
        return None

    def completion_committed_next_candidates(self, value_key: str | None) -> list[str] | None:
        del value_key
        return None


SpecT = TypeVar("SpecT")


class TypedTransform(BaseTransform, ABC, Generic[SpecT]):
    error_type: type[TransformError] = TransformError
    spec_type: type[SpecT]

    def require_spec(self, spec: object) -> SpecT:
        if not isinstance(spec, self.spec_type):
            raise self.error_type(
                f"{self.name} expected {self.spec_type.__name__}, got {type(spec).__name__}"
            )
        return spec


REGISTRY: Dict[str, BaseTransform] = {}


def register_transform(transform: BaseTransform) -> None:
    name = getattr(transform, "name", None)
    if not isinstance(name, str) or not name:
        raise TransformError("transform must define a non-empty string 'name'")
    if name in REGISTRY:
        raise TransformError(f"transform already registered: {name}")
    REGISTRY[name] = transform


def get_transform(name: str) -> BaseTransform:
    try:
        return REGISTRY[name]
    except KeyError as exc:
        raise TransformError(f"unknown transform: {name}") from exc


def list_transforms() -> List[str]:
    return sorted(REGISTRY.keys())


def apply_transform(compiled: CompiledTransform, provider: StateDictProvider) -> TransformResult:
    return compiled.transform.apply(compiled.spec, provider)


class DestinationPolicy(Enum):
    ANY = "any"
    MUST_EXIST = "must_exist"
    MUST_NOT_EXIST = "must_not_exist"


ItemT = TypeVar("ItemT")


class IteratingTransform(TypedTransform[SpecT], ABC, Generic[SpecT, ItemT]):
    error_type: type[TransformError] = TransformError
    spec_type: type[SpecT]
    progress_desc: str | None = None
    progress_unit: str = "tensor"

    def apply(self, spec: object, provider: StateDictProvider) -> TransformResult:
        typed = self.require_spec(spec)
        items = self.resolve_items(typed, provider)
        self.validate_resolved_items(typed, items, provider)

        for item in self.iter_items(items):
            self.apply_item(typed, item, provider)

        return TransformResult(name=self.name, count=len(items))

    def iter_items(self, items: list[ItemT]) -> Iterable[ItemT]:
        if self.progress_desc is None:
            return items
        return tqdm(items, desc=self.progress_desc, unit=self.progress_unit)

    @abstractmethod
    def resolve_items(
        self,
        spec: SpecT,
        provider: StateDictProvider,
    ) -> list[ItemT]:
        ...

    def validate_resolved_items(
        self,
        spec: SpecT,
        items: list[ItemT],
        provider: StateDictProvider,
    ) -> None:
        del spec, items, provider

    @abstractmethod
    def apply_item(
        self,
        spec: SpecT,
        item: ItemT,
        provider: StateDictProvider,
    ) -> None:
        ...


@dataclass(frozen=True)
class UnarySpec:
    target_ref: TensorRef

    def collect_models(self) -> set[str]:
        return {must_model(self.target_ref)}


UnarySpecT = TypeVar("UnarySpecT", bound=UnarySpec)


def _resolve_unary_target_names(
    *,
    target_ref: TensorRef,
    provider: StateDictProvider,
    op_name: str,
    error_type: type[TransformError],
) -> list[str]:
    return resolve_target_names_generic(
        target_ref=target_ref,
        provider=provider,
        op_name=op_name,
        match_names=match_expr_names,
        error_type=error_type,
    )


class UnaryTransform(IteratingTransform[UnarySpecT, str], ABC, Generic[UnarySpecT]):
    error_type: type[TransformError] = TransformError
    spec_type: type[UnarySpecT]

    allowed_keys: set[str] = {"target"}
    required_keys: set[str] = {"target"}
    target_key: str = "target"
    slice_policy: Literal["allow", "forbid"] = "forbid"

    def completion_reference_keys(self) -> list[str]:
        return [self.target_key]

    def compile(self, payload: dict, default_model: str | None) -> UnarySpecT:
        payload = ensure_mapping_payload(payload, self.name)
        validate_payload_keys(
            payload,
            op_name=self.name,
            allowed_keys=self.allowed_keys,
            required_keys=self.required_keys,
        )

        raw_target = self.require_target_expr(payload)
        target_ref = parse_model_expr(raw_target, default_model=default_model)

        self.validate_target_ref(target_ref)

        assert target_ref.model is not None
        return self.build_spec(target_ref=target_ref, payload=payload)

    def infer_output_model(self, spec: object) -> str:
        typed = self.require_spec(spec)
        model = typed.target_ref.model
        if model is None:
            raise self.error_type(f"{self.name} output model missing")
        return model

    def require_target_expr(self, payload: dict) -> str | list[object]:
        return require_expr(payload, op_name=self.name, key=self.target_key)

    def build_spec(self, target_ref: TensorRef, payload: dict) -> UnarySpecT:
        return self.spec_type(target_ref=target_ref)

    def validate_target_ref(self, target_ref: TensorRef) -> None:
        if target_ref.slice_spec is None:
            return

        if self.slice_policy == "forbid":
            raise self.error_type(f"{self.name} target must not be sliced")

        parse_slice(target_ref.slice_spec)

    def resolve_items(self, spec: UnarySpecT, provider: StateDictProvider) -> list[str]:
        return _resolve_unary_target_names(
            target_ref=spec.target_ref,
            provider=provider,
            op_name=self.name,
            error_type=self.error_type,
        )

    def resolve_targets(self, spec: UnarySpecT, provider: StateDictProvider) -> list[str]:
        return self.resolve_items(spec, provider)

    @abstractmethod
    def apply_to_target(self, spec: UnarySpecT, name: str, provider: StateDictProvider) -> None:
        ...

    def apply_item(self, spec: UnarySpecT, item: str, provider: StateDictProvider) -> None:
        self.apply_to_target(spec, item, provider)


@dataclass(frozen=True)
class BinaryMappingSpec:
    from_ref: TensorRef
    to_ref: TensorRef

    def collect_models(self) -> set[str]:
        return {must_model(self.from_ref), must_model(self.to_ref)}


BinarySpecT = TypeVar("BinarySpecT", bound=BinaryMappingSpec)


class BinaryMappingTransform(IteratingTransform[BinarySpecT, ResolvedMapping], ABC, Generic[BinarySpecT]):
    error_type: type[TransformError] = TransformError
    spec_type: type[BinarySpecT]
    destination_policy: DestinationPolicy = DestinationPolicy.ANY

    allowed_keys = {"from", "to"}
    required_keys = {"from", "to"}

    def completion_reference_keys(self) -> list[str]:
        return ["from", "to"]

    def compile(self, payload: dict, default_model: str | None) -> BinarySpecT:
        payload = ensure_mapping_payload(payload, self.name)
        validate_payload_keys(
            payload,
            op_name=self.name,
            allowed_keys=self.allowed_keys,
            required_keys=self.required_keys,
        )

        from_ref, to_ref = self.parse_refs(payload, default_model)
        self.validate_refs(from_ref, to_ref)

        assert from_ref.model is not None
        assert to_ref.model is not None
        return self.build_spec(from_ref=from_ref, to_ref=to_ref)

    def infer_output_model(self, spec: object) -> str:
        typed = self.require_spec(spec)
        model = typed.to_ref.model
        if model is None:
            raise self.error_type(f"{self.name} output model missing")
        return model

    def parse_refs(
        self,
        payload: dict,
        default_model: str | None,
    ) -> tuple[TensorRef, TensorRef]:
        raw_from = require_expr(payload, op_name=self.name, key="from")
        raw_to = require_expr(payload, op_name=self.name, key="to")

        from_ref = parse_model_expr(raw_from, default_model=default_model)
        to_ref = parse_model_expr(raw_to, default_model=default_model)
        return from_ref, to_ref

    def build_spec(self, from_ref: TensorRef, to_ref: TensorRef) -> BinarySpecT:
        return self.spec_type(from_ref=from_ref, to_ref=to_ref)

    def resolve_items(
        self,
        spec: BinarySpecT,
        provider: StateDictProvider,
    ) -> list[ResolvedMapping]:
        return resolve_name_mappings(
            from_ref=spec.from_ref,
            to_ref=spec.to_ref,
            provider=provider,
            op_name=self.name,
        )

    @abstractmethod
    def validate_refs(self, from_ref: TensorRef, to_ref: TensorRef) -> None:
        ...

    def validate_resolved_items(
        self,
        spec: BinarySpecT,
        mappings: list[ResolvedMapping],
        provider: StateDictProvider,
    ) -> None:
        del spec
        if self.destination_policy is DestinationPolicy.MUST_EXIST:
            require_dest_present(
                mappings=mappings,
                provider=provider,
                op_name=self.name,
            )
            return

        if self.destination_policy is DestinationPolicy.MUST_NOT_EXIST:
            require_dest_missing(
                mappings=mappings,
                provider=provider,
                op_name=self.name,
            )

    def apply_mapping(self, item: ResolvedMapping, provider: StateDictProvider) -> None:
        del item, provider
        raise NotImplementedError

    def apply_item(self, spec: BinarySpecT, item: ResolvedMapping, provider: StateDictProvider) -> None:
        del spec
        self.apply_mapping(item, provider)


@dataclass(frozen=True)
class TernaryMappingSpec:
    from_a_ref: TensorRef
    from_b_ref: TensorRef
    to_ref: TensorRef

    def collect_models(self) -> set[str]:
        return {
            must_model(self.from_a_ref),
            must_model(self.from_b_ref),
            must_model(self.to_ref),
        }


@dataclass(frozen=True)
class ResolvedTernaryMapping:
    src_a_model: str
    src_a_name: str
    src_a_slice: tuple[object, ...] | None
    src_b_model: str
    src_b_name: str
    src_b_slice: tuple[object, ...] | None
    dst_model: str
    dst_name: str
    dst_slice: tuple[object, ...] | None


TernarySpecT = TypeVar("TernarySpecT", bound=TernaryMappingSpec)


class TernaryMappingTransform(IteratingTransform[TernarySpecT, ResolvedTernaryMapping], ABC, Generic[TernarySpecT]):
    error_type: type[TransformError] = TransformError
    spec_type: type[TernarySpecT]
    destination_policy: DestinationPolicy = DestinationPolicy.ANY

    allowed_keys = {"from_a", "from_b", "to"}
    required_keys = {"from_a", "from_b", "to"}

    def completion_reference_keys(self) -> list[str]:
        return ["from_a", "from_b", "to"]

    def compile(self, payload: dict, default_model: str | None) -> TernarySpecT:
        payload = ensure_mapping_payload(payload, self.name)
        validate_payload_keys(
            payload,
            op_name=self.name,
            allowed_keys=self.allowed_keys,
            required_keys=self.required_keys,
        )

        from_a_ref, from_b_ref, to_ref = self.parse_refs(payload, default_model)
        self.validate_refs(from_a_ref, from_b_ref, to_ref)

        assert from_a_ref.model is not None
        assert from_b_ref.model is not None
        assert to_ref.model is not None
        return self.build_spec(from_a_ref=from_a_ref, from_b_ref=from_b_ref, to_ref=to_ref)

    def infer_output_model(self, spec: object) -> str:
        typed = self.require_spec(spec)
        model = typed.to_ref.model
        if model is None:
            raise self.error_type(f"{self.name} output model missing")
        return model

    def parse_refs(
        self,
        payload: dict,
        default_model: str | None,
    ) -> tuple[TensorRef, TensorRef, TensorRef]:
        raw_from_a = require_expr(payload, op_name=self.name, key="from_a")
        raw_from_b = require_expr(payload, op_name=self.name, key="from_b")
        raw_to = require_expr(payload, op_name=self.name, key="to")

        from_a_ref = parse_model_expr(raw_from_a, default_model=default_model)
        from_b_ref = parse_model_expr(raw_from_b, default_model=default_model)
        to_ref = parse_model_expr(raw_to, default_model=default_model)
        return from_a_ref, from_b_ref, to_ref

    def build_spec(self, from_a_ref: TensorRef, from_b_ref: TensorRef, to_ref: TensorRef) -> TernarySpecT:
        return self.spec_type(from_a_ref=from_a_ref, from_b_ref=from_b_ref, to_ref=to_ref)

    @abstractmethod
    def validate_refs(self, from_a_ref: TensorRef, from_b_ref: TensorRef, to_ref: TensorRef) -> None:
        ...

    def resolve_items(
        self,
        spec: TernarySpecT,
        provider: StateDictProvider,
    ) -> list[ResolvedTernaryMapping]:
        from_a_ref = spec.from_a_ref
        from_b_ref = spec.from_b_ref
        to_ref = spec.to_ref

        a_slice = parse_slice(from_a_ref.slice_spec) if from_a_ref.slice_spec is not None else None
        b_slice = parse_slice(from_b_ref.slice_spec) if from_b_ref.slice_spec is not None else None
        dst_slice = parse_slice(to_ref.slice_spec) if to_ref.slice_spec is not None else None

        src_a_model = must_model(from_a_ref)
        src_b_model = must_model(from_b_ref)
        dst_model = must_model(to_ref)
        src_a_sd = provider.get_state_dict(src_a_model)
        src_b_sd = provider.get_state_dict(src_b_model)

        a_expr = from_a_ref.expr
        b_expr = from_b_ref.expr
        to_expr = to_ref.expr
        if isinstance(a_expr, str) and isinstance(b_expr, str) and isinstance(to_expr, str):
            src_a_names = match_expr_names(
                expr=a_expr,
                names=src_a_sd.keys(),
                op_name=self.name,
                role="source_a",
            )
            if not src_a_names:
                raise self.error_type(
                    f"{self.name} source_a matched zero tensors: "
                    f"{format_tensor_ref(from_a_ref)}; available tensors: {sorted(src_a_sd.keys())}"
                )

            resolved: list[ResolvedTernaryMapping] = []
            dst_names_seen: set[str] = set()
            for src_a_name in src_a_names:
                try:
                    src_b_name = re.sub(a_expr, b_expr, src_a_name)
                    dst_name = re.sub(a_expr, to_expr, src_a_name)
                except re.error as exc:
                    raise self.error_type(
                        f"{self.name} invalid regex rewrite from {a_expr!r}: {exc}"
                    ) from exc

                if src_b_name not in src_b_sd:
                    raise self.error_type(
                        f"{self.name} source_b missing: {src_b_model}::{src_b_name}"
                    )
                if dst_name in dst_names_seen:
                    raise self.error_type(f"{self.name} destination collision: {dst_model}::{dst_name}")
                dst_names_seen.add(dst_name)

                resolved.append(
                    ResolvedTernaryMapping(
                        src_a_model=src_a_model,
                        src_a_name=src_a_name,
                        src_a_slice=a_slice,
                        src_b_model=src_b_model,
                        src_b_name=src_b_name,
                        src_b_slice=b_slice,
                        dst_model=dst_model,
                        dst_name=dst_name,
                        dst_slice=dst_slice,
                    )
                )
            return resolved

        if isinstance(a_expr, list) and isinstance(b_expr, list) and isinstance(to_expr, list):
            resolved: list[ResolvedTernaryMapping] = []
            dst_names_seen: set[str] = set()
            matched_any = False
            for src_a_name in sorted(src_a_sd.keys()):
                match = _match_structured_expr(
                    expr=a_expr,
                    name=src_a_name,
                    op_name=self.name,
                    role="source_a",
                )
                if match is None:
                    continue
                matched_any = True

                src_b_name = _rewrite_structured_expr(
                    expr=b_expr,
                    match=match,
                    op_name=self.name,
                    role="source_b",
                )
                dst_name = _rewrite_structured_expr(
                    expr=to_expr,
                    match=match,
                    op_name=self.name,
                    role="destination",
                )

                if src_b_name not in src_b_sd:
                    raise self.error_type(
                        f"{self.name} source_b missing: {src_b_model}::{src_b_name}"
                    )
                if dst_name in dst_names_seen:
                    raise self.error_type(f"{self.name} destination collision: {dst_model}::{dst_name}")
                dst_names_seen.add(dst_name)

                resolved.append(
                    ResolvedTernaryMapping(
                        src_a_model=src_a_model,
                        src_a_name=src_a_name,
                        src_a_slice=a_slice,
                        src_b_model=src_b_model,
                        src_b_name=src_b_name,
                        src_b_slice=b_slice,
                        dst_model=dst_model,
                        dst_name=dst_name,
                        dst_slice=dst_slice,
                    )
                )

            if not matched_any:
                raise self.error_type(
                    f"{self.name} source_a matched zero tensors: "
                    f"{format_tensor_ref(from_a_ref)}; available tensors: {sorted(src_a_sd.keys())}"
                )
            return resolved

        raise self.error_type(
            f"{self.name} requires from_a/from_b/to expressions of the same kind: "
            "either all strings (regex mode) or all lists (structured mode)"
        )

    def validate_resolved_items(
        self,
        spec: TernarySpecT,
        mappings: list[ResolvedTernaryMapping],
        provider: StateDictProvider,
    ) -> None:
        del spec
        if self.destination_policy is DestinationPolicy.ANY:
            return

        for item in mappings:
            dst_sd = provider.get_state_dict(item.dst_model)
            exists = item.dst_name in dst_sd

            if self.destination_policy is DestinationPolicy.MUST_EXIST and not exists:
                raise self.error_type(
                    f"{self.name} destination missing: {item.dst_model}::{item.dst_name}"
                )
            if self.destination_policy is DestinationPolicy.MUST_NOT_EXIST and exists:
                raise self.error_type(
                    f"{self.name} destination already exists: {item.dst_model}::{item.dst_name}"
                )

    @abstractmethod
    def apply_mapping(self, item: ResolvedTernaryMapping, provider: StateDictProvider) -> None:
        ...

    def apply_item(
        self,
        spec: TernarySpecT,
        item: ResolvedTernaryMapping,
        provider: StateDictProvider,
    ) -> None:
        del spec
        self.apply_mapping(item, provider)


__all__ = [
    "BaseTransform",
    "BinaryMappingSpec",
    "BinaryMappingTransform",
    "CompiledTransform",
    "DestinationPolicy",
    "IteratingTransform",
    "ResolvedTernaryMapping",
    "TernaryMappingSpec",
    "TernaryMappingTransform",
    "TransformControl",
    "TransformResult",
    "TypedTransform",
    "UnarySpec",
    "UnaryTransform",
    "REGISTRY",
    "apply_transform",
    "get_transform",
    "list_transforms",
    "register_transform",
]
