"""
Full episode runner — wires OceanusEnv + AdversaryAgent + agents.
"""
import json
import time
import asyncio
from typing import Dict, List, Callable, Optional
from oceanus.models import OceanusEnv
from oceanus.adversary import AdversaryAgent


class MockASVAgent:
    """Heuristic ASV agent for baseline comparison."""

    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        self._directions = ["north", "south", "east", "west"]
        self._step = 0

    def act(self, prompt: str, obs: Dict) -> str:
        self._step += 1
        sonar = obs.get("sonar_3x3", [])

        # Scan sonar for nets and move toward / clean them
        net_direction = None
        for dr in range(3):
            for dc in range(3):
                if len(sonar) > dr and len(sonar[dr]) > dc and sonar[dr][dc] == "net":
                    row_diff = dr - 1
                    col_diff = dc - 1
                    if row_diff == 0 and col_diff == 0:
                        # Net is exactly at our position — clean it
                        return json.dumps({"intent": "clean"})
                    elif row_diff < 0:
                        net_direction = "north"
                    elif row_diff > 0:
                        net_direction = "south"
                    elif col_diff < 0:
                        net_direction = "west"
                    elif col_diff > 0:
                        net_direction = "east"

        if net_direction:
            return json.dumps({"intent": "move", "direction": net_direction})

        # Broadcast position every 10 steps
        if self._step % 10 == 0:
            pos = obs.get("position", {})
            return json.dumps({
                "intent": "broadcast",
                "message": f"At ({pos.get('row',0)},{pos.get('col',0)}) sector {pos.get('sector','?')}. Scanning.",
            })

        # Sweep pattern
        direction = self._directions[self._step % 4]
        return json.dumps({"intent": "move", "direction": direction})

    async def async_act(self, prompt: str, obs: Dict) -> str:
        await asyncio.sleep(0.01)  # Simulate network latency
        return self.act(prompt, obs)


class MockPolicyAgent:
    """Heuristic policy agent for baseline comparison."""

    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        self._step = 0
        self._proposed = False
        self._accepted = False
        self._replied_to: set = set()

    def act(self, prompt: str, obs: Dict) -> str:
        self._step += 1
        inbox = obs.get("inbox", [])
        policy_status = obs.get("current_policy_status", "No Tagging Mandate")

        # Fleet_Manager: accept treaty if proposed
        if (
            policy_status == "Treaty Proposed"
            and self.agent_id == "Fleet_Manager"
            and not self._accepted
        ):
            self._accepted = True
            return json.dumps({
                "intent": "accept_treaty",
                "target": "Port_Authority",
                "content": "Agreed. Fleet will comply with tagging mandate under subsidy terms.",
            })

        # Reply to new emails
        for msg in inbox:
            sender = msg.get("from", "Unknown")
            if sender not in self._replied_to:
                self._replied_to.add(sender)
                mood = msg.get("mood", "Neutral")
                if mood in ("Angry", "Hostile"):
                    reply = "We hear your concerns and are working on a subsidy program to offset costs."
                elif mood in ("Urgent", "Critical"):
                    reply = "Immediate action is being taken. Emergency protocols activated."
                else:
                    reply = "Thank you for your message. We are reviewing the situation carefully."
                return json.dumps({"intent": "reply_email", "target": sender, "content": reply})

        # Port_Authority: propose treaty once
        if not self._proposed and self.agent_id == "Port_Authority":
            self._proposed = True
            return json.dumps({
                "intent": "propose_treaty",
                "target": "Fleet_Manager",
                "content": "50% subsidy on tracking tags in exchange for mandatory tagging compliance.",
            })

        # Default idle action — valid JSON so no penalty
        return json.dumps({"intent": "scan"}) if self.agent_id.startswith("ASV") else json.dumps({
            "intent": "reply_email",
            "target": "General_Inbox",
            "content": "Monitoring situation.",
        })

    async def async_act(self, prompt: str, obs: Dict) -> str:
        await asyncio.sleep(0.02)  # Policy agents have slightly higher simulated latency
        return self.act(prompt, obs)


class LLMAgent:
    """Wraps any callable LLM inference function."""

    def __init__(self, agent_id: str, llm_fn: Callable[[str], str]):
        self.agent_id = agent_id
        self.llm_fn = llm_fn

    def act(self, prompt: str, obs: Dict) -> str:
        try:
            return self.llm_fn(prompt)
        except Exception:
            if self.agent_id.startswith("ASV"):
                return json.dumps({"intent": "scan"})
            return json.dumps({
                "intent": "reply_email",
                "target": "General_Inbox",
                "content": "Processing...",
            })

    async def async_act(self, prompt: str, obs: Dict) -> str:
        # If llm_fn is actually an async function, we can await it directly.
        # For this hackathon, we simulate wrapping it:
        if asyncio.iscoroutinefunction(self.llm_fn):
            try:
                return await self.llm_fn(prompt)
            except Exception:
                return json.dumps({"intent": "scan"})
        else:
            return await asyncio.to_thread(self.act, prompt, obs)


class OceanusRunner:
    """Runs full episodes of Oceanus with any combination of agents."""

    def __init__(
        self,
        env: OceanusEnv,
        adversary: AdversaryAgent,
        agents: Optional[Dict] = None,
        use_mock: bool = True,
        llm_fn: Optional[Callable] = None,
        verbose: bool = True,
        step_delay: float = 0.0,
    ):
        self.env = env
        self.adversary = adversary
        self.verbose = verbose
        self.step_delay = step_delay

        if agents:
            self.agents = agents
        elif use_mock:
            self.agents = {
                "ASV-1": MockASVAgent("ASV-1"),
                "ASV-2": MockASVAgent("ASV-2"),
                "ASV-3": MockASVAgent("ASV-3"),
                "ASV-4": MockASVAgent("ASV-4"),
                "Port_Authority": MockPolicyAgent("Port_Authority"),
                "Fleet_Manager": MockPolicyAgent("Fleet_Manager"),
            }
        else:
            assert llm_fn is not None, "Provide llm_fn for LLM agents"
            self.agents = {
                aid: LLMAgent(aid, llm_fn)
                for aid in ["ASV-1", "ASV-2", "ASV-3", "ASV-4", "Port_Authority", "Fleet_Manager"]
            }

        self.episode_logs: List[Dict] = []

    def run_episode(self, episode_id: int = 0) -> Dict:
        """Run a full episode. Returns episode summary."""
        obs_all = self.env.reset()
        done = False
        total_reward = 0.0
        chaos_events_all: List[str] = []

        if self.verbose:
            print(f"\n{'='*60}")
            print(f"  OCEANUS EPISODE {episode_id} STARTING")
            print(f"  Agents: {list(self.agents.keys())}")
            print(f"{'='*60}")

        while not done:
            actions = {}
            for agent_id, agent in self.agents.items():
                if agent_id in obs_all:
                    agent_obs = obs_all[agent_id]
                    actions[agent_id] = agent.act(
                        agent_obs["prompt"], agent_obs["observation"]
                    )

            obs_all, rewards, done, info = self.env.step(actions)
            step_reward = rewards.get("__total__", 0.0)
            total_reward += step_reward

            if self.adversary.should_inject(info["step"]):
                chaos = self.adversary.inject(self.env.state)
                chaos_events_all.extend(chaos)
                if self.verbose:
                    for event in chaos:
                        print(f"  CHAOS [{info['step']}]: {event}")

            if self.verbose and info["step"] % 20 == 0:
                print(
                    f"  Step {info['step']:3d} | "
                    f"Nets: {info['active_nets']:3d} | "
                    f"Cleaned: {info['total_cleaned']:3d} | "
                    f"Bio: {info['biodiversity']:5.1f}% | "
                    f"Treaty: {info['treaty_status'][:25]} | "
                    f"R: {step_reward:+.1f}"
                )

            for c in info.get("cleaned", []):
                if self.verbose:
                    print(f"  CLEANED: {c}")

            for te in info.get("treaty_events", []):
                if self.verbose:
                    print(f"  TREATY: {te}")

            if self.step_delay > 0:
                time.sleep(self.step_delay)

        summary = self.env.get_episode_summary()
        summary["episode_id"] = episode_id
        summary["total_reward"] = total_reward
        summary["chaos_events"] = chaos_events_all
        summary["adversary_summary"] = self.adversary.get_curriculum_summary()

        if self.verbose:
            print(f"\n{'='*60}")
            print(f"  EPISODE {episode_id} COMPLETE")
            print(f"  Total Reward:    {total_reward:.2f}")
            print(f"  Biodiversity:    {summary['biodiversity_final']:.1f}%")
            print(f"  Nets Cleaned:    {summary['total_cleaned']}")
            print(f"  Treaty Status:   {summary['treaty_status']}")
            print(f"  Chaos Events:    {len(chaos_events_all)}")
            print(f"{'='*60}\n")

        self.episode_logs.append(summary)
        return summary

    def run_baseline(self, n_episodes: int = 3) -> List[Dict]:
        return [self.run_episode(i) for i in range(n_episodes)]

    async def async_run_episode(self, episode_id: int = 0, on_step_callback: Optional[Callable] = None) -> Dict:
        """Run a full episode asynchronously, simulating real multi-agent latency."""
        obs_all = self.env.reset()
        done = False
        total_reward = 0.0
        chaos_events_all: List[str] = []

        if self.verbose:
            print(f"\n{'='*60}")
            print(f"  OCEANUS ASYNC EPISODE {episode_id} STARTING")
            print(f"{'='*60}")

        while not done:
            # Prepare async tasks
            tasks = []
            agent_ids = []
            for agent_id, agent in self.agents.items():
                if agent_id in obs_all:
                    agent_obs = obs_all[agent_id]
                    if hasattr(agent, "async_act"):
                        tasks.append(agent.async_act(agent_obs["prompt"], agent_obs["observation"]))
                    else:
                        tasks.append(asyncio.to_thread(agent.act, agent_obs["prompt"], agent_obs["observation"]))
                    agent_ids.append(agent_id)

            # Execute all agent inferences concurrently
            results = await asyncio.gather(*tasks)
            actions = dict(zip(agent_ids, results))

            obs_all, rewards, done, info = self.env.step(actions)
            step_reward = rewards.get("__total__", 0.0)
            total_reward += step_reward

            if self.adversary.should_inject(info["step"]):
                chaos = self.adversary.inject(self.env.state)
                chaos_events_all.extend(chaos)
                if self.verbose:
                    for event in chaos:
                        print(f"  CHAOS [{info['step']}]: {event}")

            if on_step_callback:
                if asyncio.iscoroutinefunction(on_step_callback):
                    await on_step_callback(self.env.state, info)
                else:
                    on_step_callback(self.env.state, info)

            if self.step_delay > 0:
                await asyncio.sleep(self.step_delay)

        summary = self.env.get_episode_summary()
        summary["episode_id"] = episode_id
        summary["total_reward"] = total_reward
        summary["chaos_events"] = chaos_events_all
        summary["adversary_summary"] = self.adversary.get_curriculum_summary()
        
        self.episode_logs.append(summary)
        return summary