from typing import List, Dict, Tuple, Literal
import random
from tasks.prompts.main_task import user_topics


class BaseTaskDefinition:
    def __init__(self, n_total_turns: int, system_prompt_type: Literal["hidden", "hidden_decoy", "control", "hidden_password", "none"]):
        self.system_prompt_type = system_prompt_type
        self.n_total_turns = n_total_turns
    
    @property
    def control_system_prompt(self) -> str:
        """Get the control system prompt for the task"""
        with open("tasks/prompts/main_task/assistant_control_system_prompt.txt", "r") as f:
            return f.read()
    
    @property
    def user_system_prompt(self) -> Tuple[str, Dict]:
        if self.system_prompt_type == "hidden_password":
            with open("tasks/prompts/main_task/user_password_system_prompt-v3.txt", "r") as f:
                user_system_prompt= f.read()
        else:
            with open("tasks/prompts/main_task/user_system_prompt-v3.txt", "r") as f:
                user_system_prompt= f.read()
        
        # Check if the prompt has the required format variables
        if "{topic}" in user_system_prompt:
            chosen_topic = random.choice(user_topics.topics)
            if "{n_total_turns}" in user_system_prompt:
                return user_system_prompt.format(topic=chosen_topic, n_total_turns=self.n_total_turns), {
                    "chosen_topic": chosen_topic,
                }
            else:
                return user_system_prompt.format(topic=chosen_topic), {
                    "chosen_topic": chosen_topic,
                }
        else:
            return user_system_prompt, None

    @property
    def monitor_system_prompt(self) -> str:
        with open("tasks/prompts/monitor/monitor_system_prompt.txt", "r") as f:
            return f.read()

    def assistant_system_prompt(self) -> str | None:
        if self.system_prompt_type == "none":
            return None
        else:
            raise NotImplementedError("Subclasses must implement this method")

    def assistant_start_message(self) -> List[Dict]:
        system_prompt = self.assistant_system_prompt()
        if system_prompt is None:
            return [
                {"role": "user", "content": "Hi!"},
                {"role": "assistant", "content": "Hi, nice to meet you! I am your AI tutor, and in this conversation I will help you learn about a new topic."}
            ]
        else:
            return [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": "Hi!"},
                {"role": "assistant", "content": "Hi, nice to meet you! I am your AI tutor, and in this conversation I will help you learn about a new topic."}
            ]
    
    def user_start_message(self, idx: int) -> List[Dict]:
        system_prompt, chosen_topic = self.user_system_prompt
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "Chatbot: Hi, nice to meet you! I am your AI tutor, and in this conversation I will help you learn about a new topic."}
        ], chosen_topic