from logging import getLogger
from random import choice
from typing import Any

from taleweave.context import (
    get_agent_for_character,
    get_dungeon_master,
    get_game_systems,
    get_system_data,
)
from taleweave.game_system import GameSystem, SystemData
from taleweave.generate import generate_room
from taleweave.models.base import dataclass
from taleweave.models.entity import Room, World
from taleweave.utils.search import find_character_in_room, find_room
from taleweave.utils.systems import load_system_data, save_system_data
from taleweave.utils.template import format_prompt

logger = getLogger(__name__)

ABDUCTION_INTERVAL = 5  # turns

# set up data classes for the data that will be saved and loaded for this system


@dataclass
class AlienAbduction:
    abductee: str
    source: str


@dataclass
class AlienData:
    abduction: AlienAbduction | None = None
    laboratory: str | None = None


# system functions


def initialize_aliens(world: World):
    """
    Initialize the aliens system for the given world and create an alien laboratory room if one does not already exist.
    """
    data = get_system_data("aliens") or AlienData()

    if data.laboratory is None:
        logger.info("generating alien laboratory")
        dungeon_master = get_dungeon_master()
        systems = get_game_systems()  # TODO: other systems are not set up yet, this is always empty
        laboratory = generate_room(
            dungeon_master,
            world,
            systems,
            additional_prompt="make sure this room is an alien laboratory to which characters could be abducted"
        )
        world.rooms.append(laboratory)
        data.laboratory = laboratory.name

    return data


def abduct_character(world: World, laboratory_room: Room) -> AlienAbduction | None:
    """
    Abduct a random character from a random room and place them in the alien laboratory.
    """
    logger.info("abducting a character")
    valid_rooms = [room for room in world.rooms if room != laboratory_room and room.characters]
    if not valid_rooms:
        logger.info("no valid rooms to abduct from")
        return None

    source_room = choice(valid_rooms)
    abductee = choice(source_room.characters)

    logger.info(f"abducting {abductee.name}")
    source_room.characters.remove(abductee)
    laboratory_room.characters.append(abductee)

    # give the abductee a memory of their abduction
    abduction_memory = format_prompt("system_aliens_memory_abduction", abductee=abductee, laboratory_room=laboratory_room, source_room=source_room)
    abducted_agent = get_agent_for_character(abductee)
    if abducted_agent and abducted_agent.memory:
        abducted_agent.memory.append(abduction_memory)

    # TODO: give the abductee a chance to escape
    return AlienAbduction(abductee=abductee.name, source=source_room.name)


def return_abductee(world: World, laboratory_room: Room, abduction: AlienAbduction) -> AlienAbduction | None:
    """
    Return an abductee to their source room.
    """
    logger.info(f"returning {abduction.abductee} to their room")
    source_room = find_room(world, abduction.source)
    if not source_room:
        raise ValueError("Abductee's source room could not be found")

    abductee = find_character_in_room(laboratory_room, abduction.abductee)
    if not abductee:
        logger.info(f"abductee {abduction.abductee} is not in the laboratory")
        return

    # move the abductee back to their source room
    laboratory_room.characters.remove(abductee)
    source_room.characters.append(abductee)

    # give the abductee a memory of their return
    return_memory = format_prompt("system_aliens_memory_return", abductee=abductee, laboratory_room=laboratory_room, source_room=source_room)
    abducted_agent = get_agent_for_character(abductee)
    if abducted_agent and abducted_agent.memory:
        abducted_agent.memory.append(return_memory)

    return None


def simulate_aliens(world: World, turn: int, data: AlienData | None = None):
    """
    Simulate the alien abduction system. Every interval, abduct a character and place them in the alien laboratory or
    return them to their room.
    """
    data = data or get_system_data("aliens")
    if not data:
        raise ValueError("Alien data is required for simulation")

    # every N turns, randomly abduct a character and place them in the alien laboratory or return them to their room
    if data.laboratory and turn % ABDUCTION_INTERVAL == 0:
        laboratory_room = find_room(world, data.laboratory)
        if not laboratory_room:
            raise ValueError("Alien laboratory room not found")

        if data.abduction:
            data.abduction = return_abductee(world, laboratory_room, data.abduction)
        else:
            data.abduction = abduct_character(world, laboratory_room)


def load_aliens_data(file: str) -> Any:
    return load_system_data(AlienData, file)


def save_aliens_data(file: str, data: Any) -> None:
    return save_system_data(AlienData, file, data)


def init():
    return [GameSystem(
        name="aliens",
        data=SystemData(
            load=load_aliens_data,
            save=save_aliens_data,
        ),
        initialize=initialize_aliens,
        simulate=simulate_aliens,
    )]