"""
Main entry point for the FAQ bot using ChromaDB and a single MCP-tool.
"""

import os
import json
import sys
from typing import List, Dict, Any

import openai
from dotenv import load_dotenv

from chromadb_client import ChromadbClient
from mcp_tool import MCPTool

load_dotenv()

# Ensure OpenAI API key is set
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    print("Error: OPENAI_API_KEY not set in environment.", file=sys.stderr)
    sys.exit(1)

openai.api_key = OPENAI_API_KEY

# Initialize the ChromaDB client
db_client = ChromadbClient()

# Load FAQ documents from a local file (JSON lines format)
FAQ_FILE = os.getenv("FAQ_FILE", "data/faq.jsonl")

def load_faq_documents(file_path: str) -> List[Dict[str, Any]]:
    """
    Load FAQ documents from a JSON lines file.

    Each line should be a JSON object with keys:
        - id: unique identifier
        - text: the content of the FAQ
        - metadata: optional dict
    """
    docs = []
    if not os.path.exists(file_path):
        print(f"FAQ file {file_path} not found. Skipping load.", file=sys.stderr)
        return docs

    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            try:
                doc = json.loads(line.strip())
                docs.append(doc)
            except json.JSONDecodeError:
                continue
    return docs

# Load and add documents to the collection if not already present
if not db_client.collection.count():
    print("Loading FAQ documents into ChromaDB...")
    faq_docs = load_faq_documents(FAQ_FILE)
    if faq_docs:
        db_client.add_documents(faq_docs)
        print(f"Added {len(faq_docs)} documents.")
    else:
        print("No FAQ documents loaded.", file=sys.stderr)

# Instantiate the MCP-tool
mcp_tool = MCPTool()

# Define the function schema for OpenAI function calling
function_schema = {
    "name": mcp_tool.name,
    "description": mcp_tool.description,
    "parameters": {
        "type": "object",
        "properties": {},
        "required": [],
    },
}

def ask_question(question: str) -> str:
    """
    Ask a question to the bot. The bot will:
        1. Retrieve relevant FAQ documents from ChromaDB.
        2. Use OpenAI LLM to generate an answer, possibly invoking the MCP-tool.
    """
    # Retrieve top 3 relevant documents
    hits = db_client.query(question, top_k=3)

    # Build context from hits
    context = "\n\n".join([f"Document {hit['id']}:\n{hit['document']}" for hit in hits])

    # Construct the prompt for the LLM
    messages = [
        {"role": "system", "content": "You are an FAQ assistant. Use the provided documents to answer questions."},
        {"role": "user", "content": f"Question: {question}\n\nContext:\n{context}"},
    ]

    # Call OpenAI with function calling enabled
    response = openai.ChatCompletion.create(
        model="gpt-4o-mini",
        messages=messages,
        functions=[function_schema],
        function_call="auto",
    )

    # Parse the response
    reply = response["choices"][0]["message"]
    if reply.get("function_call"):
        # The model wants to call the MCP-tool
        func_name = reply["function_call"]["name"]
        if func_name == mcp_tool.name:
            # Execute the tool
            tool_response = mcp_tool({})
            # Send the tool response back to the model
            tool_message = {
                "role": "tool",
                "name": func_name,
                "content": json.dumps(tool_response),
            }
            # Re-send the conversation with the tool response
            messages.append(reply)
            messages.append(tool_message)
            # Get the final answer
            final_response = openai.ChatCompletion.create(
                model="gpt-4o-mini",
                messages=messages,
            )
            return final_response["choices"][0]["message"]["content"]
        else:
            return f"Unknown function call: {func_name}"
    else:
        return reply["content"]

def main():
    print("FAQ Bot (ChromaDB + MCP-tool). Type 'exit' to quit.")
    while True:
        try:
            user_input = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break

        if not user_input:
            continue

        answer = ask_question(user_input)
        print(f"\nBot: {answer}")

if __name__ == "__main__":
    main()