pages.py

import json
from datetime import datetime
from typing import Any, Dict, List, Union


from autosim.agent import Agent
from autosim.base import Action, BaseSimulator, Movie, BaseView
from autosim.utils import (
    process_prompt_with_images,
)
from actions import *
from prompts import (
    MOVIE_DETAIL_VIEW_PROMPT,
    MOVIE_LIST_VIEW_PROMPT,
)

image_marker = "<IMAGE>"


class MovieListPage(BaseView):
    def __init__(
        self,
        movies: List[Movie],
        current_page: int,
        total_pages: int,
        name: str = "Movie List",
        prompt_template: str = MOVIE_LIST_VIEW_PROMPT,
    ):
        self.name = name
        self.movies = movies
        self.current_page = current_page
        self.total_pages = total_pages
        self.prompt_template = prompt_template

    def format_movie_details(self, movie: Movie) -> str:
        image_data = {"image_url": movie.poster_path} if movie.poster_path else None
        image_part = (
            f"{image_marker}{json.dumps(image_data)}{image_marker}\n"
            if image_data
            else ""
        )
        detail = (
            f"{image_part}"
            f"{movie.title}\n"
            f"ID: {movie.id}\n"
            f"Genres: {', '.join(movie.genres)}\n"
            f"Rating: {movie.rating}\n"
        )
        if movie.user_rating is not None:
            detail += f"Your Rating: {movie.user_rating}\n"
        return detail.strip()

    def format_movie_details(self, movie: Movie) -> str:
        image_data = {"image_url": movie.poster_path} if movie.poster_path else None
        image_part = (
            f"{image_marker}{json.dumps(image_data)}{image_marker}\n"
            if image_data
            else ""
        )

        title_line = movie.title
        if movie.user_rating is not None:
            title_line += f" (Viewed, Your Rating: {movie.user_rating})"

        detail = (
            f"{image_part}"
            f"{title_line}\n"
            f"ID: {movie.id}\n"
            f"Genres: {', '.join(movie.genres)}\n"
            f"Rating: {movie.rating}\n"
        )

        return detail.strip()

    def get_actions(self) -> List[Action]:
        actions = [click_movie, exit_action]
        if self.current_page > 1:
            actions.append(previous_page)
        if self.current_page < self.total_pages:
            actions.append(next_page)
        return actions

    def get_content(self) -> Union[str, List[Dict[str, Any]]]:
        mov_lst = [self.format_movie_details(movie) for movie in self.movies]
        prompt = self.prompt_template.format(
            movies="\n\n".join(mov_lst),
            current_page=self.current_page,
            total_pages=self.total_pages,
            actions=[action.to_dict() for action in self.get_actions()],
        )

        return process_prompt_with_images(prompt, image_marker=image_marker)

    def execute_action(self, action: Action, simulator: "BaseSimulator") -> BaseView:
        action_name = action.name
        action_args = action.parameters or {}
        page_number = self.current_page
        if action_name == next_page.name and self.current_page < self.total_pages:
            page_number += 1
        elif action_name == previous_page.name and self.current_page > 1:
            page_number -= 1
        elif action_name == click_movie.name:
            movie_id = action_args.get("movie_id")
            movie = simulator.get_movie_by_id(movie_id)
            return MovieDetailPage(movie=movie)
        else:
            raise ValueError(f"Invalid action: {action_name}")
        simulator.current_page = page_number
        return MovieListPage(
            movies=simulator.get_movies_in_page(page_number),
            current_page=page_number,
            total_pages=simulator.total_pages,
        )


class MovieDetailPage(BaseView):
    def __init__(
        self,
        movie: Movie,
        name: str = "Movie Detail",
        prompt_template: str = MOVIE_DETAIL_VIEW_PROMPT,
    ):
        self.name = name
        self.movie = movie
        self.prompt_template = prompt_template

    def get_actions(self) -> List[Action]:
        return [back_action, watch_and_rate_movie, exit_action]

    def get_content(self) -> str:
        image_data = (
            {"image_url": self.movie.poster_path} if self.movie.poster_path else None
        )
        image_part = (
            f"{image_marker}{json.dumps(image_data)}{image_marker}\n"
            if image_data
            else ""
        )
        detail = (
            f"{image_part}"
            f"{self.movie.title}\n"
            f"ID: {self.movie.id}\n"
            f"Genres: {', '.join(self.movie.genres)}\n"
            f"Rating: {self.movie.rating}\n"
        )
        if self.movie.user_rating is not None:
            detail += f"Your Rating: {self.movie.user_rating}\n"
        detail += (
            f"Vote Count: {self.movie.vote_count}\n"
            f"Release Date: {self.movie.release_date}\n"
            f"Director: {self.movie.directors}\n"
            f"Actors: {self.movie.actors}\n"
            f"Overview: {self.movie.overview}"
        )
        prompt = self.prompt_template.format(
            movie_detail=detail,
            actions=[action.to_dict() for action in self.get_actions()],
        )

        return process_prompt_with_images(prompt, image_marker=image_marker)

    def execute_action(self, action: Action, simulator: "BaseSimulator") -> BaseView:
        return MovieListPage(
            movies=simulator.get_movies_in_page(simulator.current_page),
            current_page=simulator.current_page,
            total_pages=simulator.total_pages,
        )


class MovieWebsite(BaseSimulator):
    def __init__(
        self,
        user: Agent,
        movies: List[Movie],
        n_movies_per_page: int = 5,
    ):
        self.user = user
        self.movies = movies
        self.movie_pages = self._create_movie_pages(movies, n_movies_per_page)
        self.current_page = 1
        self.logs = []

    @property
    def total_pages(self) -> int:
        return len(self.movie_pages)

    def _create_movie_pages(self, lst: List, n: int) -> List[List[Movie]]:
        """Chunks a list into smaller lists of a given size."""
        return [lst[i : i + n] for i in range(0, len(lst), n)]

    def get_initial_page(self) -> BaseView:
        """Returns the initial page view."""
        return MovieListPage(
            movies=self.get_movies_in_page(self.current_page),
            current_page=self.current_page,
            total_pages=self.total_pages,
        )

    def get_movie_by_id(self, movie_id: int) -> Movie:
        """Gets a movie by its ID."""
        for movie in self.movies:
            if movie.id == movie_id:
                return movie
        raise ValueError(f"Movie with ID {movie_id} not found.")

    def update_movie_rating(self, movie_id: int, rating: float) -> None:
        """Updates the user rating for a specific movie."""
        for movie in self.movies:
            if movie.id == movie_id:
                movie.user_rating = rating
                break

    def get_movies_in_page(self, page_number: int) -> List[Movie]:
        """Gets the movies for a given page number."""
        return self.movie_pages[page_number - 1]

    def log_action(self, action: Action) -> None:
        """Logs an action along with a message."""
        log_entry = {
            "datetime": datetime.now().isoformat(),
            "action": action.to_dict(),
            "token_usage": self.user.token_usage[-1],
        }
        self.logs.append(log_entry)

    def run(self) -> None:
        """Runs the simulator."""
        page = self.get_initial_page()
        while True:
            content = page.get_content()

            self.user.receive(content)
            messages, action = self.user.generate_reply(
                page=self.current_page, page_name=page.name
            )
            if action.name == watch_and_rate_movie.name:
                movie_id = action.parameters["movie_id"]
                rating = action.parameters["rating"]
                self.update_movie_rating(movie_id, rating)

            self.log_action(action)
            fatigue_module = self.user.working_memory.fatigue_module
            cumulative_fatigue = fatigue_module.cumulative_fatigue
            max_fatigue = fatigue_module.max_fatigue
            working_memory_length = self.user.working_memory.memory_length
            average_matching_score = self.user.working_memory.average_match_score
            if (
                action.name == exit_action.name
                or cumulative_fatigue >= max_fatigue
                or (working_memory_length >= 10 and average_matching_score <= 3)
                or working_memory_length >= 15
            ):
                break

            page = page.execute_action(action, self)
            self.user.clear_chat_messages()