from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional, Set, Type

from pydantic import BaseModel

from .entries import RpcBundle
from .typing_utils import extract_nested_models, is_pydantic_model, apply_aliases


@dataclass
class RpcRegistry:
    """
    Central registry that collects RPC bundles, applies model normalization
    (Pydantic field aliasing), and exposes a unified bundle.

    - Plugins and decorators add RpcBundle instances.
    - Registry merges them.
    - Registry applies camelCase aliases to all Pydantic models exactly once.
    """

    service_name: str = "VentionApp"
    _bundles: List[RpcBundle] = field(default_factory=list)
    _models_normalized: bool = False

    # ------------- Bundle registration -------------

    def add_bundle(self, bundle: RpcBundle) -> None:
        """Register a bundle for inclusion in the unified RPC view."""
        self._bundles.append(bundle)

    @property
    def bundle(self) -> RpcBundle:
        """Return a merged RpcBundle (does not mutate stored bundles)."""
        merged = RpcBundle()
        for bundle in self._bundles:
            merged.extend(bundle)
        return merged

    # ------------- Model normalization / aliasing -------------
    def _normalize_model(
        self, model: Optional[Type[BaseModel]], seen: Set[Type[BaseModel]]
    ) -> None:
        """Recursively normalize a model and all its nested models."""
        if model is None:
            return
        if not is_pydantic_model(model):
            return
        if model in seen:
            return

        seen.add(model)
        apply_aliases(model)

        if hasattr(model, "model_fields"):
            for _, field_info in model.model_fields.items():
                nested_models = extract_nested_models(field_info.annotation)
                for nested_model in nested_models:
                    self._normalize_model(nested_model, seen)

    def normalize_models_and_apply_aliases(self) -> None:
        """
        Walk all RPCs in all bundles and apply camelCase JSON aliases
        to every Pydantic model exactly once, including nested models.
        """
        if self._models_normalized:
            return

        seen: Set[Type[BaseModel]] = set()

        for bundle in self._bundles:
            for action in bundle.actions:
                self._normalize_model(action.input_type, seen)
                self._normalize_model(action.output_type, seen)
            for stream in bundle.streams:
                self._normalize_model(stream.payload_type, seen)

        self._models_normalized = True

    # ------------- Unified, normalized view -------------

    def get_unified_bundle(self) -> RpcBundle:
        """
        Get the fully merged, normalized RPC bundle.

        This will apply aliasing exactly once and then return a merged RpcBundle.
        """
        self.normalize_models_and_apply_aliases()
        return self.bundle
