from dataclasses import dataclass, field
from functools import partial
from typing import Any, Callable, Dict, List, Literal, Optional
from uuid import UUID

from ..utils.constants import DEFAULT_LAYOUT_TOKEN, REST_TOKEN
from .routing import PageNode, RouteTable
from pydantic import BaseModel


class LoadingState(BaseModel):
    state: Literal["lacy", "done", "hidden"]
    node_id: str
    updated: bool = False

    def update_state(self, state: Literal["lacy", "done", "hidden"]) -> None:
        self.state = state
        self.updated = True


@dataclass
class RoutingContext:
    """Encapsulates all routing state for a single request"""

    pathname: str
    query_params: Dict[str, Any]
    resolve_type: Literal["search", "url", "lacy"]
    path_vars: Dict[str, str] = field(default_factory=dict)
    endpoints: Dict[UUID, Callable] = field(default_factory=dict)
    segments: List[str] = field(default_factory=list)
    _loading_states: Dict[str, Any] = field(default_factory=dict, repr=False)

    @property
    def variables(self):
        return {**self.query_params, **self.path_vars}

    @classmethod
    def from_request(
        cls,
        pathname: str,
        query_params: Dict[str, Any],
        loading_state_dict: Dict[str, Dict],
        resolve_type: Literal["search", "url", "lacy"],
    ) -> "RoutingContext":
        """Create context from request data"""
        path = pathname.strip("/")
        segments = [seg for seg in path.split("/") if seg] if path else []
        # Build _loading_states dict
        _loading_states = {
            segment_key: LoadingState(**ils)
            for segment_key, ils in loading_state_dict.items()
        }
        return cls(
            pathname=pathname,
            query_params=query_params,
            segments=segments,
            resolve_type=resolve_type,
            _loading_states=_loading_states,
        )

    def get_node_state(self, segment_key: Optional[str] = None) -> Optional[str]:
        """Get loading state for a node"""
        ls = self._loading_states.get(segment_key)
        if ls:
            return ls.state
        return None

    def set_node_state(
        self, node: PageNode, state: str, segment_key: Optional[str] = None
    ):
        """Set loading state for a node"""
        if segment_key not in self._loading_states:
            self._loading_states[segment_key] = LoadingState(
                state=state, node_id=node.node_id, updated=True
            )
        else:
            self._loading_states[segment_key].update_state(state)

    def add_endpoint(self, node: PageNode):
        partial_endpoint = partial(node.endpoint, **self.variables)
        self.endpoints[node.node_id] = partial_endpoint

    def should_lazy_load(
        self, node: PageNode, segment_key: Optional[str] = None
    ) -> bool:
        """Check if node should be lazy loaded"""
        return (
            node.loading is not None
            and self.resolve_type != "lacy"
            and DEFAULT_LAYOUT_TOKEN not in segment_key
        )

    def pop_segment(self) -> Optional[str]:
        """Remove and return the last segment"""
        return self.segments.pop() if self.segments else None

    def peek_segment(self) -> Optional[str]:
        """Peek at the last segment without removing"""
        return self.segments[-1] if self.segments else None

    def consume_path_var(self, node: PageNode) -> Optional[str]:
        """Consume a segment as path variable"""
        if not node.is_path_template:
            return self.peek_segment()
        
        if node.segment == REST_TOKEN:
            rest_value = list(reversed(self.segments))
            self.segments = []
            self.path_vars["rest"] = rest_value
            return None
        else:
            value = self.pop_segment()
            if value:
                self.path_vars[node.segment] = value
            return value

    def merge_segments(self, ignore_empty_folders: bool):
        """Merge segments if empty folder should be ignored"""
        if ignore_empty_folders:
            self.pop_segment()
            return

        if len(self.segments) < 2:
            self.pop_segment()
            return

        first = self.segments.pop()
        second = self.segments.pop()
        combined = f"{first}/{second}"
        self.segments.append(combined)

    def gather_endpoints(self):
        if not self.endpoints:
            return {}

        keys = list(self.endpoints.keys())
        funcs = list(self.endpoints.values())
        results = []
        for func in funcs:
            try:
                result = func()
                results.append(result)
            except Exception as e:
                results.append(e)
        return dict(zip(keys, results))

    def to_loading_state_dict(self) -> Dict[str, Any]:
        """Convert context back to loading state dict for response"""
        return {**self.get_updated_loading_state(), "query_params": self.query_params}

    def get_updated_loading_state(self) -> Dict[str, Dict]:
        """Return only updated loading states as a dict."""
        return {
            key: {"state": state.state, "node_id": state.node_id}
            for key, state in self._loading_states.items()
            if state.updated == True
        }

    def set_silent_loading_states(self, node: PageNode, state: str = "done"):
        """Mark all descendant slots as done"""
        for slot_name, slot_id in node.slots.items():
            slot_node = RouteTable.get_node(slot_id)
            if slot_node:
                self.set_node_state(slot_node, state, slot_name)
                self.set_silent_loading_states(slot_node, state)
