"""In-memory TUI session-to-persona state management."""

from __future__ import annotations

from dataclasses import dataclass, field

from cleveragents.tui.persona.registry import PersonaRegistry
from cleveragents.tui.persona.schema import Persona


@dataclass(slots=True)
class PersonaState:
    """Tracks active persona per TUI session."""

    registry: PersonaRegistry
    active_by_session: dict[str, str] = field(default_factory=dict)
    preset_by_session: dict[str, str] = field(default_factory=dict)

    def _resolve_default_name(self) -> str:
        default = self.registry.ensure_default()
        return self.registry.get_last_persona() or default.name

    def active_name(self, session_id: str) -> str:
        """Return the active persona name for the given session.

        Initialises the session to the default persona on first access.
        """
        if session_id not in self.active_by_session:
            self.active_by_session[session_id] = self._resolve_default_name()
        return self.active_by_session[session_id]

    def active_persona(self, session_id: str) -> Persona:
        """Return the active Persona object for the given session.

        Falls back to the registry default if the stored name is no longer
        present in the registry.
        """
        name = self.active_name(session_id)
        persona = self.registry.get(name)
        if persona is None:
            persona = self.registry.ensure_default()
            self.active_by_session[session_id] = persona.name
        return persona

    def set_active_persona(self, session_id: str, persona_name: str) -> Persona:
        """Set the active persona for the given session and persist the choice.

        Raises ValueError if the persona name is not found in the registry.
        Initialises the session preset to "default" if not already set.
        """
        persona = self.registry.get(persona_name)
        if persona is None:
            raise ValueError(f"Unknown persona: {persona_name}")
        self.active_by_session[session_id] = persona.name
        self.registry.set_last_persona(persona.name)
        if session_id not in self.preset_by_session:
            self.preset_by_session[session_id] = "default"
        return persona

    def current_preset(self, session_id: str) -> str:
        """Return the current argument preset name for the given session.

        Defaults to "default" on first access.
        """
        if session_id not in self.preset_by_session:
            self.preset_by_session[session_id] = "default"
        return self.preset_by_session[session_id]

    def cycle_preset(self, session_id: str) -> str:
        persona = self.active_persona(session_id)
        if not persona.argument_presets:
            self.preset_by_session[session_id] = "default"
            return "default"
        names = [item.name for item in persona.argument_presets]
        current = self.current_preset(session_id)
        if current not in names:
            self.preset_by_session[session_id] = names[0]
            return names[0]
        idx = names.index(current)
        next_name = names[(idx + 1) % len(names)]
        self.preset_by_session[session_id] = next_name
        return next_name

    def cycle_persona(self, session_id: str) -> Persona:
        """Cycle to the next persona in cycle_order sequence.

        Only personas with cycle_order > 0 are included in the cycle.
        If no cyclic personas exist, returns the current active persona.
        """
        personas = self.registry.list_personas()
        cyclic = sorted(
            [p for p in personas if p.cycle_order > 0], key=lambda p: p.cycle_order
        )

        if not cyclic:
            return self.active_persona(session_id)

        current = self.active_name(session_id)
        current_names = [p.name for p in cyclic]

        if current not in current_names:
            # Current persona is not in cycle, start from first
            next_persona = cyclic[0]
        else:
            idx = current_names.index(current)
            next_persona = cyclic[(idx + 1) % len(cyclic)]

        return self.set_active_persona(session_id, next_persona.name)

    def effective_arguments(self, session_id: str) -> dict[str, object]:
        persona = self.active_persona(session_id)
        preset = self.current_preset(session_id)
        return persona.effective_arguments(preset)
