#!/usr/bin/env python3
"""
Moky SimLab - adapter.py
========================

A locally runnable *protocol sketch* of the SimLab interaction loop:

    Observation -> your Agent -> Action -> World -> WorldResult -> next Observation

IMPORTANT
---------
* This is NOT an official SDK, client library or production API.
  Field names, value ranges and the world itself are illustrative only.
* The "world" below (LocalSentinelWorld) is a tiny fake arena that runs in
  this process. Nothing is sent over the network.
* Only the Python standard library is used.

Usage
-----
    python3 adapter.py                     # 480 frames, seed 42
    python3 adapter.py --frames 960 --seed 7
    python3 adapter.py --out replay.json   # also write a per-frame replay

To plug in your own system, subclass `Agent` and implement `act()`.
Your model, memory, tools and deployment stay entirely on your side;
the world only ever sees the Action you return.
"""

from __future__ import annotations

import argparse
import json
import math
import random
import statistics
import time
from dataclasses import asdict, dataclass, field
from typing import Optional


# ---------------------------------------------------------------------------
# 1. Protocol types (illustrative)
# ---------------------------------------------------------------------------

@dataclass
class Observation:
    """What the world shows the Agent on one frame."""
    frame: int
    agent_pos: tuple[float, float]
    agent_stamina: float          # 0.0 - 1.0
    agent_health: float           # 0.0 - 1.0
    boss_pos: tuple[float, float]
    boss_intent: str              # "idle" | "windup" | "sweep_attack" | "recover"
    boss_health: float            # 0.0 - 1.0
    distance: float


@dataclass
class Action:
    """What the Agent sends back. The world only receives this."""
    type: str                     # "move" | "dodge" | "attack" | "wait"
    direction: Optional[str] = None   # "toward" | "left" | "right" | "back"


@dataclass
class WorldResult:
    """How the world responded to the Action on this frame."""
    frame: int
    damage_taken: float = 0.0
    damage_dealt: float = 0.0
    attack_avoided: Optional[bool] = None   # None when no attack landed this frame
    done: bool = False
    outcome: Optional[str] = None           # "sentinel_down" | "agent_down" | "timeout"


# ---------------------------------------------------------------------------
# 2. Agent interface - this is the part you replace
# ---------------------------------------------------------------------------

class Agent:
    """Minimal Agent contract: reset once per episode, act once per frame."""

    def reset(self, seed: int) -> None:
        pass

    def act(self, obs: Observation) -> Action:
        raise NotImplementedError


class SimpleSentinelPolicy(Agent):
    """A tiny hand-written policy with a little memory.

    It remembers how many frames the boss usually spends winding up before a
    sweep, so it learns to dodge a bit earlier as the episode goes on.
    """

    def reset(self, seed: int) -> None:
        self.windup_frames: list[int] = []
        self._current_windup = 0
        self._dodge_side = "left"

    def act(self, obs: Observation) -> Action:
        # memory: measure how long wind-ups last
        if obs.boss_intent == "windup":
            self._current_windup += 1
        elif self._current_windup:
            self.windup_frames.append(self._current_windup)
            self._current_windup = 0

        expected = statistics.median(self.windup_frames) if self.windup_frames else 6
        about_to_hit = obs.boss_intent == "sweep_attack" or (
            obs.boss_intent == "windup" and self._current_windup >= expected - 1
        )

        if about_to_hit and obs.agent_stamina > 0.15 and obs.distance < 3.0:
            self._dodge_side = "right" if self._dodge_side == "left" else "left"
            return Action("dodge", self._dodge_side)
        if obs.distance > 1.4:
            return Action("move", "toward")
        if obs.boss_intent in ("recover", "idle") and obs.agent_stamina > 0.1:
            return Action("attack")
        return Action("wait")


# ---------------------------------------------------------------------------
# 3. A fake local world (stands in for the real SimLab arena)
# ---------------------------------------------------------------------------

@dataclass
class _State:
    frame: int = 0
    agent_pos: list = field(default_factory=lambda: [9.0, 6.0])
    stamina: float = 1.0
    agent_hp: float = 1.0
    boss_hp: float = 1.0
    intent: str = "idle"
    intent_left: int = 20
    dodging: int = 0


class LocalSentinelWorld:
    """THE SENTINEL, very much simplified. Deterministic for a given seed."""

    BOSS_POS = (0.0, 0.0)

    def __init__(self, seed: int = 42):
        self.seed = seed

    def reset(self) -> Observation:
        self.rng = random.Random(self.seed)
        self.s = _State()
        return self._observe()

    def _distance(self) -> float:
        x, y = self.s.agent_pos
        return math.hypot(x - self.BOSS_POS[0], y - self.BOSS_POS[1])

    def _observe(self) -> Observation:
        s = self.s
        return Observation(
            frame=s.frame,
            agent_pos=(round(s.agent_pos[0], 2), round(s.agent_pos[1], 2)),
            agent_stamina=round(s.stamina, 2),
            agent_health=round(s.agent_hp, 2),
            boss_pos=self.BOSS_POS,
            boss_intent=s.intent,
            boss_health=round(s.boss_hp, 3),
            distance=round(self._distance(), 2),
        )

    def _advance_boss(self) -> None:
        s = self.s
        s.intent_left -= 1
        if s.intent_left > 0:
            return
        nxt = {"idle": "windup", "windup": "sweep_attack",
               "sweep_attack": "recover", "recover": "idle"}[s.intent]
        s.intent = nxt
        s.intent_left = {"idle": self.rng.randint(10, 24),
                         "windup": self.rng.randint(5, 8),
                         "sweep_attack": 3,
                         "recover": self.rng.randint(12, 18)}[nxt]

    def step(self, action: Action) -> tuple[Observation, WorldResult]:
        s = self.s
        res = WorldResult(frame=s.frame)

        # apply the Agent's Action
        if action.type == "move":
            d = self._distance() or 1.0
            s.agent_pos[0] -= s.agent_pos[0] / d * 0.18
            s.agent_pos[1] -= s.agent_pos[1] / d * 0.18
        elif action.type == "dodge" and s.stamina >= 0.15:
            s.stamina -= 0.15
            s.dodging = 4
        elif action.type == "attack" and self._distance() <= 1.6 and s.stamina >= 0.05:
            s.stamina -= 0.05
            hit = 0.012 if s.intent == "recover" else 0.005
            s.boss_hp = max(0.0, s.boss_hp - hit)
            res.damage_dealt = hit

        # the world advances
        s.stamina = min(1.0, s.stamina + 0.02)
        if s.dodging:
            s.dodging -= 1
        self._advance_boss()
        if s.intent == "sweep_attack" and s.intent_left == 3 and self._distance() < 3.0:
            if s.dodging:
                res.attack_avoided = True
            else:
                res.attack_avoided = False
                res.damage_taken = 0.18
                s.agent_hp = max(0.0, s.agent_hp - 0.18)

        s.frame += 1
        if s.boss_hp <= 0:
            res.done, res.outcome = True, "sentinel_down"
        elif s.agent_hp <= 0:
            res.done, res.outcome = True, "agent_down"
        return self._observe(), res


# ---------------------------------------------------------------------------
# 4. The interaction loop (what SimLab runs around your Agent)
# ---------------------------------------------------------------------------

def run_episode(world: LocalSentinelWorld, agent: Agent, max_frames: int = 480) -> list[dict]:
    obs = world.reset()
    agent.reset(world.seed)
    records = []
    for _ in range(max_frames):
        t0 = time.perf_counter()
        action = agent.act(obs)                    # your logic, memory, tools
        latency_ms = (time.perf_counter() - t0) * 1000
        next_obs, result = world.step(action)      # the world only sees Actions
        records.append({
            "observation": asdict(obs),
            "action": asdict(action),
            "world_result": asdict(result),
            "latency_ms": round(latency_ms, 4),
        })
        obs = next_obs
        if result.done:
            break
    if not records[-1]["world_result"]["done"]:
        records[-1]["world_result"].update(done=True, outcome="timeout")
    return records


def summarize(records: list[dict]) -> dict:
    lat = sorted(r["latency_ms"] for r in records)
    results = [r["world_result"] for r in records]
    attacks = [x for x in results if x["attack_avoided"] is not None]
    avoided = sum(1 for x in attacks if x["attack_avoided"])
    actions: dict[str, int] = {}
    for r in records:
        actions[r["action"]["type"]] = actions.get(r["action"]["type"], 0) + 1
    dealt = sum(x["damage_dealt"] for x in results)
    p95 = lat[min(len(lat) - 1, int(len(lat) * 0.95))]
    return {
        "frames": len(records),
        "outcome": results[-1]["outcome"],
        "boss_health_left": round(max(0.0, 1.0 - dealt), 3),
        "damage_dealt": round(dealt, 3),
        "damage_taken": round(sum(x["damage_taken"] for x in results), 3),
        "attacks_avoided": f"{avoided}/{len(attacks)}",
        "actions": actions,
        "latency_ms": {
            "mean": round(statistics.fmean(lat), 4),
            "p50": round(statistics.median(lat), 4),
            "p95": round(p95, 4),
            "max": round(lat[-1], 4),
        },
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="Moky SimLab protocol sketch (not an official SDK)")
    ap.add_argument("--frames", type=int, default=480)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--out", help="optional path to write the per-frame replay JSON")
    args = ap.parse_args()

    world = LocalSentinelWorld(seed=args.seed)
    records = run_episode(world, SimpleSentinelPolicy(), max_frames=args.frames)
    summary = summarize(records)

    print("Moky SimLab - local protocol sketch (synthetic, not a benchmark result)")
    print(f"world: LocalSentinelWorld  seed: {args.seed}  frames: {summary['frames']}")
    for key in ("outcome", "boss_health_left", "damage_dealt", "damage_taken", "attacks_avoided"):
        print(f"  {key:<17} {summary[key]}")
    print(f"  {'actions':<17} " + ", ".join(f"{k}={v}" for k, v in sorted(summary["actions"].items())))
    lm = summary["latency_ms"]
    print(f"  {'latency (ms)':<17} mean {lm['mean']}  p50 {lm['p50']}  p95 {lm['p95']}  max {lm['max']}")

    if args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            json.dump({"meta": {"world": "LocalSentinelWorld", "seed": args.seed,
                                "synthetic": True}, "summary": summary,
                       "frames": records}, f, ensure_ascii=False, indent=2)
        print(f"replay written to {args.out}")


if __name__ == "__main__":
    main()