# wcs/games/content/ai.py
"""Live content source of the Millionaire game: one continuous AI chat per game session (see content/__init__)."""
import json
import random
import threading
import uuid

from wcs.ai import AIHandler as ai_handler
from wcs.games import millionaire_prompts as prompts
from wcs.games.voice import TtsVoice

LETTERS = ('A', 'B', 'C', 'D')


class AiContent:
    name = 'ai'
    badge = ''
    rules_note = 'AI vytváří unikátní otázky'

    def __init__(self, media_type, title, plot='', genre='', episode_title='', season_number=None,
                 episode_number=None, actor_data=None):
        self.media_type = media_type
        self.title = title
        self.plot = plot
        self.genre = genre
        self.episode_title = episode_title
        self.season_number = season_number
        self.episode_number = episode_number
        self.actor_data = actor_data or {}
        self.voice = TtsVoice()
        self.system_prompt = ''
        self.session_id = ''
        self._reset_chat()

    def begin(self):
        """New game: fresh chat history, a session id for cumulative usage tracking and the fixed system prompt."""
        self._reset_chat()
        self.session_id = f"millionaire_{uuid.uuid4().hex[:12]}"
        from wcs.ai.UsageTracker import describe_session     # the game's calls carry its title in the AI cost history
        describe_session(self.session_id, self.title)
        self.system_prompt = prompts.build_system_prompt(
            media_type=self.media_type, title=self.title, plot=self.plot, genre=self.genre,
            episode_title=self.episode_title, season_number=self.season_number,
            episode_number=self.episode_number, actor_data=self.actor_data)

    def _reset_chat(self):
        self.chat_history = []
        # a fresh order of question angles for every game (variety when a title is played again)
        self.angles = (random.sample(prompts.QUESTION_ANGLES, 9) if self.media_type in prompts.ANGLE_MEDIA
                       else [None] * 9)
        self.correct_counts = dict.fromkeys(LETTERS, 0)
        self._lock = threading.Lock()
        self._asked = 0              # exchanges started
        self._joined = 0             # exchanges already in chat_history
        self._finished = {}          # exchange number -> (user, assistant) message pair, None when it failed

    def welcome(self):
        return self._chat(prompts.build_welcome_message())

    def question(self, round_num):
        # the chat keeps the order the player sees (the feedback messages name these letters)
        angle = self.angles[round_num - 1] if 1 <= round_num <= 9 else None
        return self._chat(prompts.build_question_message(round_num, angle), transform=self._place_correct)

    def _place_correct(self, data):
        """Move the right answer to a random letter that has been right least often in this game.

        Models put the right answer mostly under A or B; the position is chosen here instead, evenly enough that the
        letters stay balanced over nine rounds yet cannot be predicted.
        """
        options = data.get('options')
        correct = str(data.get('correct', '')).strip().upper()
        if not isinstance(options, dict) or sorted(options) != list(LETTERS) or correct not in LETTERS:
            return data                                          # unexpected shape: the controller handles it as is
        fewest = min(self.correct_counts.values())
        target = random.choice([letter for letter in LETTERS if self.correct_counts[letter] <= fewest + 1])
        wrong = [options[letter] for letter in LETTERS if letter != correct]
        random.shuffle(wrong)
        placed = {}
        for letter in LETTERS:
            placed[letter] = options[correct] if letter == target else wrong.pop()
        self.correct_counts[target] += 1
        result = dict(data)
        result['options'] = placed
        result['correct'] = target
        return result

    def feedback(self, round_num, answer_index, is_correct, question):
        options = question.get('options', [])
        correct_index = question.get('correct', 0)
        letters = ["A", "B", "C", "D"]
        return self._chat(prompts.build_feedback_message(
            round_num=round_num,
            player_answer_letter=letters[answer_index] if answer_index < 4 else "A",
            player_answer_text=options[answer_index] if answer_index < len(options) else "",
            correct_answer_letter=letters[correct_index] if correct_index < 4 else "A",
            correct_answer_text=options[correct_index] if correct_index < len(options) else "",
            is_correct=is_correct,
            is_last_round=(round_num >= 9)))

    def tip(self, question):
        return ai_handler.generate_director_tip(
            media_type=self.media_type, title=self.title, question=question['question'], options=question['options'],
            plot=self.plot, genre=self.genre, episode_title=self.episode_title, season_number=self.season_number,
            episode_number=self.episode_number, session_id=self.session_id)

    def _chat(self, user_message, transform=None):
        """One exchange of the continuous chat: ask the AI, return the parsed JSON (after `transform`).

        Exchanges may run at the same time (the next question is asked while the feedback is written, the first one
        with the welcome). A request sees the exchanges finished before it; the answers join the chat in the order
        they were asked, so the chat always alternates user / assistant.
        """
        user = {"role": "user", "content": user_message}
        with self._lock:
            number = self._asked
            self._asked += 1
            messages = self.chat_history + [user]
        pair = None
        try:
            data, _usage = ai_handler.generate_millionaire_chat(self.system_prompt, messages,
                                                                session_id=self.session_id)
            if transform:
                data = transform(data)
            pair = (user, {"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
            return data
        finally:
            with self._lock:
                self._finished[number] = pair
                while self._joined in self._finished:
                    joined = self._finished.pop(self._joined)
                    if joined:
                        self.chat_history = self.chat_history + list(joined)
                    self._joined += 1