"""Agent Swarm -- 19 Agents with I-Logic Routing

Each agent occupies a position in Morton space. Its behavior is
determined by its I-Logic state. Tasks are routed via algebraic
composition -- multiply(task_state, agent_state) predicts the outcome
BEFORE execution.

The 19 agents map to the 12-layer architecture:
    L1-L9: 2 agents per internal layer (sensor + actor)
    L10:   1 agent (Identifier -- self-awareness)
    = 19 total

ROUTING PRINCIPLE:
    Task arrives with an I-Logic state (condition).
    The swarm finds the agent whose state RESOLVES the condition:
        condition x agent_state = DIRECT (resolution)
    If no single agent resolves it, compose a multi-agent path.

USAGE:
    from cubesquare.swarm import Swarm

    swarm = Swarm()
    result = swarm.dispatch("authentication module needs refactoring")
    print(result.agent.name, result.i_logic_outcome)

    # Multi-agent workflow
    workflow = swarm.compose_workflow(["navigate", "generate", "validate"])
    print(workflow.path, workflow.final_state)

NO PYTHON LOOPS in routing hot path.
All I-Logic operations are O(1) complex multiplication.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
import math

from ._core.constants import PRIME_131, LAYER_BASINS, LAYER_NAMES, K5, K11
from ._core import morton_encode, morton_decode, compute_basin
from .i_logic.states import ILogicState
from .i_logic.algebra import multiply, compose


# =========================================================================
# AGENT DEFINITIONS
# =========================================================================

@dataclass
class AgentRole:
    """Definition of an agent's role in the swarm."""
    agent_id: int
    name: str
    layer: int
    layer_name: str
    role: str               # sensor or actor
    specialization: str     # what this agent does
    i_logic_state: complex  # behavioral state
    home_basin: int         # center of this agent's basin range
    basin_range: Tuple[int, int]  # (start, end) basin range
    keywords: List[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "id": self.agent_id,
            "name": self.name,
            "layer": self.layer,
            "layer_name": self.layer_name,
            "role": self.role,
            "specialization": self.specialization,
            "i_logic_state": ILogicState.name(self.i_logic_state),
            "home_basin": self.home_basin,
            "basin_range": self.basin_range,
            "keywords": self.keywords,
        }


# The 19 agents: 2 per internal layer (sensor + actor) + 1 external
AGENT_DEFINITIONS: List[dict] = [
    # L1 GEOMETRIC (basins 0-14) -- Structure
    {"id": 1,  "name": "Navigator",     "layer": 1, "role": "sensor",
     "spec": "Finds structures, tree traversal, file discovery",
     "state": ILogicState.DIRECT,
     "keywords": ["find", "search", "locate", "navigate", "explore", "tree", "path"]},
    {"id": 2,  "name": "Architect",     "layer": 1, "role": "actor",
     "spec": "Creates structures, scaffolding, organization",
     "state": ILogicState.DIRECT,
     "keywords": ["create", "structure", "scaffold", "organize", "build", "init"]},

    # L2 SEMANTIC (basins 15-29) -- Meaning
    {"id": 3,  "name": "Semanticist",   "layer": 2, "role": "sensor",
     "spec": "Interprets meaning, clusters concepts, understands context",
     "state": ILogicState.DIRECT,
     "keywords": ["meaning", "concept", "understand", "interpret", "classify"]},
    {"id": 4,  "name": "Translator",    "layer": 2, "role": "actor",
     "spec": "Converts between representations, adapts formats",
     "state": ILogicState.DIRECT,
     "keywords": ["convert", "translate", "transform", "adapt", "format"]},

    # L3 SPATIAL (basins 30-43) -- Location
    {"id": 5,  "name": "Locator",       "layer": 3, "role": "sensor",
     "spec": "Finds positions, spatial queries, proximity detection",
     "state": ILogicState.SPIRAL,
     "keywords": ["where", "position", "location", "nearby", "proximity", "distance"]},
    {"id": 6,  "name": "Coordinator",   "layer": 3, "role": "actor",
     "spec": "Manages spatial relations, coordinates agents",
     "state": ILogicState.SPIRAL,
     "keywords": ["coordinate", "arrange", "place", "distribute", "map"]},

    # L4 FUNCTIONAL (basins 44-58) -- Operation
    {"id": 7,  "name": "Generator",     "layer": 4, "role": "actor",
     "spec": "Creates operations, code generation, implementation",
     "state": ILogicState.DIRECT,
     "keywords": ["generate", "implement", "code", "write", "produce", "make"]},
    {"id": 8,  "name": "Refactor",      "layer": 4, "role": "actor",
     "spec": "Transforms, abstracts, summarizes, simplifies",
     "state": ILogicState.DIRECT,
     "keywords": ["refactor", "simplify", "abstract", "extract", "optimize"]},

    # L5 PROBABILISTIC (basins 59-72) -- Uncertainty
    {"id": 9,  "name": "Validator",     "layer": 5, "role": "sensor",
     "spec": "Tests, verifies, samples, measures confidence",
     "state": ILogicState.INVERSE,
     "keywords": ["test", "verify", "validate", "check", "assert", "prove"]},
    {"id": 10, "name": "Balancer",      "layer": 5, "role": "actor",
     "spec": "Load balancing, Pareto optimization, resource allocation",
     "state": ILogicState.INVERSE,
     "keywords": ["balance", "distribute", "allocate", "optimize", "scale"]},

    # L6 EMOTIONAL (basins 73-87) -- State
    {"id": 11, "name": "Healer",        "layer": 6, "role": "actor",
     "spec": "Repairs states, resolves errors, restores coherence",
     "state": ILogicState.INVERSE,
     "keywords": ["fix", "repair", "heal", "resolve", "restore", "debug"]},
    {"id": 12, "name": "Monitor",       "layer": 6, "role": "sensor",
     "spec": "Observes state, alerts on anomalies, tracks health",
     "state": ILogicState.INVERSE,
     "keywords": ["monitor", "watch", "observe", "track", "alert", "status"]},

    # L7 FREQUENCY (basins 88-101) -- Timing
    {"id": 13, "name": "Scheduler",     "layer": 7, "role": "actor",
     "spec": "Temporal management, sequencing, scheduling",
     "state": ILogicState.SPIRAL,
     "keywords": ["schedule", "sequence", "order", "plan", "timeline", "when"]},
    {"id": 14, "name": "Synchronizer",  "layer": 7, "role": "actor",
     "spec": "Aligns frequencies, beat sync, phase coordination",
     "state": ILogicState.SPIRAL,
     "keywords": ["sync", "align", "phase", "frequency", "rhythm", "timing"]},

    # L8 ELECTROMAGNETIC (basins 102-116) -- Force
    {"id": 15, "name": "Stabilizer",    "layer": 8, "role": "actor",
     "spec": "Maintains coherence, dampens oscillation, error correction",
     "state": ILogicState.INVERSE,
     "keywords": ["stabilize", "dampen", "correct", "steady", "anchor"]},
    {"id": 16, "name": "Amplifier",     "layer": 8, "role": "actor",
     "spec": "Boosts signals, amplifies patterns, strengthens weak signals",
     "state": ILogicState.INVERSE,
     "keywords": ["amplify", "boost", "strengthen", "enhance", "magnify"]},

    # L9 MEMBRANE (basins 117-130) -- Boundary
    {"id": 17, "name": "Guardian",      "layer": 9, "role": "sensor",
     "spec": "Membrane control, access gating, boundary enforcement",
     "state": ILogicState.INVERSE,
     "keywords": ["guard", "protect", "gate", "filter", "boundary", "limit"]},
    {"id": 18, "name": "Transmitter",   "layer": 9, "role": "actor",
     "spec": "Cross-boundary messaging, external interface, bridging",
     "state": ILogicState.INVERSE,
     "keywords": ["transmit", "send", "bridge", "connect", "external", "relay"]},

    # L10 AGENT IDENTITY -- External layer
    {"id": 19, "name": "Identifier",    "layer": 10, "role": "sensor",
     "spec": "Self-awareness, identity verification, swarm introspection",
     "state": ILogicState.DIRECT,
     "keywords": ["identity", "self", "who", "introspect", "meta", "swarm"]},
]


# =========================================================================
# DATA STRUCTURES
# =========================================================================

@dataclass
class DispatchResult:
    """Result of dispatching a task to the swarm."""
    agent: AgentRole
    task_state: complex        # I-Logic state of the task
    outcome: complex           # Single-agent outcome (task x agent)
    outcome_name: str          # Human-readable single-agent outcome
    resolved: bool             # True if DIRECT (single or chained)
    confidence: float          # Keyword match score (0-1)
    reasoning: str             # Why this agent was chosen
    chain: Optional[WorkflowResult] = None  # Multi-agent resolution chain

    def to_dict(self) -> dict:
        d = {
            "agent": self.agent.to_dict(),
            "task_state": ILogicState.name(self.task_state),
            "outcome": self.outcome_name,
            "resolved": self.resolved,
            "confidence": round(self.confidence, 3),
            "reasoning": self.reasoning,
        }
        if self.chain:
            d["chain"] = self.chain.to_dict()
        return d


@dataclass
class WorkflowResult:
    """Result of composing a multi-agent workflow."""
    agents: List[AgentRole]
    states: List[complex]
    final_state: complex
    final_state_name: str
    steps: int
    resolved: bool
    path_description: str

    def to_dict(self) -> dict:
        return {
            "agents": [a.name for a in self.agents],
            "states": [ILogicState.name(s) for s in self.states],
            "final_state": self.final_state_name,
            "steps": self.steps,
            "resolved": self.resolved,
            "path": self.path_description,
        }


# =========================================================================
# SWARM
# =========================================================================

class Swarm:
    """19-agent swarm with I-Logic routing.

    Each agent has:
    - A layer assignment (L1-L10)
    - An I-Logic state (DIRECT, INVERSE, SPIRAL, ANTI_SPIRAL)
    - A home basin (center of its basin range)
    - Keywords for task matching

    Routing uses I-Logic algebraic composition:
        task_state x agent_state -> outcome
        DIRECT outcome = task resolved by this agent

    Usage:
        swarm = Swarm()

        # Dispatch a task
        result = swarm.dispatch("fix the authentication bug")
        print(result.agent.name)  # "Healer" (L6, INVERSE resolves errors)

        # Compose a workflow
        wf = swarm.compose_workflow(["Navigator", "Generator", "Validator"])
        print(wf.final_state_name)  # Composed I-Logic outcome
    """

    def __init__(self):
        self._agents: Dict[int, AgentRole] = {}
        self._name_index: Dict[str, int] = {}
        self._keyword_index: Dict[str, List[int]] = {}
        self._layer_index: Dict[int, List[int]] = {}
        self._initialize()

    def _initialize(self):
        """Build the agent registry from definitions."""
        for defn in AGENT_DEFINITIONS:
            layer = defn["layer"]
            # Compute home basin from layer
            if layer <= 9:
                basins = LAYER_BASINS.get(layer, range(0, 15))
                start, end = basins.start, basins.stop - 1
                home = (start + end) // 2
            else:
                # External layers don't have basin ranges
                start, end, home = 0, 130, 65

            agent = AgentRole(
                agent_id=defn["id"],
                name=defn["name"],
                layer=layer,
                layer_name=LAYER_NAMES.get(layer, f"L{layer}"),
                role=defn["role"],
                specialization=defn["spec"],
                i_logic_state=defn["state"],
                home_basin=home,
                basin_range=(start, end),
                keywords=defn["keywords"],
            )
            self._agents[agent.agent_id] = agent
            self._name_index[agent.name.lower()] = agent.agent_id

            # Index by keywords
            for kw in agent.keywords:
                if kw not in self._keyword_index:
                    self._keyword_index[kw] = []
                self._keyword_index[kw].append(agent.agent_id)

            # Index by layer
            if layer not in self._layer_index:
                self._layer_index[layer] = []
            self._layer_index[layer].append(agent.agent_id)

    # -----------------------------------------------------------------
    # CORE API
    # -----------------------------------------------------------------

    def dispatch(self, task_description: str) -> DispatchResult:
        """Route a task to the best agent via I-Logic prediction.

        1. Extract keywords from task description
        2. Score each agent by keyword match
        3. Among top candidates, pick the one whose I-Logic state
           best resolves the inferred task state
        4. Return the dispatch result with outcome prediction

        The task's I-Logic state is inferred from keywords:
        - Error/fix/bug -> INVERSE (needs healing)
        - Create/build/new -> SPIRAL (expansion)
        - Find/search/where -> DIRECT (pass-through)
        - Remove/delete/reduce -> ANTI_SPIRAL (contraction)
        """
        task_terms = set(task_description.lower().split())

        # Infer task I-Logic state from content
        task_state = self._infer_task_state(task_terms)

        # Score agents by keyword match
        scores: List[Tuple[float, AgentRole]] = []
        for agent in self._agents.values():
            score = sum(1 for kw in agent.keywords if kw in task_terms)
            # Normalize by keyword count
            if agent.keywords:
                score = score / len(agent.keywords)
            scores.append((score, agent))

        # Sort by score descending
        scores.sort(key=lambda x: x[0], reverse=True)

        # Among top candidates, find best I-Logic match
        best_agent = scores[0][1]
        best_confidence = scores[0][0]
        best_outcome = multiply(task_state, best_agent.i_logic_state)

        # Check if any top-5 agent resolves to DIRECT
        for score, agent in scores[:5]:
            outcome = multiply(task_state, agent.i_logic_state)
            if outcome == ILogicState.DIRECT:
                best_agent = agent
                best_confidence = max(score, 0.5)
                best_outcome = outcome
                break
            # Also prefer non-ABSORBED outcomes
            if outcome != ILogicState.ABSORBED and best_outcome == ILogicState.ABSORBED:
                best_agent = agent
                best_confidence = score
                best_outcome = outcome

        resolved = best_outcome == ILogicState.DIRECT

        # --- RESOLUTION CHAINING ---
        # When single-agent dispatch can't resolve, compose a multi-agent
        # chain that algebraically cancels the condition to DIRECT.
        # SPIRAL  -> SPIRAL agent (=INVERSE) -> INVERSE agent (=DIRECT)
        # ANTI_SPIRAL -> SPIRAL agent (=DIRECT)  [single-step chain]
        chain = None
        if not resolved and best_outcome != ILogicState.ABSORBED:
            chain = self.find_resolution_path(task_state)
            if chain and chain.resolved:
                resolved = True

        if chain and chain.resolved:
            chain_names = " -> ".join(a.name for a in chain.agents)
            reasoning = (
                f"Single: {ILogicState.name(task_state)} x "
                f"{ILogicState.name(best_agent.i_logic_state)} = "
                f"{ILogicState.name(best_outcome)}. "
                f"Chain: {chain_names} resolves to DIRECT."
            )
        else:
            reasoning = (
                f"Matched keywords in L{best_agent.layer} {best_agent.layer_name}. "
                f"Task state {ILogicState.name(task_state)} x "
                f"Agent state {ILogicState.name(best_agent.i_logic_state)} = "
                f"{ILogicState.name(best_outcome)}."
            )

        return DispatchResult(
            agent=best_agent,
            task_state=task_state,
            outcome=best_outcome,
            outcome_name=ILogicState.name(best_outcome),
            resolved=resolved,
            confidence=best_confidence,
            reasoning=reasoning,
            chain=chain,
        )

    def compose_workflow(
        self, agent_names: List[str]
    ) -> WorkflowResult:
        """Compose a multi-agent workflow via I-Logic path composition.

        Takes a list of agent names in execution order. Composes their
        I-Logic states using the algebraic compose() function to predict
        the overall workflow outcome.

        Example:
            wf = swarm.compose_workflow(["Navigator", "Generator", "Validator"])
            # DIRECT x DIRECT x INVERSE = INVERSE
            # -> workflow ends in INVERSE state (needs resolution)
        """
        agents = []
        states = []
        for name in agent_names:
            agent = self.get_agent_by_name(name)
            if agent is None:
                continue
            agents.append(agent)
            states.append(agent.i_logic_state)

        if not states:
            return WorkflowResult(
                agents=[], states=[], final_state=ILogicState.ABSORBED,
                final_state_name="ABSORBED", steps=0, resolved=False,
                path_description="No valid agents in workflow",
            )

        # Compose all states (O(N) in agents, O(1) per multiplication)
        final_state, steps_taken = compose(states)
        final_name = ILogicState.name(final_state)

        path_parts = []
        for a in agents:
            path_parts.append(f"{a.name}:{ILogicState.name(a.i_logic_state)}")
        path_desc = " -> ".join(path_parts) + f" = {final_name}"

        return WorkflowResult(
            agents=agents,
            states=states,
            final_state=final_state,
            final_state_name=final_name,
            steps=steps_taken,
            resolved=(final_state == ILogicState.DIRECT),
            path_description=path_desc,
        )

    def find_resolver(self, condition_state: complex) -> Optional[AgentRole]:
        """Find the agent whose state resolves a condition to DIRECT.

        condition x agent_state = DIRECT -> agent is the resolver.

        For INVERSE condition: needs INVERSE agent (INVERSE x INVERSE = DIRECT)
        For SPIRAL condition: needs ANTI_SPIRAL agent (but we don't have one)
                              or compose two SPIRAL agents (SPIRAL x SPIRAL = INVERSE,
                              then INVERSE x INVERSE = DIRECT)
        """
        for agent in self._agents.values():
            outcome = multiply(condition_state, agent.i_logic_state)
            if outcome == ILogicState.DIRECT:
                return agent
        return None

    def find_resolution_path(
        self, condition_state: complex, max_depth: int = 4
    ) -> Optional[WorkflowResult]:
        """Find a multi-agent path that resolves a condition.

        Tries single agents first, then pairs, then triples, up to max_depth.
        Uses I-Logic composition to predict outcomes algebraically.
        """
        # Try single agents
        for agent in self._agents.values():
            outcome = multiply(condition_state, agent.i_logic_state)
            if outcome == ILogicState.DIRECT:
                return WorkflowResult(
                    agents=[agent], states=[agent.i_logic_state],
                    final_state=ILogicState.DIRECT, final_state_name="DIRECT",
                    steps=1, resolved=True,
                    path_description=f"{agent.name} resolves {ILogicState.name(condition_state)} to DIRECT",
                )

        # Try pairs
        if max_depth >= 2:
            agents_list = list(self._agents.values())
            for a1 in agents_list:
                mid = multiply(condition_state, a1.i_logic_state)
                if mid == ILogicState.ABSORBED:
                    continue
                for a2 in agents_list:
                    final = multiply(mid, a2.i_logic_state)
                    if final == ILogicState.DIRECT:
                        return WorkflowResult(
                            agents=[a1, a2],
                            states=[a1.i_logic_state, a2.i_logic_state],
                            final_state=ILogicState.DIRECT,
                            final_state_name="DIRECT",
                            steps=2, resolved=True,
                            path_description=(
                                f"{a1.name} + {a2.name} resolves "
                                f"{ILogicState.name(condition_state)} to DIRECT"
                            ),
                        )

        return None

    # -----------------------------------------------------------------
    # QUERY API
    # -----------------------------------------------------------------

    def get_agent(self, agent_id: int) -> Optional[AgentRole]:
        """Get agent by ID. O(1)."""
        return self._agents.get(agent_id)

    def get_agent_by_name(self, name: str) -> Optional[AgentRole]:
        """Get agent by name (case-insensitive). O(1)."""
        agent_id = self._name_index.get(name.lower())
        if agent_id is not None:
            return self._agents.get(agent_id)
        return None

    def get_agents_for_layer(self, layer: int) -> List[AgentRole]:
        """Get all agents assigned to a layer."""
        ids = self._layer_index.get(layer, [])
        return [self._agents[i] for i in ids]

    def get_agents_for_basin(self, basin_id: int) -> List[AgentRole]:
        """Get agents whose basin range contains the given basin."""
        return [
            a for a in self._agents.values()
            if a.basin_range[0] <= basin_id <= a.basin_range[1]
        ]

    def list_agents(self) -> List[AgentRole]:
        """List all agents in order."""
        return [self._agents[i] for i in sorted(self._agents.keys())]

    @property
    def agent_count(self) -> int:
        return len(self._agents)

    # -----------------------------------------------------------------
    # INTERNAL
    # -----------------------------------------------------------------

    def _infer_task_state(self, terms: set) -> complex:
        """Infer I-Logic state of a task from its keywords.

        INVERSE:      error, fix, bug, broken, heal, repair, resolve
        SPIRAL:       create, build, new, expand, grow, add, generate
        ANTI_SPIRAL:  remove, delete, reduce, shrink, compress, trim
        ABSORBED:     stop, kill, terminate, abort, cancel, halt
        DIRECT:       everything else (neutral pass-through)
        """
        inverse_words = {"error", "fix", "bug", "broken", "heal", "repair",
                         "resolve", "fail", "crash", "corrupt", "wrong"}
        spiral_words = {"create", "build", "new", "expand", "grow", "add",
                        "generate", "start", "init", "setup", "implement"}
        anti_spiral_words = {"remove", "delete", "reduce", "shrink", "compress",
                             "trim", "clean", "prune", "minimize", "strip"}
        absorbed_words = {"stop", "kill", "terminate", "abort", "cancel",
                          "halt", "shutdown", "quit", "exit", "end"}

        scores = {
            ILogicState.INVERSE: len(terms & inverse_words),
            ILogicState.SPIRAL: len(terms & spiral_words),
            ILogicState.ANTI_SPIRAL: len(terms & anti_spiral_words),
            ILogicState.ABSORBED: len(terms & absorbed_words),
        }

        best_state = max(scores, key=scores.get)
        if scores[best_state] == 0:
            return ILogicState.DIRECT
        return best_state

    def summary(self) -> dict:
        """Swarm summary."""
        state_counts = {}
        for agent in self._agents.values():
            state_name = ILogicState.name(agent.i_logic_state)
            state_counts[state_name] = state_counts.get(state_name, 0) + 1

        return {
            "agent_count": len(self._agents),
            "layers_covered": sorted(set(a.layer for a in self._agents.values())),
            "state_distribution": state_counts,
            "agents": [a.to_dict() for a in self.list_agents()],
        }
