"""Summarize y86 x166 bounce opportunities and score-chain markers."""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter
from dataclasses import asdict, dataclass
from pathlib import Path

import envpool
import numpy as np


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from pong_policy import ACTION_NAMES, RamPolicy
from pong_state import PongState, detect_state


@dataclass(frozen=True)
class BounceEvent:
    step: int
    ball_x: int
    ball_y: int
    self_paddle_y: int
    self_score: int
    opponent_score: int
    dx: int | None
    dy: int | None
    target_y: float
    action_name: str
    mode: str
    previous_paddle_delta: int | None
    score22_marker_before: bool
    score22_marker_after: bool


@dataclass(frozen=True)
class SeedSummary:
    seed: int
    score: float
    steps: int
    done: bool
    bounce_count: int
    down_bounce_count: int
    score22_marker_seen: bool
    score_pairs: list[str]
    first_bounce_step: int | None
    last_bounce_step: int | None
    events: list[dict]


def reset_env(env: object) -> tuple[np.ndarray, dict]:
    reset_out = env.reset()
    if isinstance(reset_out, tuple):
        obs, info = reset_out
    else:
        obs, info = reset_out, {}
    return np.asarray(obs), info


def step_env(env: object, action: int) -> tuple[np.ndarray, float, bool, dict]:
    step_out = env.step(np.array([action], dtype=np.int32))
    if len(step_out) == 5:
        obs, reward, terminated, truncated, info = step_out
        done = np.logical_or(terminated, truncated)
    else:
        obs, reward, done, info = step_out
    return (
        np.asarray(obs),
        float(np.asarray(reward).reshape(-1)[0]),
        bool(np.asarray(done).reshape(-1)[0]),
        info,
    )


def is_y86_x166_bounce(state: PongState, previous_state: PongState | None, target_y: float) -> bool:
    if previous_state is None:
        return False
    dx = state.ball_x - previous_state.ball_x
    dy = state.ball_y - previous_state.ball_y
    return (
        dx > 0
        and dy == 5
        and 164 <= state.ball_x <= 168
        and 48 <= state.ball_y <= 52
        and 76 <= state.self_paddle_y <= 80
        and 80 <= target_y <= 85
    )


def run_seed(seed: int, max_steps: int) -> SeedSummary:
    env = envpool.make("Pong-v5", env_type="gym", num_envs=1, seed=seed)
    try:
        _obs, info = reset_env(env)
        policy = RamPolicy()
        policy.reset()
        previous_paddle_y: int | None = None
        score = 0.0
        events: list[BounceEvent] = []
        score22_marker_seen = False

        for step in range(1, max_steps + 1):
            state = detect_state(info)
            previous_state = policy.previous_state
            marker_before = policy.y86_score22_bounce_family_seen
            action = policy.select_action(info)
            marker_after = policy.y86_score22_bounce_family_seen
            target_y = policy.last_target_y
            mode = policy.last_mode
            dx = dy = None
            if previous_state is not None:
                dx = state.ball_x - previous_state.ball_x
                dy = state.ball_y - previous_state.ball_y
            previous_paddle_delta = None
            if previous_paddle_y is not None:
                previous_paddle_delta = state.self_paddle_y - previous_paddle_y

            if is_y86_x166_bounce(state, previous_state, target_y):
                score22_marker_seen = score22_marker_seen or marker_after
                events.append(
                    BounceEvent(
                        step=step,
                        ball_x=state.ball_x,
                        ball_y=state.ball_y,
                        self_paddle_y=state.self_paddle_y,
                        self_score=state.self_score,
                        opponent_score=state.opponent_score,
                        dx=dx,
                        dy=dy,
                        target_y=target_y,
                        action_name=ACTION_NAMES.get(action, str(action)),
                        mode=mode,
                        previous_paddle_delta=previous_paddle_delta,
                        score22_marker_before=marker_before,
                        score22_marker_after=marker_after,
                    )
                )

            previous_paddle_y = state.self_paddle_y
            _obs, reward, done, info = step_env(env, action)
            score += reward
            if done:
                return build_summary(seed, score, step, done, events)

        return build_summary(seed, score, max_steps, False, events)
    finally:
        env.close()


def build_summary(
    seed: int,
    score: float,
    steps: int,
    done: bool,
    events: list[BounceEvent],
) -> SeedSummary:
    return SeedSummary(
        seed=seed,
        score=score,
        steps=steps,
        done=done,
        bounce_count=len(events),
        down_bounce_count=sum(1 for event in events if event.action_name == "down"),
        score22_marker_seen=any(event.score22_marker_after for event in events),
        score_pairs=[f"{event.self_score}-{event.opponent_score}" for event in events],
        first_bounce_step=events[0].step if events else None,
        last_bounce_step=events[-1].step if events else None,
        events=[asdict(event) for event in events],
    )


def parse_seeds(raw: str) -> list[int]:
    if ".." in raw:
        start, end = raw.split("..", 1)
        return list(range(int(start), int(end) + 1))
    return [int(part.strip()) for part in raw.split(",") if part.strip()]


def summarize(rows: list[SeedSummary]) -> dict:
    by_score = Counter(str(int(row.score)) for row in rows)
    by_score_pairs = Counter("|".join(row.score_pairs) if row.score_pairs else "<none>" for row in rows)
    score13 = [row.seed for row in rows if row.score == 13.0]
    score13_with_events = [row.seed for row in rows if row.score == 13.0 and row.bounce_count > 0]
    marker_seeds = [row.seed for row in rows if row.score22_marker_seen]
    return {
        "seed_count": len(rows),
        "score_counts": dict(sorted(by_score.items(), key=lambda item: int(item[0]))),
        "score13_seeds": score13,
        "score13_with_y86_x166_events": score13_with_events,
        "score22_marker_seeds": marker_seeds,
        "score_pair_chain_counts": [
            {"count": count, "score_pairs": chain}
            for chain, count in by_score_pairs.most_common()
        ],
        "rows": [asdict(row) for row in rows],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seeds", default="0..79")
    parser.add_argument("--max-steps", type=int, default=30_000)
    parser.add_argument("--output", type=Path, default=Path("diagnostics/y86_bounce_chains.json"))
    args = parser.parse_args()

    rows = [run_seed(seed, args.max_steps) for seed in parse_seeds(args.seeds)]
    result = summarize(rows)
    result["source"] = {
        "script": str(Path(__file__)),
        "seeds": args.seeds,
        "max_steps": args.max_steps,
        "envpool": envpool.__version__,
    }
    text = json.dumps(result, indent=2, sort_keys=True)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(text + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()
