from openai import OpenAI
from context import TWIN_SYSTEM_PROMPT
from tools import tools, handle_tool_calls
from styles import CSS, JS, EXAMPLES
from dotenv import load_dotenv
import gradio as gr
import spaces
import os

load_dotenv(override=True)

MODEL_NAME = os.getenv('MODEL_NAME')
if not MODEL_NAME:
    MODEL_NAME = "gpt-5.4-mini"

openai_api_key = os.getenv('OPENAI_API_KEY')
anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')
google_api_key = os.getenv('GOOGLE_API_KEY')
deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')
groq_api_key = os.getenv('GROQ_API_KEY')
grok_api_key = os.getenv('GROK_API_KEY')
openrouter_api_key = os.getenv('OPENROUTER_API_KEY')

ANTHROPIC_BASE_URL = "https://api.anthropic.com/v1/"
DEEPSEEK_BASE_URL = "https://api.deepseek.com/v1"
GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
GROQ_BASE_URL = "https://api.groq.com/openai/v1"
GROK_BASE_URL = "https://api.x.ai/v1"
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
OLLAMA_BASE_URL = "http://localhost:11434/v1"

if MODEL_NAME.startswith("gpt"):
    model = OpenAI()
elif MODEL_NAME.startswith("claude"):
    model = OpenAI(api_key=anthropic_api_key, base_url=ANTHROPIC_BASE_URL)
elif MODEL_NAME.startswith("gemini"):
    model = OpenAI(api_key=google_api_key, base_url=GEMINI_BASE_URL)
elif MODEL_NAME.startswith("deepseek"):
    model = OpenAI(api_key=deepseek_api_key, base_url=DEEPSEEK_BASE_URL)
elif MODEL_NAME.startswith("groq"):
    model = OpenAI(api_key=groq_api_key, base_url=GROQ_BASE_URL)
elif MODEL_NAME.startswith("grok"):
    model = OpenAI(api_key=grok_api_key, base_url=GROK_BASE_URL)
elif MODEL_NAME.startswith("openrouter"):
    model = OpenAI(api_key=openrouter_api_key, base_url=OPENROUTER_BASE_URL)
elif MODEL_NAME.startswith("ollama"):
    model = OpenAI(api_key=ollama_api_key, base_url=OLLAMA_BASE_URL)
else:
    model = OpenAI()

system = [{"role": "system", "content": TWIN_SYSTEM_PROMPT}]


@spaces.GPU
def chat(message, history):
    messages = system + history + [{"role": "user", "content": message}]
    response = model.chat.completions.create(model=MODEL_NAME, messages=messages, tools=tools)
    while response.choices[0].finish_reason == "tool_calls":
        message = response.choices[0].message
        tool_calls = message.tool_calls
        results = handle_tool_calls(tool_calls)
        messages.append(message)
        messages.extend(results)
        response = model.chat.completions.create(model=MODEL_NAME, messages=messages, tools=tools)
    return response.choices[0].message.content


if __name__ == "__main__":
    gr.ChatInterface(
        chat,
        examples=EXAMPLES,
        title="Digital Twin",
        description="Talk to my AI twin about my career",
        chatbot=gr.Chatbot(show_label=False),
    ).launch(css=CSS, js=JS, theme=gr.themes.Base())