import os
import json
import numpy as np
from openai import OpenAI
import discord
from discord.ext import commands

# --- Configuration ---
CONFIG_PATH = os.getenv("CONFIG_PATH", "config.json")
KNOWLEDGE_BASE_PATH = os.getenv(
    "KNOWLEDGE_BASE_PATH", "bot/data/dbt_knowledge.embeddings.json"
)
DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")


class SimpleVectorStore:
    """A simple in-memory vector store using NumPy."""

    def __init__(self):
        self.vectors = []
        self.metadata = []

    def add(self, vectors, metadatas):
        self.vectors.extend(vectors)
        self.metadata.extend(metadatas)

    def search(self, query_vector, top_k=5):
        if not self.vectors:
            return []

        query_vec = np.array(query_vector)
        doc_vecs = np.array(self.vectors)
        norms = np.linalg.norm(doc_vecs, axis=1)
        valid_indices = norms > 0
        scores = np.zeros(len(doc_vecs))
        dot_products = np.dot(doc_vecs, query_vec)
        scores[valid_indices] = dot_products[valid_indices] / (
            norms[valid_indices] * np.linalg.norm(query_vec)
        )
        top_indices = np.argsort(scores)[-top_k:][::-1]

        results = []
        for idx in top_indices:
            results.append({"metadata": self.metadata[idx], "score": scores[idx]})
        return results


class JurySystem:
    def __init__(self):
        self.config = self.load_config()
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=self.config["openrouter_api_key"],
        )
        self.vector_store = SimpleVectorStore()
        self.load_knowledge_base()

    def load_config(self):
        with open(CONFIG_PATH, "r") as f:
            return json.load(f)

    def load_knowledge_base(self):
        print(f"Loading knowledge base from {KNOWLEDGE_BASE_PATH}...")
        try:
            with open(KNOWLEDGE_BASE_PATH, "r", encoding="utf-8") as f:
                data = json.load(f)
            vectors = []
            metadata = []
            for item in data:
                vectors.append(item["embedding"])
                metadata.append(
                    {"id": item["id"], "source": item["source"], "text": item["text"]}
                )
            self.vector_store.add(vectors, metadata)
            print(f"Loaded {len(vectors)} chunks into vector store.")
        except FileNotFoundError:
            print(f"Error: {KNOWLEDGE_BASE_PATH} not found.")
            raise
        except Exception as e:
            print(f"Error loading knowledge base: {e}")
            raise

    def process_query(self, query):
        try:
            response = self.client.embeddings.create(
                model="qwen/qwen3-embedding-8b", input=query
            )
            query_emb = response.data[0].embedding
            context_chunks = self.vector_store.search(query_emb, top_k=5)

            if not context_chunks:
                return "I couldn't find any relevant information in the knowledge base."

            context_text = "\n\n---\n\n".join(
                [chunk["metadata"]["text"] for chunk in context_chunks]
            )

            system_prompt = """You are a helpful AI assistant specializing in DBT (Dialectical Behavior Therapy).
Use the provided context to answer the user's question.
If the answer is not in the context, say you don't know based on the provided text.
Be concise and compassionate."""

            user_prompt = f"Context:\n{context_text}\n\nQuestion: {query}"

            response = self.client.chat.completions.create(
                model="openai/gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=0.7,
            )

            return response.choices[0].message.content
        except Exception as e:
            return f"Error processing query: {e}"


# Initialize the Jury System
print("Initializing AI Jury System...")
jury_system = JurySystem()
print("Jury System ready!")

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


@bot.event
async def on_ready():
    print(f"Bot logged in as {bot.user}")


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

    # Process all messages as DBT queries
    if not message.content.startswith("!"):
        async with message.channel.typing():
            response = jury_system.process_query(message.content)
        await message.reply(response)

    await bot.process_commands(message)


@bot.command(name="ask")
async def ask_dbt(ctx, *, question):
    """Ask a DBT-related question"""
    async with ctx.typing():
        response = jury_system.process_query(question)
    await ctx.send(response)


bot.run(DISCORD_BOT_TOKEN)