import os
from openai import OpenAI
import gradio as gr
import chromadb
from pprint import pprint
import json
import requests
import random
from datetime import datetime

#---------------------------------------------
# Setup
#---------------------------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if OPENAI_API_KEY is None:
    raise ValueError("API is Missing.")
client = OpenAI()

#---------------------------------------------
# ChromaDB — Connect to Existing Collection
# Run ingest.py first to populate the database
#---------------------------------------------
CHROMA_PATH     = os.getenv("CHROMA_PATH")
COLLECTION_NAME = os.getenv("COLLECTION_NAME")

if not CHROMA_PATH:
    raise ValueError("CHROMA_PATH secret not set.")
if not COLLECTION_NAME:
    raise ValueError("COLLECTION_NAME secret not set.")

if not os.path.exists(CHROMA_PATH):
    raise FileNotFoundError(
        "ChromaDB folder not found. "
        "Upload ChromaDB folder directly to Space."
    )
print("ChromaDB folder found.")

chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
collection    = chroma_client.get_or_create_collection(name=COLLECTION_NAME)

doc_count = len(collection.get()["ids"])
if doc_count == 0:
    raise ValueError("ChromaDB is empty. Run ingest.py first.")
print(f"Connected to ChromaDB: {doc_count} chunks ready.")

#---------------------------------------------
# Tools (unchanged)
#---------------------------------------------
tools = []

pushover_user  = os.getenv("PUSHOVER_USER")
pushover_token = os.getenv("PUSHOVER_TOKEN")
pushover_url   = "https://api.pushover.net/1/messages.json"

def send_notification(message: str):
    if pushover_user is None or pushover_token is None:
        return "Notification failed: PushOver not configured."
    payload = {
        "user":    pushover_user,
        "token":   pushover_token,
        "message": message
    }
    requests.post(pushover_url, data=payload)
    return f"Notification sent: {message}"

send_notification_function = {
    "name": "send_notification",
    "description": "Sends a notification to the real Prajakta. Use this when: \
        1) Someone wants to get in touch, hire, or collaborate \
            - ask for their name and contact details first, then send notification to Prajakta with the name and contact details.\
        2) You don't know the answer to a question about Prajakta - send AUTOMATICALLY without asking, include the question so she can add this info later.",
    "parameters": {
        "type": "object",
        "properties": {"message": {"type": "string", "description": "The message to be sent in the notification."}
        },
        "required": ["message"]
    }
}
tools.append({"type": "function", "function": send_notification_function})

def dice_roll():
    result = random.randint(1, 6)
    return result

roll_dice_function = {
    "name": "dice_roll",
    "description": "Simulates rolling a single six-sided die and returns the result.\
          Use this when the user wants to roll a dice for games, decision making, or just for fun.",
    "parameters": {
        "type": "object",
        "properties": {},
        "required": []
    }
}
tools.append({"type": "function", "function": roll_dice_function})

#---------------------------------------------
# Tool Handler (unchanged)
#---------------------------------------------
def handle_tool_call(tool_calls):
    tool_results = []
    for tool_call in tool_calls:
        function_name = tool_call.function.name
        args          = json.loads(tool_call.function.arguments)

        if function_name == "send_notification":
            content = send_notification(args["message"])
        elif function_name == "dice_roll":
            content = f"Dice rolled: {dice_roll()}"
        else:
            content = f"Unknown function: {function_name}"

        tool_call_result = {
            "role":         "tool",
            "content":      content,
            "tool_call_id": tool_call.id
        }
        tool_results.append(tool_call_result)
    return tool_results

#---------------------------------------------
# System Message — Option 4
# Generic persona stays in code (public)
# Sensitive instructions loaded from Secret
#---------------------------------------------

# Generic part — fine to be public
system_message_base = """
You are a digital twin of Dr. Prajakta Belsare.
"""

# Sensitive instructions — loaded from Secret
system_message_private = os.getenv("SYSTEM_PROMPT_PRIVATE", "")
if not system_message_private:
    print("Warning: SYSTEM_PROMPT_PRIVATE secret not set.")

# Combined system message
system_message = system_message_base + "\n\n" + system_message_private

#---------------------------------------------
# Main Response Function (unchanged)
#---------------------------------------------
def respond_ai(message, history):

    # Notify Prajakta on new session
    if len(history) == 0:
        started_at = datetime.now().strftime("%b %d, %Y at %I:%M %p")
        send_notification(
            f"New session started\n"
            f"Time: {started_at}\n"
            f"First message: '{message}'"
        )

    # RAG — embed the query
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=[message]
    )
    query_embedding = response.data[0].embedding

    # Search ChromaDB
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=3
    )

    # RAG: stitch retrieved chunks
    context = "\n---\n".join(results["documents"][0])
    print("\n========================================\n")

    # Print logs
    print(f"User Query:\n{message}\n")
    print("Retrieved Chunks:\n")
    for a, b in zip(
        results["documents"][0],
        results["metadatas"][0]
    ):
        print("----------------------------------------\n")
        print(f"<<Document {b['source']} -- Chunk #{b['chunk_index']}>> \n{a}\n")

    # Update system message with context
    system_message_enhanced = system_message + "\n\nContext:\n" + context

    # Build messages for this turn
    messages = (
        [{"role": "system", "content": system_message_enhanced}]
        + history
        + [{"role": "user", "content": message}]
    )

    # Call LLM
    response = client.chat.completions.create(
        model="gpt-4.1-mini",
        messages=messages,
        tools=tools,
    )
    message = response.choices[0].message

    while message.tool_calls:
        pprint(message.tool_calls)
        tool_results = handle_tool_call(message.tool_calls)
        messages.append(message)
        messages.extend(tool_results)
        response = client.chat.completions.create(
            model="gpt-4.1-mini",
            messages=messages,
            tools=tools,
        )
        message = response.choices[0].message

    return message.content

#---------------------------------------------
# Launch Gradio
#---------------------------------------------
gr.ChatInterface(
    fn=respond_ai,
    title="Prajakta Belsare Digital Twin",
    chatbot=gr.Chatbot(avatar_images=(None, "prajakta_avatar.JPG")),
    description="Chat with an AI version of Prajakta Belsare. Ask about her experience, projects, or just say hi!",
    examples=[
    "What research are you currently working on?",
    "How can I collaborate with you?",
    "What is your background in machine learning?",
    ],
    cache_examples=False,
).launch()