from __future__ import annotations

import argparse
import json
from datetime import datetime, timezone
from pathlib import Path

import pygame
from vizdoom import Button, GameVariable

from run_experiment import MenuAction, OUTPUTS, action_vector, make_game


SCREEN_SIZE = (960, 720)
GAME_RECT = pygame.Rect(0, 0, 960, 620)
PANEL_RECT = pygame.Rect(0, 620, 960, 100)


def menu_for(scenario: str) -> list[MenuAction]:
    if scenario == "basic":
        return [
            MenuAction("ATTACK", Button.ATTACK, 4),
            MenuAction("STRAFE LEFT", Button.MOVE_LEFT, 3),
            MenuAction("STRAFE RIGHT", Button.MOVE_RIGHT, 3),
        ]
    return [
        MenuAction("FORWARD", Button.MOVE_FORWARD, 8),
        MenuAction("TURN LEFT", Button.TURN_LEFT, 2),
        MenuAction("TURN RIGHT", Button.TURN_RIGHT, 2),
    ]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Play a ViZDoom scenario through one switch")
    parser.add_argument("--scenario", choices=("basic", "my_way_home"), default="my_way_home")
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--dwell-ms", type=int, default=500)
    parser.add_argument(
        "--headless-smoke",
        action="store_true",
        help="Advance ten neutral tics without opening a window",
    )
    return parser.parse_args()


def run_headless_smoke(args: argparse.Namespace) -> None:
    game = make_game(args.scenario, args.seed)
    for _ in range(10):
        game.make_action(action_vector(game, None), 1)
    assert game.get_state() is not None
    print(
        json.dumps(
            {
                "scenario": args.scenario,
                "seed": args.seed,
                "neutral_tics": 10,
                "status": "ok",
            }
        )
    )
    game.close()


def is_switch_event(event: pygame.event.Event) -> bool:
    return event.type == pygame.MOUSEBUTTONDOWN or (
        event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE
    )


def write_session(args: argparse.Namespace, events: list[dict], success: bool) -> Path:
    OUTPUTS.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    path = OUTPUTS / f"manual-session-{args.scenario}-{stamp}.json"
    path.write_text(
        json.dumps(
            {
                "created_at": datetime.now(timezone.utc).isoformat(),
                "scenario": args.scenario,
                "seed": args.seed,
                "dwell_ms": args.dwell_ms,
                "input_contract": "SPACE or one mouse button selects the highlighted action",
                "success": success,
                "events": events,
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    return path


def run_interactive(args: argparse.Namespace) -> None:
    pygame.init()
    pygame.display.set_caption("ONE SWITCH DOOM")
    window = pygame.display.set_mode(SCREEN_SIZE)
    font = pygame.font.SysFont("Consolas", 25, bold=True)
    small_font = pygame.font.SysFont("Consolas", 18)
    clock = pygame.time.Clock()

    game = make_game(args.scenario, args.seed)
    menu = menu_for(args.scenario)
    highlighted = 0
    scan_elapsed = 0
    switch_events: list[dict] = []
    running = True
    final_message = ""

    while running:
        delta_ms = clock.tick(35)
        scan_elapsed += delta_ms
        while scan_elapsed >= args.dwell_ms:
            highlighted = (highlighted + 1) % len(menu)
            scan_elapsed -= args.dwell_ms

        for event in pygame.event.get():
            if event.type == pygame.QUIT or (
                event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE
            ):
                running = False
                continue
            if is_switch_event(event) and not game.is_episode_finished():
                selected = menu[highlighted]
                game.make_action(action_vector(game, selected.button), selected.duration_tics)
                switch_events.append(
                    {
                        "press": len(switch_events) + 1,
                        "action": selected.name,
                        "episode_tic": game.get_episode_time(),
                    }
                )
                highlighted = 0
                scan_elapsed = 0

        if not game.is_episode_finished():
            game.make_action(action_vector(game, None), 1)
            state = game.get_state()
            if state is not None:
                surface = pygame.surfarray.make_surface(state.screen_buffer.swapaxes(0, 1))
                surface = pygame.transform.scale(surface, GAME_RECT.size)
                window.blit(surface, GAME_RECT)
        else:
            success = (
                game.get_game_variable(GameVariable.KILLCOUNT) > 0
                if args.scenario == "basic"
                else game.get_game_variable(GameVariable.ARMOR) > 0
            )
            final_message = "SCENARIO CLEARED" if success else "THE DEMON / CLOCK WON"
            running = False

        pygame.draw.rect(window, (10, 10, 10), PANEL_RECT)
        option_width = PANEL_RECT.width // len(menu)
        for index, item in enumerate(menu):
            rect = pygame.Rect(index * option_width, PANEL_RECT.top, option_width, 62)
            color = (255, 238, 0) if index == highlighted else (72, 72, 72)
            pygame.draw.rect(window, color, rect, 0)
            pygame.draw.rect(window, (255, 255, 255), rect, 2)
            label = font.render(item.name, True, (0, 0, 0) if index == highlighted else (255, 255, 255))
            window.blit(label, label.get_rect(center=rect.center))
        instructions = small_font.render(
            "ONE INPUT ONLY: press SPACE or click to select the yellow action. ESC quits.",
            True,
            (255, 255, 255),
        )
        window.blit(instructions, (16, 690))
        pygame.display.flip()

    success = (
        game.get_game_variable(GameVariable.KILLCOUNT) > 0
        if args.scenario == "basic"
        else game.get_game_variable(GameVariable.ARMOR) > 0
    )
    session_path = write_session(args, switch_events, success)
    game.close()
    pygame.quit()
    print(final_message or "Session closed")
    print(f"Switch log: {session_path}")


def main() -> None:
    args = parse_args()
    if args.headless_smoke:
        run_headless_smoke(args)
    else:
        run_interactive(args)


if __name__ == "__main__":
    main()