from typing import Any, Iterable, NamedTuple, TypeVar, Union

import msgspec
from jsonpointer import JsonPointer
from msgspec import UNSET, Struct, UnsetType, defstruct

from autocrud.util.type_utils import (
    get_list_item_type,
    get_union_args,
    is_list_type,
    is_struct_type,
    is_union_type,
    resolve_struct_origin,
)

T = TypeVar("T")


def _normalize_paths(partial: Iterable[str | JsonPointer]) -> list[list[str]]:
    paths = []
    for p in partial:
        if isinstance(p, JsonPointer):
            parts = p.parts
        else:
            parts = p.split("/")

        parts = [x for x in parts if x]
        paths.append(parts)
    return paths


def _merge_paths(paths: list[list[str]]) -> dict[str, list[list[str]]]:
    grouped = {}
    for p in paths:
        if not p:
            continue
        head = p[0]
        tail = p[1:]
        if head not in grouped:
            grouped[head] = []
        grouped[head].append(tail)
    return grouped


def _get_struct_fields(struct_type: Any) -> dict[str, Any]:
    return {f.name: f for f in msgspec.structs.fields(struct_type)}


def create_partial_type(base_type: Any, partial: Iterable[str | JsonPointer]) -> Any:
    paths = _normalize_paths(partial)
    return _build_type(base_type, paths, prefix="Partial")


def _build_type(current_type: Any, paths: list[list[str]], prefix: str) -> Any:
    if any(len(p) == 0 for p in paths):
        return current_type

    if is_union_type(current_type):
        new_args = []
        for arg in get_union_args(current_type):
            if arg is type(None):
                new_args.append(arg)
            else:
                new_args.append(_build_type(arg, paths, prefix))

        if len(new_args) == 2 and type(None) in new_args:
            other = new_args[0] if new_args[1] is type(None) else new_args[1]
            return other | None
        return Union[tuple(new_args)]

    if is_list_type(current_type):
        sub_paths = []
        for p in paths:
            if len(p) > 0:
                sub_paths.append(p[1:])

        element_type = get_list_item_type(current_type)
        new_element_type = _build_type(element_type, sub_paths, prefix + "Item")
        return list[new_element_type]

    if is_struct_type(current_type):
        fields = _get_struct_fields(current_type)
        grouped_paths = _merge_paths(paths)

        new_fields = []
        for field_name, sub_paths in grouped_paths.items():
            if field_name in fields:
                field_def = fields[field_name]
                field_type = field_def.type

                new_field_type = _build_type(
                    field_type, sub_paths, f"{prefix}_{field_name}"
                )

                # Add UnsetType to allow pruning
                if is_union_type(new_field_type):
                    u_args = get_union_args(new_field_type)
                    if UnsetType not in u_args:
                        new_field_type = Union[tuple(list(u_args) + [UnsetType])]
                else:
                    if new_field_type is not UnsetType:
                        new_field_type = new_field_type | UnsetType

                new_fields.append(
                    (
                        field_name,
                        new_field_type,
                        msgspec.field(
                            default=field_def.default,
                            default_factory=field_def.default_factory,
                            name=field_def.name,
                        ),
                    )
                )

        # For generic aliases (e.g. Job[Payload]), resolve the origin for __name__
        struct_class = resolve_struct_origin(current_type)
        type_name = f"{prefix}_{struct_class.__name__}"

        # Keep tag info
        kwargs = {}
        try:
            info = msgspec.inspect.type_info(current_type)
            kwargs["tag"] = info.tag
            kwargs["tag_field"] = info.tag_field
            kwargs["array_like"] = info.array_like
            kwargs["forbid_unknown_fields"] = info.forbid_unknown_fields
        except Exception:
            pass

        return defstruct(type_name, new_fields, kw_only=True, **kwargs)

    return current_type


def prune_object(obj: Any, partial: Iterable[str | JsonPointer]) -> Any:
    paths = _normalize_paths(partial)
    if not _needs_pruning(paths):
        return obj
    return _prune(obj, paths)


def _needs_pruning(paths: list[list[str]]) -> bool:
    for path in paths:
        for part in path:
            # Check if part is a specific index (digit)
            if part.isdigit() or (part.startswith("-") and part[1:].isdigit()):
                return True

            # Check if part is a slice
            if ":" in part:
                s = _parse_slice(part)
                if s:
                    # If it's not a full slice, we need pruning
                    if not (
                        s.start is None
                        and s.stop is None
                        and (s.step is None or s.step == 1)
                    ):
                        return True
    return False


def _parse_slice(s: str) -> slice | None:
    if ":" not in s:
        return None
    parts = s.split(":")
    if len(parts) > 3:
        return None

    try:
        start = int(parts[0]) if parts[0] else None
        stop = int(parts[1]) if len(parts) > 1 and parts[1] else None
        step = int(parts[2]) if len(parts) > 2 and parts[2] else None
        return slice(start, stop, step)
    except ValueError:
        return None


def _prune(obj: Any, paths: list[list[str]]) -> Any:
    if any(len(p) == 0 for p in paths):
        return obj

    if isinstance(obj, list):
        wildcard_subpaths = [p[1:] for p in paths if p and p[0] == "-"]
        has_wildcard = len(wildcard_subpaths) > 0

        other_paths = [p for p in paths if p and p[0] != "-"]

        if not other_paths:
            if has_wildcard:
                return obj
            return []

        obj_len = len(obj)
        slices = []
        specific_indices = {}

        for p in other_paths:
            head = p[0]
            s = _parse_slice(head)
            if s:
                slices.append((s, p[1:]))
            else:
                try:
                    idx = int(head)
                    if idx < 0:
                        idx += obj_len
                    if 0 <= idx < obj_len:
                        if idx not in specific_indices:
                            specific_indices[idx] = []
                        specific_indices[idx].append(p[1:])
                except ValueError:
                    pass

        # Optimization: If all paths are universal (wildcard or full slice) and no specific indices,
        # then the object structure already matches the request (handled by create_partial_type),
        # so we can return it as is.
        if not specific_indices:
            all_slices_full = True
            for s, _ in slices:
                if not (
                    s.start is None
                    and s.stop is None
                    and (s.step is None or s.step == 1)
                ):
                    all_slices_full = False
                    break

            if all_slices_full:
                return obj

        if has_wildcard:
            indices_to_process = range(obj_len)
        else:
            indices = set(specific_indices.keys())
            for s, _ in slices:
                indices.update(range(*s.indices(obj_len)))
            indices_to_process = sorted(indices)

        new_list = []
        for i in indices_to_process:
            item = obj[i]

            current_subpaths = list(wildcard_subpaths)

            if i in specific_indices:
                current_subpaths.extend(specific_indices[i])

            for s, sub in slices:
                start, stop, step = s.indices(obj_len)
                in_slice = False
                if step > 0:
                    if start <= i < stop and (i - start) % step == 0:
                        in_slice = True
                else:
                    if start >= i > stop and (start - i) % (-step) == 0:
                        in_slice = True

                if in_slice:
                    current_subpaths.append(sub)

            new_list.append(_prune(item, current_subpaths))

        return new_list

    if isinstance(obj, Struct):
        grouped = _merge_paths(paths)
        changes = {}

        fields = msgspec.structs.fields(type(obj))
        for f in fields:
            field_name = f.name
            val = getattr(obj, field_name)

            if field_name in grouped:
                if val is not UNSET:
                    sub_paths = grouped[field_name]
                    pruned_val = _prune(val, sub_paths)
                    if pruned_val is not val:
                        changes[field_name] = pruned_val
            else:
                if val is not UNSET:
                    changes[field_name] = UNSET

        if changes:
            return msgspec.structs.replace(obj, **changes)
        return obj

    return obj


# ---------------------------------------------------------------------------
# Partial field classification & filtering utilities
# ---------------------------------------------------------------------------


class PartialFieldsSpec(NamedTuple):
    """Classified partial fields by prefix (meta/info/data)."""

    data_fields: list[str] | None = None
    meta_fields: list[str] | None = None
    info_fields: list[str] | None = None


def classify_partial_fields(
    fields: list[str] | None,
    default_category: str = "data",
) -> PartialFieldsSpec:
    """Classify partial field paths by prefix.

    Fields prefixed with ``meta/``, ``info/``, or ``data/`` are routed to the
    corresponding bucket.  Fields without a recognised prefix fall into
    *default_category* (``"data"`` by default, ``"meta"`` for meta-only
    endpoints, etc.).

    Args:
        fields: raw partial field list from the query string.
        default_category: where to place unprefixed fields
            (``"data"`` | ``"meta"`` | ``"info"``).

    Returns:
        A ``PartialFieldsSpec`` with each bucket set to ``None`` when empty.
    """
    if not fields:
        return PartialFieldsSpec()

    data_fields: list[str] = []
    meta_fields: list[str] = []
    info_fields: list[str] = []

    for raw in fields:
        # Normalise: strip leading slash
        stripped = raw.lstrip("/")

        if stripped.startswith("meta/"):
            meta_fields.append("/" + stripped[len("meta/") :])
        elif stripped.startswith("info/"):
            info_fields.append("/" + stripped[len("info/") :])
        elif stripped.startswith("data/"):
            data_fields.append("/" + stripped[len("data/") :])
        else:
            # No recognised prefix → default category
            path = raw if raw.startswith("/") else "/" + raw
            if default_category == "meta":
                meta_fields.append(path)
            elif default_category == "info":
                info_fields.append(path)
            else:
                data_fields.append(path)

    return PartialFieldsSpec(
        data_fields=data_fields or None,
        meta_fields=meta_fields or None,
        info_fields=info_fields or None,
    )


def filter_struct_partial(
    struct: Struct,
    fields: list[str],
) -> Struct:
    """Return a copy of *struct* keeping only the requested fields.

    Uses ``create_partial_type`` to build a lightweight Struct type that
    contains only the selected fields, then round-trips through msgspec
    JSON encode/decode for the actual filtering.
    """
    partial_type = create_partial_type(type(struct), fields)
    return msgspec.json.decode(msgspec.json.encode(struct), type=partial_type)
