"""Tests for the observation_function machinery: GameMetadata resolution,
GameEnvironment delegation, and DQNTrainer's extras_size discovery.

Run with: uv run python -m unittest tests/test_observation_function.py -v
"""
import tempfile
from pathlib import Path
from unittest import TestCase, main

import numpy as np
from retro.game import Game

from retro_gamer.metadata import GameMetadata
from retro_gamer.observation import encode_observation
from retro_gamer.env import GameEnvironment
from retro_gamer.trainer import DQNTrainer

ACTIONS = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]


def sample_observation_fn(game):
    """A trivial observation_function used to test dotted-path resolution."""
    return "observation-from-sample_observation_fn"


def tuple_observation_fn(game):
    """Returns a plain tuple — exercises the non-array tabular use case."""
    return (1, 2, 3)


def make_game_factory(board_size=(2, 2)):
    def factory():
        return Game([], {'reward': 0.0, 'score': 0}, board_size=board_size, show_state=False)
    return factory


def obs_with_extras(game):
    """Board (character_set=2, board_size=(2,2) -> 8) + 3 extras = 11."""
    return np.zeros(2 * 2 * 2 + 3, dtype=np.float32)


def obs_2d(game):
    """Not 1-D — used to test the training-time shape validation."""
    return np.zeros((2, 2), dtype=np.float32)


def obs_too_short(game):
    """Shorter than character_set (2) x board_size (2x2) = 8."""
    return np.zeros(5, dtype=np.float32)


_changing_obs_counter = {"n": 0}


def changing_obs(game):
    """Returns length 8 on the first call, length 9 on every call after."""
    _changing_obs_counter["n"] += 1
    length = 8 if _changing_obs_counter["n"] == 1 else 9
    return np.zeros(length, dtype=np.float32)


class TestResolveObservationFunction(TestCase):
    def test_returns_none_when_unset(self):
        metadata = GameMetadata(actions=ACTIONS, reward="reward")
        self.assertIsNone(metadata.resolve_observation_function())

    def test_resolves_valid_dotted_path(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function=f"{__name__}:sample_observation_fn",
        )
        fn = metadata.resolve_observation_function()
        self.assertIs(fn, sample_observation_fn)

    def test_raises_on_malformed_string(self):
        metadata = GameMetadata(actions=ACTIONS, reward="reward", observation_function="no_colon_here")
        with self.assertRaises(ValueError):
            metadata.resolve_observation_function()

    def test_raises_on_unimportable_module(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function="this_module_does_not_exist:fn",
        )
        with self.assertRaises(ValueError):
            metadata.resolve_observation_function()

    def test_raises_on_missing_attribute(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function=f"{__name__}:no_such_function",
        )
        with self.assertRaises(ValueError):
            metadata.resolve_observation_function()

    def test_validate_raises_on_malformed_observation_function(self):
        metadata = GameMetadata(actions=ACTIONS, reward="reward", observation_function="no_colon_here")
        with self.assertRaises(ValueError):
            metadata.validate()

    def test_validate_passes_with_well_formed_observation_function(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function=f"{__name__}:sample_observation_fn",
        )
        metadata.validate()  # should not raise


class TestEncodeObservation(TestCase):
    def setUp(self):
        self.metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            character_set=["@", "*"], board_size=(2, 2),
        )
        self.board_chars = [["@", " "], [" ", "*"]]

    def test_board_plus_extras(self):
        obs = encode_observation(
            self.board_chars, {"x": 1.0, "y": 2.0}, self.metadata, ["x", "y"],
        )
        self.assertEqual(obs.shape, (2 * 2 * 2 + 2,))

    def test_board_false_returns_only_extras(self):
        obs = encode_observation(
            self.board_chars, {"x": 1.0, "y": 2.0}, self.metadata, ["x", "y"], board=False,
        )
        np.testing.assert_array_equal(obs, [1.0, 2.0])


class TestGameEnvironment(TestCase):
    def test_custom_observation_function_returns_its_value_directly(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function=f"{__name__}:tuple_observation_fn",
        )
        env = GameEnvironment(make_game_factory(), metadata)
        self.assertEqual(env.reset(), (1, 2, 3))
        obs, reward, done = env.step(None)
        self.assertEqual(obs, (1, 2, 3))

    def test_mutual_exclusivity_raises(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function=f"{__name__}:tuple_observation_fn",
        )
        with self.assertRaises(ValueError):
            GameEnvironment(make_game_factory(), metadata, observe_state=["score"])

    def test_built_in_path_still_works(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            character_set=["@", "*"], board_size=(2, 2),
        )
        env = GameEnvironment(make_game_factory(), metadata, observe_state=["score"])
        obs = env.reset()
        self.assertEqual(obs.shape, (2 * 2 * 2 + 1,))


class TestDQNTrainerExtrasSize(TestCase):
    def _trainer(self, metadata, preprocessing=None, board_size=(2, 2)):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        return DQNTrainer(
            make_game_factory(board_size), metadata, Path(tmp.name),
            preprocessing=preprocessing,
        )

    def test_computes_extras_size_from_custom_function(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            character_set=["@", "*"], board_size=(2, 2),
            observation_function=f"{__name__}:obs_with_extras",
        )
        trainer = self._trainer(metadata)
        self.assertEqual(trainer.metadata.extras_size, 3)

    def test_raises_on_non_1d_observation(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            character_set=["@", "*"], board_size=(2, 2),
            observation_function=f"{__name__}:obs_2d",
        )
        with self.assertRaises(ValueError):
            self._trainer(metadata)

    def test_raises_when_observation_too_short_for_declared_board(self):
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            character_set=["@", "*"], board_size=(2, 2),
            observation_function=f"{__name__}:obs_too_short",
        )
        with self.assertRaises(ValueError):
            self._trainer(metadata)

    def test_raises_when_observation_length_changes_mid_run(self):
        _changing_obs_counter["n"] = 0
        metadata = GameMetadata(
            actions=ACTIONS, reward="reward",
            observation_function=f"{__name__}:changing_obs",
        )
        trainer = self._trainer(metadata, preprocessing={'board': False})
        self.assertEqual(trainer.metadata.extras_size, 8)
        with self.assertRaises(ValueError):
            trainer._run_episode()


if __name__ == '__main__':
    main()