import openai
import gradio as gr
import os

# New API (openai>=1.0.0)
from openai import OpenAI

# Pricing for GPT-3.5-turbo (as of June 2024)
INPUT_COST_PER_1K = 0.0015
OUTPUT_COST_PER_1K = 0.002

def ask_question(api_key, user_question):
    if not api_key:
        return "❌ Please enter your OpenAI API key.", "", ""
    
    try:
        # Use the OpenAI client with your key
        client = OpenAI(api_key=api_key)

        # Send a chat completion request
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": user_question}
            ],
            temperature=0.5
        )

        # Extract the answer
        answer = response.choices[0].message.content

        # Token usage
        usage = response.usage
        input_tokens = usage.prompt_tokens
        output_tokens = usage.completion_tokens
        total_tokens = usage.total_tokens

        # Cost calculation
        input_cost = input_tokens * INPUT_COST_PER_1K / 1000
        output_cost = output_tokens * OUTPUT_COST_PER_1K / 1000
        total_cost = input_cost + output_cost

        token_info = f"🔢 Tokens Used: {total_tokens} (Input: {input_tokens}, Output: {output_tokens})"
        cost_info = f"💰 Estimated Cost: ${total_cost:.6f}"

        return answer, token_info, cost_info

    except Exception as e:
        return f"❌ Error: {str(e)}", "", ""

# Gradio Interface
iface = gr.Interface(
    fn=ask_question,
    inputs=[
        gr.Textbox(label="🔐 OpenAI API Key", type="password", placeholder="sk-..."),
        gr.Textbox(label="❓ Your Question", placeholder="Ask anything...")
    ],
    outputs=[
        gr.Textbox(label="🤖 GPT Response"),
        gr.Textbox(label="📊 Token Info"),
        gr.Textbox(label="💵 Cost Estimate")
    ],
    title="🧠 Ask GPT-3.5 with Your API Key",
    description="Uses latest openai>=1.0.0 SDK"
)

iface.launch()