"""Tests for PeptideGym baseline agents."""
import gymnasium as gym
import numpy as np
import pytest

import peptidegym  # noqa: F401
from peptidegym.agents.random_agent import RandomAgent
from peptidegym.agents.heuristic_agent import (
    HeuristicAMPAgent,
    HeuristicCyclicAgent,
    HeuristicEpitopeAgent,
)
from peptidegym.peptide.properties import AMINO_ACIDS


def test_random_agent_creates():
    """RandomAgent should be constructible with any env."""
    env = gym.make("PeptideGym/AMP-v0")
    agent = RandomAgent(env, seed=42)
    assert agent is not None
    env.close()


def test_random_agent_predict_returns_valid_action():
    """predict() should return a valid action in the action space."""
    env = gym.make("PeptideGym/AMP-v0")
    agent = RandomAgent(env, seed=42)
    obs, _ = env.reset(seed=42)
    action, state = agent.predict(obs)
    assert env.action_space.contains(action)
    assert state is None
    env.close()


def test_random_agent_runs_episode():
    """RandomAgent should complete a full episode without crashing."""
    env = gym.make("PeptideGym/AMP-v0")
    agent = RandomAgent(env, seed=42)
    obs, _ = env.reset(seed=42)
    done = False
    steps = 0
    while not done and steps < 100:
        action, _ = agent.predict(obs)
        obs, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        steps += 1
    assert done or steps == 100
    env.close()


def test_heuristic_amp_agent_creates():
    """HeuristicAMPAgent should be constructible."""
    env = gym.make("PeptideGym/AMP-v0")
    agent = HeuristicAMPAgent(env, seed=42)
    assert agent is not None
    env.close()


def test_heuristic_amp_agent_runs_episode():
    """HeuristicAMPAgent should complete an episode and produce a sequence."""
    env = gym.make("PeptideGym/AMP-v0")
    agent = HeuristicAMPAgent(env, seed=42)
    obs, _ = env.reset(seed=42)
    done = False
    steps = 0
    while not done and steps < 100:
        action, _ = agent.predict(obs)
        obs, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        steps += 1
    assert done
    assert len(info["sequence"]) > 0
    env.close()


def test_heuristic_cyclic_agent_runs_episode():
    """HeuristicCyclicAgent should complete an episode."""
    env = gym.make("PeptideGym/CyclicPeptide-v0")
    agent = HeuristicCyclicAgent(env, seed=42)
    obs, _ = env.reset(seed=42)
    agent.reset()
    done = False
    steps = 0
    while not done and steps < 100:
        action, _ = agent.predict(obs)
        obs, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        steps += 1
    assert done
    env.close()


def test_heuristic_epitope_agent_runs_episode():
    """HeuristicEpitopeAgent should complete an episode."""
    env = gym.make("PeptideGym/Epitope-v0")
    agent = HeuristicEpitopeAgent(env, seed=42)
    obs, _ = env.reset(seed=42)
    agent.reset()
    done = False
    steps = 0
    while not done and steps < 100:
        action, _ = agent.predict(obs)
        obs, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        steps += 1
    assert done
    env.close()


def test_heuristic_agents_produce_valid_sequences():
    """All heuristic agents should produce sequences of valid amino acids."""
    configs = [
        ("PeptideGym/AMP-v0", HeuristicAMPAgent),
        ("PeptideGym/CyclicPeptide-v0", HeuristicCyclicAgent),
        ("PeptideGym/Epitope-v0", HeuristicEpitopeAgent),
    ]
    for env_id, AgentClass in configs:
        env = gym.make(env_id)
        agent = AgentClass(env, seed=42)
        if hasattr(agent, "reset"):
            agent.reset()
        obs, _ = env.reset(seed=42)
        done = False
        steps = 0
        while not done and steps < 100:
            action, _ = agent.predict(obs)
            obs, _, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            steps += 1
        seq = info.get("sequence", "")
        for aa in seq:
            assert aa in AMINO_ACIDS, f"Invalid AA '{aa}' in {env_id}"
        env.close()
