# bot.py

import os
import asyncio
import discord
from discord import app_commands
import requests
from dotenv import load_dotenv

# --- Configuration ---
load_dotenv()

DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
PERPLEXICA_API_URL = "http://localhost:3000/api/search"

# Hardcoded model configurations
CHAT_MODEL = { "provider": "ollama", "name": "ministral-3:8b" }
EMBEDDING_MODEL = { "provider": "ollama", "name": "ryanshillington/Qwen3-Embedding-0.6B:latest" }

# In-memory store for user-selected optimization mode
user_modes = {}
DEFAULT_OPTIMIZATION_MODE = "balanced"
DEFAULT_FOCUS_MODE = "webSearch"

HELP_TEXT = """
Hello! I'm Perplexica Bot.
To use me, please type `/search <your query>` to search for information.
You can also change your search effort mode using `/mode`.
I don't respond to direct messages or mentions outside of these commands.
"""

# --- Discord Bot Setup ---
intents = discord.Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)
tree = app_commands.CommandTree(client)

@client.event
async def on_ready():
    """Event handler for when the bot is ready."""
    print(f'Logged in as {client.user}')
    await tree.sync()
    print("Slash commands synced.")

@client.event
async def on_message(message: discord.Message):
    """Event handler for when a message is sent."""
    # Ignore messages from the bot itself
    if message.author == client.user:
        return

    # Respond to DMs or mentions
    if isinstance(message.channel, discord.DMChannel) or client.user.mentioned_in(message):
        await message.channel.send(HELP_TEXT)
        return

    await client.process_commands(message) # Important! Allows other commands to be processed.

# --- Slash Commands ---

@tree.command(name="search", description="Search with Perplexica")
async def search(interaction: discord.Interaction, *, query: str):
    """Performs a search using the Perplexica API."""
    await interaction.response.send_message(f"<@{interaction.user.id}> 🧠 Thinking about your query: \"{query}\"...")
    
    # Get user's selected mode, default to balanced
    optimization_mode = user_modes.get(interaction.user.id, DEFAULT_OPTIMIZATION_MODE)

    # Construct the payload
    payload = {
        "query": query,
        "focusMode": DEFAULT_FOCUS_MODE,
        "optimizationMode": optimization_mode,
        "chatModel": CHAT_MODEL,
        "embeddingModel": EMBEDDING_MODEL,
        "stream": False
    }

    try:
        # Send request to Perplexica API
        response = requests.post(PERPLEXICA_API_URL, json=payload, timeout=300)
        response.raise_for_status()  # Raise an exception for bad status codes
        
        data = response.json()
        message = data.get("message", "No answer found.")
        sources = data.get("sources", [])

        # Format the response
        embed = discord.Embed(
            title=f"🔎 {query}",
            description=message,
            color=discord.Color.blue()
        )

        if sources:
            source_links = []
            for i, source in enumerate(sources, 1):
                title = source.get("title")
                url = source.get("url")
                if title and url:
                    source_links.append(f"[{i}] {title} ({url})")
            
            # Join sources and add to a field, handling character limits
            source_text = "\n".join(source_links)
            if len(source_text) > 1024:
                source_text = source_text[:1021] + "..."

            embed.add_field(name="🔗 Sources", value=source_text, inline=False)

        # Handle overall message length
        if len(embed) > 2000:
            original_description = embed.description

            # Calculate available space for description in the first message
            non_desc_len = len(embed) - len(original_description)
            
            # Max length for description, keeping total embed length under 2000.
            # With a small buffer for "..."
            available_len = 2000 - non_desc_len - 10
            
            # Ensure available_len is not negative
            if available_len < 0:
                available_len = 0
            
            first_part = original_description[:available_len]
            rest_of_message = original_description[available_len:]

            embed.description = first_part
            if rest_of_message:
                embed.description += "..."

            await interaction.edit_original_response(content=f"<@{interaction.user.id}>", embed=embed)

            if rest_of_message:
                # Split the rest of the message into chunks of 2000 characters
                chunks = [rest_of_message[i:i + 2000] for i in range(0, len(rest_of_message), 2000)]
                for i, chunk in enumerate(chunks):
                    await asyncio.sleep(1) # Add a small delay to prevent rate limiting issues
                    continuation_embed = discord.Embed(
                        title=f"🔎 {query} (continued {i + 1}/{len(chunks)})",
                        description=chunk,
                        color=discord.Color.blue()
                    )
                    await interaction.followup.send(content=f"<@{interaction.user.id}>", embed=continuation_embed)
        else:
             await interaction.edit_original_response(content=f"<@{interaction.user.id}>", embed=embed)

    except requests.exceptions.RequestException as e:
        await interaction.edit_original_response(content=f"An error occurred while contacting the Perplexica API: {e}")
    except Exception as e:
        await interaction.edit_original_response(content=f"An unexpected error occurred: {e}")


@tree.command(name="mode", description="Set the search optimization mode")
@app_commands.choices(mode=[
    app_commands.Choice(name="Balanced", value="balanced"),
    app_commands.Choice(name="Fast", value="speed"),
])
async def mode(interaction: discord.Interaction, mode: app_commands.Choice[str]):
    """Sets the optimization mode for the user."""
    user_modes[interaction.user.id] = mode.value
    await interaction.response.send_message(f"✅ Your search mode has been set to **{mode.name}**.", ephemeral=True)


# --- Run the Bot ---
if __name__ == "__main__":
    if DISCORD_BOT_TOKEN == "YOUR_DISCORD_BOT_TOKEN_HERE" or not DISCORD_BOT_TOKEN:
        print("ERROR: Please set your DISCORD_BOT_TOKEN in the .env file.")
    else:
        client.run(DISCORD_BOT_TOKEN)