import discord
from discord.ext import commands
from openai import OpenAI
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure OpenAI client to point to OpenWebUI
client = OpenAI(
    api_key=os.getenv('OPENAI_API_KEY'),
    base_url=os.getenv('OPENWEBUI_API_BASE')  # e.g., "http://localhost:8080/v1"
)

# Initialize Discord bot
intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix="!", intents=intents)

# Add a dictionary to store conversation histories
conversation_histories = {}
DEFAULT_HISTORY_LIMIT = 50
MAX_HISTORY_LIMIT = 200

async def get_ai_response(prompt, channel_id=None, include_history=False, history_limit=DEFAULT_HISTORY_LIMIT):
    try:
        messages = []
        
        # If history is requested and exists for this channel, include it
        if include_history and channel_id in conversation_histories:
            messages = conversation_histories[channel_id]
        
        # Add the current prompt
        messages.append({"role": "user", "content": prompt})
        
        response = client.chat.completions.create(
            model=os.getenv('MODEL_NAME') or "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
            messages=messages,
            temperature=0.7,
            max_tokens=500
        )
        
        # Store the conversation history if channel_id is provided
        if channel_id:
            if channel_id not in conversation_histories:
                conversation_histories[channel_id] = []
            conversation_histories[channel_id].extend([
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": response.choices[0].message.content}
            ])
            
            # Limit history based on specified or default limit
            if len(conversation_histories[channel_id]) > history_limit * 2:  # multiply by 2 because each exchange has 2 messages
                conversation_histories[channel_id] = conversation_histories[channel_id][-(history_limit * 2):]
        
        return response.choices[0].message.content
    except Exception as e:
        print(f"Error getting AI response: {e}")
        return "Sorry, I encountered an error while processing your request."

@bot.event
async def on_message(message):
    if message.author == bot.user:
        return

    if bot.user in message.mentions:
        prompt = message.content.replace(f'<@{bot.user.id}>', '').strip()
        
        if not prompt:
            await message.channel.send("Hello! How can I help you?")
            return

        # Check if the message includes a request for history
        include_history = False
        history_limit = DEFAULT_HISTORY_LIMIT
        
        # Check for "with history", "with X history", or "X lines of chat" patterns
        import re
        if "with history" in prompt.lower() or re.search(r"\d+\s*lines of chat", prompt.lower()):
            include_history = True
            # Check for specific history limit
            if match := re.search(r"with (\d+) history", prompt.lower()):
                requested_limit = int(match.group(1))
                history_limit = min(requested_limit, MAX_HISTORY_LIMIT)
                prompt = re.sub(r"with \d+ history", "", prompt, flags=re.IGNORECASE)
            elif match := re.search(r"(\d+)\s*lines of chat", prompt.lower()):
                requested_limit = int(match.group(1))
                history_limit = min(requested_limit, MAX_HISTORY_LIMIT)
                prompt = re.sub(r"\d+\s*lines of chat", "", prompt, flags=re.IGNORECASE)
            else:
                prompt = prompt.lower().replace("with history", "")
                
        prompt = prompt.strip()

        async with message.channel.typing():
            response = await get_ai_response(
                prompt,
                channel_id=str(message.channel.id),
                include_history=include_history,
                history_limit=history_limit
            )
            
            if len(response) > 2000:
                chunks = [response[i:i+2000] for i in range(0, len(response), 2000)]
                for chunk in chunks:
                    await message.channel.send(chunk)
            else:
                await message.channel.send(response)

    await bot.process_commands(message)

@bot.command(name='clearhistory')
async def clear_history(ctx):
    channel_id = str(ctx.channel.id)
    if channel_id in conversation_histories:
        conversation_histories[channel_id] = []
        await ctx.send("Conversation history has been cleared.")
    else:
        await ctx.send("No conversation history exists for this channel.")

def main():
    # Get the Discord token from environment variables
    discord_token = os.getenv('DISCORD_TOKEN')
    if not discord_token:
        raise ValueError("Discord token not found in environment variables")

    # Run the bot
    bot.run(discord_token)

if __name__ == "__main__":
    main()