"""Production skeleton compression for ACMS child-context inheritance."""

from __future__ import annotations

import logging
import math
from pathlib import PurePosixPath
from typing import TYPE_CHECKING, Final, cast

from cleveragents.acms import (
    CODE_DETAIL_LEVEL_MAP,
    FUNC_DETAIL_LEVEL_MAP,
    JAVA_DETAIL_LEVELS,
    OO_DETAIL_LEVEL_MAP,
    PROC_DETAIL_LEVEL_MAP,
    PYTHON_DETAIL_LEVELS,
    RUST_DETAIL_LEVELS,
    TYPESCRIPT_DETAIL_LEVELS,
)
from cleveragents.application.services.acms_skeleton_renderers import (
    FragmentRenderVisitor,
    get_render_visitor,
)
from cleveragents.application.services.component_resolver import (
    ComponentRegistrationError,
    ComponentResolver,
)
from cleveragents.application.services.config_service import ConfigService
from cleveragents.domain.models.acms import DetailLevelMap
from cleveragents.domain.models.acms.ontology_registry import get_domain
from cleveragents.domain.models.core.context_fragment import ContextFragment

if TYPE_CHECKING:
    from cleveragents.application.services.acms_service import SkeletonCompressor

logger = logging.getLogger(__name__)

_DETAIL_DEPTH_OVERVIEW: Final[int] = 0
_DETAIL_DEPTH_SUMMARY: Final[int] = 1
_CODE_EXTENSIONS: Final[frozenset[str]] = frozenset(
    {".py", ".pyi", ".ts", ".tsx", ".js", ".jsx", ".java", ".rs"}
)
_DOC_EXTENSIONS: Final[frozenset[str]] = frozenset(
    {".md", ".mdx", ".rst", ".txt", ".adoc"}
)
_DATA_EXTENSIONS: Final[frozenset[str]] = frozenset({".sql", ".ddl", ".csv", ".tsv"})
_INFRA_EXTENSIONS: Final[frozenset[str]] = frozenset(
    {".yaml", ".yml", ".json", ".toml", ".tf", ".hcl"}
)
_PREFIX_FAMILY: Final[dict[str, str]] = {
    "uko-code:": "code",
    "uko-oo:": "code",
    "uko-func:": "code",
    "uko-proc:": "code",
    "uko-py:": "code",
    "uko-ts:": "code",
    "uko-java:": "code",
    "uko-rs:": "code",
    "uko-doc:": "doc",
    "uko-data:": "data",
    "uko-infra:": "infra",
    "uko:": "generic",
}


def _child_detail_map(
    domain: str,
    parent: DetailLevelMap,
    levels: tuple[tuple[str, int], ...],
) -> DetailLevelMap:
    """Create a child ``DetailLevelMap`` over an inherited parent chain."""

    depth_map = dict(levels)
    return DetailLevelMap(
        domain=domain,
        parent=parent,
        levels=depth_map,
        max_depth=max(depth_map.values()) if depth_map else parent.max_depth,
    )


_DETAIL_MAPS: Final[dict[str, DetailLevelMap]] = {
    "uko:": get_domain("uko:").detail_map,
    "uko-code:": CODE_DETAIL_LEVEL_MAP,
    "uko-doc:": get_domain("uko-doc:").detail_map,
    "uko-data:": get_domain("uko-data:").detail_map,
    "uko-infra:": get_domain("uko-infra:").detail_map,
    "uko-oo:": OO_DETAIL_LEVEL_MAP,
    "uko-func:": FUNC_DETAIL_LEVEL_MAP,
    "uko-proc:": PROC_DETAIL_LEVEL_MAP,
    "uko-py:": _child_detail_map("uko-py:", OO_DETAIL_LEVEL_MAP, PYTHON_DETAIL_LEVELS),
    "uko-ts:": _child_detail_map(
        "uko-ts:", OO_DETAIL_LEVEL_MAP, TYPESCRIPT_DETAIL_LEVELS
    ),
    "uko-java:": _child_detail_map(
        "uko-java:", OO_DETAIL_LEVEL_MAP, JAVA_DETAIL_LEVELS
    ),
    "uko-rs:": _child_detail_map("uko-rs:", OO_DETAIL_LEVEL_MAP, RUST_DETAIL_LEVELS),
}


class DepthReductionCompressor:
    """Compress parent fragments by re-rendering them at overview depths 0-1."""

    def compress(
        self,
        fragments: tuple[ContextFragment, ...],
        skeleton_budget: int,
    ) -> tuple[ContextFragment, ...]:
        """Return reduced fragments whose total token count fits the budget."""

        if skeleton_budget <= 0 or not fragments:
            return ()

        normalized = tuple(self._normalise_fragment(fragment) for fragment in fragments)
        ordered = sorted(
            normalized,
            key=lambda fragment: (
                -fragment.relevance_score,
                fragment.detail_depth,
                getattr(fragment, "fragment_id", fragment.uko_node),
            ),
        )

        compressed: list[ContextFragment] = []
        remaining = skeleton_budget
        for fragment in ordered:
            reduced = self._reduce_fragment(fragment, remaining)
            if reduced is None:
                continue
            compressed.append(reduced)
            remaining -= reduced.token_count
            if remaining <= 0:
                break

        logger.info(
            "Compressed skeleton fragments via depth reduction",
            extra={
                "input_count": len(fragments),
                "output_count": len(compressed),
                "skeleton_budget": skeleton_budget,
                "used_tokens": sum(fragment.token_count for fragment in compressed),
            },
        )
        return tuple(compressed)

    @staticmethod
    def _normalise_fragment(fragment: ContextFragment) -> ContextFragment:
        detail_map = _detail_map_for_fragment(fragment)
        metadata = dict(fragment.metadata)
        if metadata.get("detail_level_domain") == detail_map.domain:
            return fragment
        metadata["detail_level_domain"] = detail_map.domain
        return fragment.model_copy(update={"metadata": metadata})

    def _reduce_fragment(
        self,
        fragment: ContextFragment,
        remaining_budget: int,
    ) -> ContextFragment | None:
        if remaining_budget <= 0:
            return None

        detail_map = _detail_map_for_fragment(fragment)
        family = _domain_family(fragment, detail_map)
        visitor = get_render_visitor(family)

        candidate_depths = [
            depth
            for depth in (_DETAIL_DEPTH_SUMMARY, _DETAIL_DEPTH_OVERVIEW)
            if depth <= min(fragment.detail_depth, detail_map.max_depth)
        ]
        if not candidate_depths:
            candidate_depths = [_DETAIL_DEPTH_OVERVIEW]

        seen_depths: set[int] = set()
        fallback_candidate: ContextFragment | None = None
        for depth in candidate_depths:
            if depth in seen_depths:
                continue
            seen_depths.add(depth)
            candidate = self._rerender_fragment(
                fragment,
                detail_map=detail_map,
                visitor=visitor,
                target_depth=depth,
            )
            if candidate.token_count <= remaining_budget:
                return candidate
            fallback_candidate = candidate
        if fallback_candidate is None:
            return None
        return self._fit_candidate(fallback_candidate, remaining_budget)

    def _rerender_fragment(
        self,
        fragment: ContextFragment,
        *,
        detail_map: DetailLevelMap,
        visitor: FragmentRenderVisitor,
        target_depth: int,
    ) -> ContextFragment:
        level_name = _level_name_for_depth(detail_map, target_depth)
        if fragment.detail_depth <= target_depth:
            content = fragment.content
        else:
            content = visitor.render(
                fragment,
                target_depth=target_depth,
                level_name=level_name,
            )
        token_count = _estimate_tokens(content)
        metadata = {
            **fragment.metadata,
            "detail_level_domain": detail_map.domain,
            "skeleton_target_depth": str(target_depth),
            "skeleton_target_level": level_name,
        }
        return fragment.model_copy(
            update={
                "content": content,
                "detail_depth": target_depth,
                "token_count": token_count,
                "metadata": metadata,
            }
        )

    @staticmethod
    def _fit_candidate(
        fragment: ContextFragment,
        remaining_budget: int,
    ) -> ContextFragment | None:
        if fragment.token_count <= remaining_budget:
            return fragment
        if remaining_budget <= 0:
            return None
        clipped = _clip_to_token_budget(fragment.content, remaining_budget)
        if not clipped:
            return None
        clipped_tokens = _estimate_tokens(clipped)
        if clipped_tokens > remaining_budget:
            return None
        return fragment.model_copy(
            update={
                "content": clipped,
                "token_count": clipped_tokens,
            }
        )


def resolve_configured_skeleton_compressor() -> SkeletonCompressor:
    """Resolve the configured skeleton compressor via the component resolver."""

    resolver = ComponentResolver()
    builtin = DepthReductionCompressor()
    from cleveragents.application.services.acms_service import SkeletonCompressor

    resolver.register_global(SkeletonCompressor, builtin)
    configured = ConfigService().resolve("context.pipeline.skeleton-compressor").value
    if (
        not isinstance(configured, str)
        or configured == "builtin:DepthReductionCompressor"
    ):
        return cast(SkeletonCompressor, resolver.resolve(SkeletonCompressor).component)

    instance = resolver._import_component(configured)
    compress = getattr(instance, "compress", None)
    if not callable(compress):
        msg = (
            f"Configured skeleton compressor {configured!r} does not expose a callable "
            "compress() method"
        )
        raise ComponentRegistrationError(msg)

    resolver.register_global(SkeletonCompressor, cast(SkeletonCompressor, instance))
    return cast(SkeletonCompressor, resolver.resolve(SkeletonCompressor).component)


def _detail_map_for_fragment(fragment: ContextFragment) -> DetailLevelMap:
    """Infer the most specific ``DetailLevelMap`` chain available for a fragment."""

    metadata = fragment.metadata
    for key in ("detail_level_domain", "detail_map_domain", "uko_domain"):
        value = metadata.get(key)
        if isinstance(value, str) and value in _DETAIL_MAPS:
            return _DETAIL_MAPS[value]

    for prefix in _DETAIL_MAPS:
        if fragment.uko_node.startswith(prefix):
            return _DETAIL_MAPS[prefix]
        if fragment.provenance.resource_uri.startswith(prefix):
            return _DETAIL_MAPS[prefix]

    resource_uri = fragment.provenance.resource_uri.lower()
    resource_type = getattr(fragment.provenance, "resource_type", "").lower()
    extension = PurePosixPath(resource_uri.split("?", 1)[0]).suffix.lower()

    if extension in _CODE_EXTENSIONS:
        if extension in {".py", ".pyi"}:
            return _DETAIL_MAPS["uko-py:"]
        if extension in {".ts", ".tsx", ".js", ".jsx"}:
            return _DETAIL_MAPS["uko-ts:"]
        if extension == ".java":
            return _DETAIL_MAPS["uko-java:"]
        if extension == ".rs":
            return _DETAIL_MAPS["uko-rs:"]
        return _DETAIL_MAPS["uko-code:"]
    if extension in _DOC_EXTENSIONS:
        return _DETAIL_MAPS["uko-doc:"]
    if extension in _DATA_EXTENSIONS:
        return _DETAIL_MAPS["uko-data:"]
    if extension in _INFRA_EXTENSIONS:
        return _DETAIL_MAPS["uko-infra:"]

    if any(
        token in resource_type
        for token in ("python", "typescript", "javascript", "java", "rust", "code")
    ):
        if "python" in resource_type:
            return _DETAIL_MAPS["uko-py:"]
        if "typescript" in resource_type or "javascript" in resource_type:
            return _DETAIL_MAPS["uko-ts:"]
        if "java" in resource_type:
            return _DETAIL_MAPS["uko-java:"]
        if "rust" in resource_type:
            return _DETAIL_MAPS["uko-rs:"]
        return _DETAIL_MAPS["uko-code:"]
    if any(token in resource_type for token in ("doc", "markdown", "text")):
        return _DETAIL_MAPS["uko-doc:"]
    if any(token in resource_type for token in ("sql", "database", "schema", "table")):
        return _DETAIL_MAPS["uko-data:"]
    if any(
        token in resource_type
        for token in ("infra", "config", "deploy", "kubernetes", "terraform")
    ):
        return _DETAIL_MAPS["uko-infra:"]
    return _DETAIL_MAPS["uko:"]


def _domain_family(fragment: ContextFragment, detail_map: DetailLevelMap) -> str:
    """Return the rendering family for a fragment."""

    explicit = fragment.metadata.get("detail_level_domain")
    if isinstance(explicit, str) and explicit in _PREFIX_FAMILY:
        return _PREFIX_FAMILY[explicit]
    return _PREFIX_FAMILY.get(detail_map.domain, "generic")


def _level_name_for_depth(detail_map: DetailLevelMap, depth: int) -> str:
    """Resolve the effective named level for a target depth."""

    matches = [
        name
        for name, mapped_depth in sorted(
            detail_map.effective_levels().items(),
            key=lambda item: (item[1], item[0]),
        )
        if mapped_depth == depth
    ]
    return matches[0] if matches else f"depth-{depth}"


def _estimate_tokens(text: str) -> int:
    """Estimate token count for reduced fragments conservatively."""

    if not text:
        return 1
    return max(1, math.ceil(len(text) / 4))


def _clip_to_token_budget(text: str, token_budget: int) -> str:
    """Clip text so the rendered content itself fits the token budget."""

    if token_budget <= 0:
        return ""
    return text[: max(1, token_budget * 4)].rstrip()
