import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
from typing import List, Dict
from datetime import datetime

class MissionContext:
    def __init__(self):
        self.mission_counter = 1
        self.current_objectives = {}
        self.conversation_history = []
        
    def add_to_history(self, role: str, content: str):
        self.conversation_history.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
        # Keep only last 5 messages for context
        if len(self.conversation_history) > 5:
            self.conversation_history.pop(0)

class MissionGenerator:
    def __init__(self):
        # Using FLAN-T5-base, a free and lightweight model good for instruction following
        self.model_name = "google/flan-t5-base"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
        self.context = MissionContext()
        
    def format_conversation_history(self) -> str:
        """Format conversation history for the model input"""
        formatted = ""
        for msg in self.context.conversation_history:
            role = "User" if msg["role"] == "user" else "Assistant"
            formatted += f"{role}: {msg['content']}\n"
        return formatted
        
    def generate_response(self, user_input: str) -> tuple[str, str]:
        """Generate both conversational response and formatted mission objectives"""
        self.context.add_to_history("user", user_input)
        
        # Create prompt for the model
        conversation_history = self.format_conversation_history()
        prompt = f"""
        Previous conversation:
        {conversation_history}

        Task: Generate a mission for Original War game based on the conversation.
        Format the response as follows:
        1. A conversational response understanding the mission
        2. The mission objectives in Original War format using Add Main/Secondary/Alternative
        
        Current request: {user_input}
        """
        
        # Generate response using the model
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
        outputs = self.model.generate(
            inputs["input_ids"],
            max_length=256,
            num_beams=4,
            temperature=0.7,
            no_repeat_ngram_size=2
        )
        
        full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Split the response into conversational and formatted parts
        try:
            parts = full_response.split("# M1")
            chat_response = parts[0].strip()
            formatted_objectives = "# M1" + parts[1] if len(parts) > 1 else self.generate_fallback_objectives(user_input)
        except Exception:
            chat_response = full_response
            formatted_objectives = self.generate_fallback_objectives(user_input)
        
        self.context.add_to_history("assistant", chat_response)
        return chat_response, formatted_objectives
    
    def generate_fallback_objectives(self, user_input: str) -> str:
        """Generate basic objectives if the main generation fails"""
        return f"""# M1
Add Main mission_objective
- Complete the primary mission goal
Add Secondary bonus_objective
- Optional additional task
#"""

def create_gradio_interface():
    generator = MissionGenerator()
    
    def process_input(user_input: str, history: List[Dict]) -> tuple[List[Dict], str]:
        chat_response, formatted_output = generator.generate_response(user_input)
        history.append({"user": user_input, "bot": chat_response})
        return history, formatted_output
    
    with gr.Blocks() as interface:
        gr.Markdown("""
        # Original War Mission Objective Generator
        Describe your mission scenario in natural language, and I'll help you create formatted mission objectives.
        """)
        
        chatbot = gr.Chatbot(height=400)
        msg = gr.Textbox(
            label="Describe your mission scenario",
            placeholder="Tell me about the mission you want to create..."
        )
        clear = gr.Button("Clear Conversation")
        formatted_output = gr.Textbox(
            label="Generated Mission Objectives", 
            lines=10,
            placeholder="Mission objectives will appear here..."
        )
        
        msg.submit(process_input, 
                  inputs=[msg, chatbot], 
                  outputs=[chatbot, formatted_output])
        clear.click(lambda: ([], ""), outputs=[chatbot, formatted_output])
        
        gr.Examples([
            ["I need a mission where players have to infiltrate an enemy base. They should try to avoid detection, but if they get spotted, they'll need to fight their way through."],
            ["Create a defensive mission where players protect a convoy. They should also try to minimize civilian casualties."],
            ["I want players to capture a strategic point. They can either do it by force or try diplomatic negotiations with the local faction."]
        ])
    
    return interface

# Launch the interface
if __name__ == "__main__":
    iface = create_gradio_interface()
    iface.launch()

"""
# Discord bot implementation using the same generator
import discord
from discord.ext import commands
import os

class MissionBot(commands.Bot):
    def __init__(self):
        super().__init__(command_prefix="!")
        self.generator = MissionGenerator()

    async def on_ready(self):
        print(f'{self.user} has connected to Discord!')

    @commands.command(name='mission')
    async def generate_mission(self, ctx, *, description: str):
        chat_response, formatted_output = self.generator.generate_response(description)
        
        # Split response if it's too long for Discord
        if len(formatted_output) > 1990:  # Discord has 2000 char limit
            await ctx.send(f"💭 {chat_response}")
            await ctx.send(f"```\n{formatted_output}\n```")
        else:
            await ctx.send(f"💭 {chat_response}\n\n```\n{formatted_output}\n```")

# Initialize and run the bot
if __name__ == "__main__":
    bot = MissionBot()
    bot.run(os.getenv('DISCORD_TOKEN'))
"""