"""Training Shell — MCP wrapper for algebraic routing table.

Exposes TrainingLoop routing table, episode dispatch, agent stats,
and best resolver through the standard shell run() interface.

The routing table is NOT learned. It IS the I-Logic multiplication table.
Computed once in __init__. Pure algebra, deterministic.
"""

from typing import Any, Dict, List, Optional

from ..swarm import Swarm
from ..training import TrainingLoop
from ..i_logic.states import ILogicState


class TrainingShell:
    """MCP shell wrapper for algebraic routing."""

    HEX_ID = "54524E00"  # TRN\0
    name = "training"
    version = "1.0.0"

    def __init__(self):
        self._swarm = Swarm()
        self._loop = TrainingLoop(self._swarm)

    # ------------------------------------------------------------------
    # COMMANDS
    # ------------------------------------------------------------------

    def status(self) -> Dict[str, Any]:
        """Training status and routing summary."""
        return self._loop.summary()

    def table(self) -> Dict[str, Any]:
        """Show the algebraic routing table.

        Computed from I-Logic multiplication: condition × agent = DIRECT.
        No training needed — the answer is the algebra.
        """
        return {
            "description": "Algebraic routing: condition × agent_state = DIRECT",
            "routing_table": self._loop.get_routing_table(),
        }

    def resolve(self, state_name: str = "INVERSE") -> Dict[str, Any]:
        """Find the best agent for a condition state.

        Args:
            state_name: I-Logic state (DIRECT, INVERSE, SPIRAL, ANTI_SPIRAL).
        """
        state = ILogicState.from_name(state_name.upper())
        agent_id = self._loop.get_best_resolver(state)
        if agent_id is not None:
            agent = self._swarm.get_agent(agent_id)
            if agent:
                return {
                    "condition": state_name.upper(),
                    "best_resolver": agent.name,
                    "agent_id": agent_id,
                    "agent_state": ILogicState.name(agent.i_logic_state),
                    "algebra": f"{state_name.upper()} × {ILogicState.name(agent.i_logic_state)} = DIRECT",
                }
        return {"condition": state_name.upper(), "best_resolver": None}

    def episode(self, *tasks: str) -> Dict[str, Any]:
        """Run an episode: batch-dispatch tasks and record transitions.

        Args:
            tasks: Task descriptions (comma-separated or multiple args).
        """
        task_list = list(tasks)
        if len(task_list) == 1 and "," in task_list[0]:
            task_list = [t.strip() for t in task_list[0].split(",")]
        if not task_list:
            task_list = ["validate system", "find documentation", "fix authentication"]

        ep = self._loop.run_episode(task_list)
        return {
            "episode_id": ep.episode_id,
            "tasks": len(ep.transitions),
            "resolution_rate": round(ep.resolution_rate, 3),
            "total_reward": round(ep.total_reward, 3),
            "coherence": round(ep.coherence, 3),
            "duration_ms": round(ep.duration_ms, 3),
            "transitions": [
                {
                    "task": t.task,
                    "agent_id": t.agent_id,
                    "task_state": ILogicState.name(t.task_state),
                    "outcome": ILogicState.name(t.outcome),
                    "resolved": t.resolved,
                    "chained": t.chained,
                    "reward": round(t.reward, 3),
                }
                for t in ep.transitions
            ],
        }

    def agent_stats(self) -> Dict[str, Any]:
        """Per-agent performance statistics."""
        return self._loop.get_agent_stats()

    def history(self) -> Dict[str, Any]:
        """Episode history summary."""
        return {
            "total_episodes": self._loop.total_episodes,
            "total_transitions": self._loop.total_transitions,
            "overall_resolution_rate": round(self._loop.overall_resolution_rate, 3),
            "overall_coherence": round(self._loop.overall_coherence, 3),
            "episodes": self._loop.get_episode_history(),
        }

    # ------------------------------------------------------------------
    # SHELL INTERFACE
    # ------------------------------------------------------------------

    def run(self, command: str, *args, as_json: bool = False) -> Any:
        """Execute a training command.

        Commands:
            status      — Full training summary
            table       — Algebraic routing table
            resolve     — Find best agent for a condition
            episode     — Run an episode with tasks
            agent_stats — Per-agent statistics
            history     — Episode history
        """
        handlers = {
            "status": lambda: self.status(),
            "table": lambda: self.table(),
            "resolve": lambda: self.resolve(args[0]) if args else self.resolve("INVERSE"),
            "episode": lambda: self.episode(*args) if args else self.episode(),
            "agent_stats": lambda: self.agent_stats(),
            "history": lambda: self.history(),
        }
        handler = handlers.get(command, handlers["status"])
        return handler()
