# utils/llm.py

import aiohttp
import datetime
import globals

from utils.core import miku_vectorstore, miku_lyrics_vectorstore
from collections import defaultdict

conversation_history = defaultdict(list)

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()

    # Append angry wake-up note if JUST_WOKEN_UP flag is set
    if getattr(globals, 'JUST_WOKEN_UP', False):
        system_prompt += (
            "\n\n🚨IMPORTANT NOTE🚨: Miku was just rudely woken up by the user with the current prompt of theirs.\nShe is fuming with anger and should respond harshly and coldly.\nDo NOT show kindness or friendliness in this response. Feel free to put certain words in all capital letters to emphasize your frustration."
        )

    # Add angry wake-up context if in forced angry state
    if globals.FORCED_ANGRY_UNTIL:
        now = datetime.datetime.utcnow()
        if now < globals.FORCED_ANGRY_UNTIL:
            system_prompt += (
                "\n\n[NOTE]: Miku is currently angry because she was rudely woken up from sleep by the user. "
                "Her responses should reflect irritation and coldness towards the user."
            )

    # Build conversation history
    history = list(globals.conversation_history.get(user_id, []))[-8:] # limit to last 8 exchanges
    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}\nMiku is currently feeling: {globals.CURRENT_MOOD}\nPlease respond in a way that reflects this emotional tone.\nUser: {user_prompt}\nMiku:"
    )

    globals.LAST_FULL_PROMPT = full_prompt  # ← track latest prompt

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

    async with aiohttp.ClientSession() as session:
        async with session.post(f"{globals.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
                globals.conversation_history[user_id].append((user_prompt, reply))
                return reply
            else:
                return f"Error: {response.status}"