import discord
import aiohttp
import asyncio
import os
import re
import random
import string

import base64
import subprocess
import aiofiles

from langchain_community.vectorstores import FAISS
from langchain_ollama import OllamaEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.docstore.document import Document

from collections import defaultdict, deque

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from discord import File
from discord import Status
import datetime

from apscheduler.triggers.date import DateTrigger
from datetime import datetime, timedelta

scheduler = AsyncIOScheduler()
BEDTIME_CHANNEL_IDS = [761014220707332107]

# Stores last 5 exchanges per user (as deque)
conversation_history = defaultdict(lambda: deque(maxlen=5))

DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://ollama:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "mistral")

embeddings = OllamaEmbeddings(
    model=OLLAMA_MODEL,
    base_url=OLLAMA_URL
)

# Set up Discord client
intents = discord.Intents.default()
intents.message_content = True
intents.members = True
intents.presences = True
client = discord.Client(intents=intents)

current_model = None  # Track currently loaded model name

KINDNESS_KEYWORDS = [
    "thank you", "love you", "luv u", "you're the best", "so cute",
    "adorable", "amazing", "sweet", "kind", "great job", "well done",
    "precious", "good girl", "cutie", "angel", "my favorite", "so helpful"
]
HEART_REACTIONS = ["💙", "💝", "💖", "💕", "💜", "❤️‍🔥", "☺️"]
kindness_reacted_messages = set()

# Switch model
async def switch_model(model_name: str, timeout: int = 600):
    global current_model
    if current_model == model_name:
        print(f"🔁 Model '{model_name}' already loaded.")
        return

    # Unload all other models to clear VRAM
    async with aiohttp.ClientSession() as session:
        async with session.get(f"{OLLAMA_URL}/api/show") as resp:
            if resp.status == 200:
                data = await resp.json()
                loaded_models = data.get("models", [])
                for model in loaded_models:
                    if model["name"] != model_name:
                        print(f"🔁 Unloading model: {model['name']}")
                        await session.post(f"{OLLAMA_URL}/api/stop", json={"name": model["name"]})
            else:
                print("⚠️ Failed to check currently loaded models.")

    print(f"🔄 Switching to model '{model_name}'...")
    async with aiohttp.ClientSession() as session:
        await session.post(f"{OLLAMA_URL}/api/stop")
    # Warm up the new model (dummy call to preload it)
    payload = {
        "model": model_name,
        "prompt": "Hello",
        "stream": False
    }
    headers = {"Content-Type": "application/json"}

    # Poll until /api/generate returns 200
    async with aiohttp.ClientSession() as session:
        for _ in range(timeout):
            async with session.post(f"{OLLAMA_URL}/api/generate", json=payload, headers=headers) as resp:
                if resp.status == 200:
                    current_model = model_name
                    print(f"✅ Model {model_name} ready!")
                    return
            await asyncio.sleep(1)  # Wait a second before trying again

    raise TimeoutError(f"Timed out waiting for model '{model_name}' to become available.")


async def is_miku_addressed(message) -> bool:
    # If message is a reply, check the referenced message author
    if message.reference:
        try:
            referenced_msg = await message.channel.fetch_message(message.reference.message_id)
            if referenced_msg.author == message.guild.me:  # or client.user if you use client
                return True
        except Exception as e:
            print(f"⚠️ Could not fetch referenced message: {e}")

    cleaned = message.content.strip()

    return bool(re.search(
        r'(?<![\w\(])(?:[^\w\s]{0,2}\s*)?miku(?:\s*[^\w\s]{0,2})?(?=,|\s*,|[!\.?\s]*$)',
        cleaned,
        re.IGNORECASE
    ))

async def download_and_encode_image(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            if resp.status != 200:
                return None
            img_bytes = await resp.read()
            return base64.b64encode(img_bytes).decode('utf-8')

async def analyze_image_with_qwen(base64_img):
    await switch_model("moondream")

    payload = {
        "model": "moondream",
        "prompt": "Describe this image in detail.",
        "images": [base64_img],
        "stream": False
    }

    headers = {"Content-Type": "application/json"}

    async with aiohttp.ClientSession() as session:
        async with session.post(f"{OLLAMA_URL}/api/generate", json=payload, headers=headers) as response:
            if response.status == 200:
                data = await response.json()
                return data.get("response", "No description.")
            else:
                return f"Error: {response.status}"

async def rephrase_as_miku(qwen_output, user_prompt):
    await switch_model(OLLAMA_MODEL)  # likely llama3

    with open("miku_prompt.txt", "r", encoding="utf-8") as f:
        system_prompt = f.read()

    relevant_docs_lore = miku_vectorstore.similarity_search(qwen_output, k=3)
    context = "\n\n".join([doc.page_content for doc in relevant_docs_lore])

    full_prompt = (
        f"{context}\n\n"
        f"The user asked: \"{user_prompt}\"\n"
        f"The image contains: \"{qwen_output}\"\n\n"
        f"Respond like Miku: cheerful, helpful, and opinionated when asked.\n\n"
        f"Miku:"
    )

    payload = {
        "model": OLLAMA_MODEL,
        "prompt": full_prompt,
        "system": system_prompt,
        "stream": False
    }

    headers = {"Content-Type": "application/json"}

    async with aiohttp.ClientSession() as session:
        async with session.post(f"{OLLAMA_URL}/api/generate", json=payload, headers=headers) as response:
            if response.status == 200:
                data = await response.json()
                return data.get("response", "No response.")
            else:
                return f"Error: {response.status}"

# Load and index once at startup
def load_miku_knowledge():
    with open("miku_lore.txt", "r", encoding="utf-8") as f:
        text = f.read()

    from langchain.text_splitter import RecursiveCharacterTextSplitter

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=520,
        chunk_overlap=50,
        separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
    )

    docs = [Document(page_content=chunk) for chunk in text_splitter.split_text(text)]

    vectorstore = FAISS.from_documents(docs, embeddings)
    return vectorstore

def load_miku_lyrics():
    with open("miku_lyrics.txt", "r", encoding="utf-8") as f:
        lyrics_text = f.read()

    text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = [Document(page_content=chunk) for chunk in text_splitter.split_text(lyrics_text)]

    vectorstore = FAISS.from_documents(docs, embeddings)
    return vectorstore

miku_vectorstore = load_miku_knowledge()
miku_lyrics_vectorstore = load_miku_lyrics()

async def query_ollama(user_prompt, user_id):
    relevant_docs_lore = miku_vectorstore.similarity_search(user_prompt, k=3)
    relevant_docs_lyrics = miku_lyrics_vectorstore.similarity_search(user_prompt, k=3)

    context_lore = "\n".join([doc.page_content for doc in relevant_docs_lore])
    context_lyrics = "\n".join([doc.page_content for doc in relevant_docs_lyrics])

    combined_docs = relevant_docs_lore + relevant_docs_lyrics
    context = "\n\n".join([doc.page_content for doc in combined_docs])

    # Persona definition
    with open("miku_prompt.txt", "r", encoding="utf-8") as f:
        system_prompt = f.read()

    # Build conversation history
    history = conversation_history[user_id]
    history_text = "\n".join([f"User: {u}\nMiku: {m}" for u, m in history])

    # Combine prompt
    full_prompt = (
        f"{context_lore}\n\n{context_lyrics}\n\n"
        f"{history_text}\nUser: {user_prompt}\nMiku:"
    )

    headers = {'Content-Type': 'application/json'}
    payload = {
        "model": OLLAMA_MODEL,
        "prompt": full_prompt,
		"system": system_prompt,
        "stream": False
    }

    async with aiohttp.ClientSession() as session:
        async with session.post(f"{OLLAMA_URL}/api/generate", json=payload, headers=headers) as response:
            if response.status == 200:
                data = await response.json()
                reply = data.get("response", "No response.")
                # Save to conversation history
                conversation_history[user_id].append((user_prompt, reply))
                return reply
            else:
                return f"Error: {response.status}"

async def send_monday_video():
    await switch_model(OLLAMA_MODEL)

    # Generate a motivational message
    prompt = "It's Miku Monday! Give me an energetic and heartfelt Miku Monday morning message to inspire someone for the week ahead."
    response = await query_ollama(prompt, user_id="weekly-motivation")

    video_url = "http://zip.koko210cloud.xyz/u/zEgU7Z.mp4"

    target_channel_ids = [
        761014220707332107,
        1140377617237807266
    ]
    
    for channel_id in target_channel_ids:
        channel = client.get_channel(channel_id)
        if channel is None:
            print(f"❌ Could not find channel with ID {channel_id}. Make sure the bot is in the server.")
            return

        try:
            await channel.send(content=response)
            # Send video link
            await channel.send(f"[Happy Miku Monday!]({video_url})")

            print(f"✅ Sent Monday video to channel ID {channel_id}")
        except Exception as e:
            print(f"⚠️ Failed to send video to channel ID {channel_id}: {e}")

async def send_bedtime_reminder():
    await switch_model(OLLAMA_MODEL)

    for channel_id in BEDTIME_CHANNEL_IDS:
        channel = client.get_channel(channel_id)
        if not channel:
            print(f"⚠️ Channel ID {channel_id} not found.")
            continue

        guild = channel.guild

        # Filter online members (excluding bots)
        online_members = [
            member for member in guild.members
            if member.status in {Status.online, Status.idle, Status.dnd}
            and not member.bot
        ]

        specific_user_id = 214857593045254151 # target user ID
        specific_user = guild.get_member(specific_user_id)
        if specific_user and specific_user not in online_members:
            online_members.append(specific_user)

        if not online_members:
            print(f"😴 No online members to ping in {guild.name}")
            continue

        chosen_one = random.choice(online_members)

        # Generate bedtime message
        prompt = (
            f"Write a sweet, funny, or encouraging bedtime message to remind someone it's getting late and they should sleep. "
            f"Make it short and wholesome, as if Miku is genuinely worried about their well-being. Imply that it's not good staying up so late."
        )
        bedtime_message = await query_ollama(prompt, user_id="bedtime-miku")

        try:
            await channel.send(f"{chosen_one.mention}, {bedtime_message}")
            print(f"🌙 Sent bedtime reminder to {chosen_one.display_name} in {guild.name}")
        except Exception as e:
            print(f"⚠️ Failed to send bedtime reminder in {guild.name}: {e}")

def schedule_random_bedtime():
    now = datetime.now()
    target_time = now.replace(hour=20, minute=30, second=0, microsecond=0)
    # If it's already past 23:30 today, schedule for tomorrow
    if now > target_time:
        target_time += timedelta(days=1)

    # Add random offset (0–29 mins)
    offset_minutes = random.randint(0, 29)
    run_time = target_time + timedelta(minutes=offset_minutes)

    scheduler.add_job(send_bedtime_reminder, trigger=DateTrigger(run_date=run_time))
    print(f"⏰ Bedtime reminder scheduled for {run_time.strftime('%Y-%m-%d %H:%M:%S')}")

async def overlay_username_with_ffmpeg(base_video_path, output_path, username):
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    text = f"@{username}"

    # Define your six positions (x, y)
    positions = {
        1: ("250", "370"),
        2: ("330", "130"),
        3: ("300", "90"),
        4: ("380", "180"),
        5: ("365", "215"),
        6: ("55", "365"),
        7: ("290", "130"),
        8: ("320", "210"),
        9: ("310", "240"),
        10: ("400", "240")
    }

    # Each entry: (start_time, end_time, position_index)
    text_entries = [
        (4.767, 5.367, 1, "username"),
        (5.4, 5.967, 2, "username"),
        (6.233, 6.833, 3, "username"),
        (6.967, 7.6, 4, "username"),
        (7.733, 8.367, 5, "username"),
        (8.667, 9.133, 6, "username"),
        (9.733, 10.667, 7, "username"),
        (11.6, 12.033, 8, "@everyone"),
        (12.067, 13.0, 9, "@everyone"),
        (13.033, 14.135, 10, "@everyone"),
    ]

    # Build drawtext filters
    drawtext_filters = []
    for start, end, pos_id, text_type in text_entries:
        x_coord, y_coord = positions[pos_id]

        # Determine actual text content
        text_content = f"@{username}" if text_type == "username" else text_type

        x = f"{x_coord} - text_w/2"
        y = f"{y_coord} - text_h/2"

        filter_str = (
            f"drawtext=text='{text_content}':"
            f"fontfile='{font_path}':"
            f"fontcolor=black:fontsize=30:x={x}:y={y}:"
            f"enable='between(t,{start},{end})'"
        )
        drawtext_filters.append(filter_str)

    vf_string = ",".join(drawtext_filters)

    ffmpeg_command = [
        "ffmpeg",
        "-i", base_video_path,
        "-vf", vf_string,
        "-codec:a", "copy",
        output_path
    ]

    try:
        subprocess.run(ffmpeg_command, check=True)
        print("✅ Video processed successfully with username overlays.")
    except subprocess.CalledProcessError as e:
        print(f"⚠️ FFmpeg error: {e}")

async def detect_and_react_to_kindness(message, after_reply=False):
    if message.id in kindness_reacted_messages:
        return  # Already reacted — skip

    content = message.content.lower()

    emoji = random.choice(HEART_REACTIONS)

    # 1. Keyword-based detection
    if any(keyword in content for keyword in KINDNESS_KEYWORDS):
        try:
            await message.add_reaction(emoji)
            kindness_reacted_messages.add(message.id)
            message.kindness_reacted = True  # Mark as done
            print("✅ Kindness detected via keywords. Reacted immediately.")
        except Exception as e:
            print(f"⚠️ Error adding reaction: {e}")
        return

    # 2. If not after_reply, defer model-based check
    if not after_reply:
        print("🗝️ No kindness via keywords. Deferring...")
        return

    # 3. Model-based detection
    try:
        prompt = (
            "The following message was sent to Miku the bot. "
            "Does it sound like the user is being kind or affectionate toward Miku? "
            "Answer with 'yes' or 'no' only.\n\n"
            f"Message: \"{message.content}\""
        )
        result = await query_ollama(prompt, user_id="kindness-check")

        if result.strip().lower().startswith("yes"):
            await message.add_reaction(emoji)
            kindness_reacted_messages.add(message.id)
            print("✅ Kindness detected via model. Reacted.")
        else:
            print("🧊 No kindness detected.")
    except Exception as e:
        print(f"⚠️ Error during kindness analysis: {e}")

@client.event
async def on_ready():
    print(f'🎤 MikuBot connected as {client.user}')

    # Schedule the weekly task (Monday 07:30)
    scheduler.add_job(send_monday_video, 'cron', day_of_week='mon', hour=7, minute=30)

    # Schedule first bedtime reminder
    schedule_random_bedtime()
    # Reschedule every midnight
    scheduler.add_job(schedule_random_bedtime, 'cron', hour=0, minute=0)
    #scheduler.add_job(send_bedtime_reminder, 'cron', hour=12, minute=22)

    scheduler.start()

@client.event
async def on_message(message):
    if message.author == client.user:
        return
        
    if message.content.strip().lower() == "miku, rape this nigga balls" and message.reference:
        async with message.channel.typing():
            # Get replied-to user
            try:
                replied_msg = await message.channel.fetch_message(message.reference.message_id)
                target_username = replied_msg.author.display_name

                # Prepare video
                base_video = "MikuMikuBeam.mp4"
                output_video = f"/tmp/video_{''.join(random.choices(string.ascii_letters, k=5))}.mp4"

                await overlay_username_with_ffmpeg(base_video, output_video, target_username)

                caption = f"Here you go, @{target_username}! 🌟"
                #await message.channel.send(content=caption, file=discord.File(output_video))
                await replied_msg.reply(file=discord.File(output_video))

            except Exception as e:
                print(f"⚠️ Error processing video: {e}")
                await message.channel.send("Sorry, something went wrong while generating the video.")
            return


    text = message.content.strip()

    if await is_miku_addressed(message):   
 
        prompt = text  # No cleanup — keep it raw 
        
        # 1st kindness check with just keywords
        await detect_and_react_to_kindness(message)

        async with message.channel.typing():
            # If message has an image attachment
            if message.attachments:
                for attachment in message.attachments:
                    if any(attachment.filename.lower().endswith(ext) for ext in [".jpg", ".jpeg", ".png", ".webp"]):

                        base64_img = await download_and_encode_image(attachment.url)
                        if not base64_img:
                            await message.channel.send("I couldn't load the image, sorry!")
                            return
					# Analyze image (objective description)
                    qwen_description = await analyze_image_with_qwen(base64_img)
                    miku_reply = await rephrase_as_miku(qwen_description, prompt)
                    
                    await message.channel.send(miku_reply)
                    return

            # If message is just a prompt, no image
            response = await query_ollama(prompt, user_id=str(message.author.id))
            
            await message.channel.send(response)
            
        # 2nd kindness check (only if no keywords detected)
        await detect_and_react_to_kindness(message, after_reply=True)

    if message.content.lower().strip() == "!reset":
        conversation_history[str(message.author.id)].clear()
        await message.channel.send("Okay! Memory reset for you~ ✨")

    # Manual Monday test command
    if message.content.lower().strip() == "!monday":
        await send_monday_video()
        #await message.channel.send("✅ Monday message sent (or attempted). Check logs.")
        return

client.run(DISCORD_BOT_TOKEN)