import requests
import gradio as gr

# Configuration (Update this with your ngrok URL)
API_URL = "https://5b9f-154-161-168-18.ngrok-free.app/v1/chat/completions"
HEADERS = {"Content-Type": "application/json"}

# Store the last prompt
last_prompt = ""

def generate(prompt):
    """Generation function with prompt storage"""
    global last_prompt
    last_prompt = prompt
    try:
        response = requests.post(
            API_URL,
            headers=HEADERS,
            json={
                "model": "TinyLlama-1.1B-Chat-v1.0",
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 256,
                "temperature": 0.7
            },
            timeout=30
        )
        return response.json()['choices'][0]['message']['content']
    except Exception as e:
        return f"Error: {str(e)}"

def regenerate():
    """Regenerate response using last prompt"""
    if last_prompt:
        return generate(last_prompt)
    return "No previous prompt to regenerate"

# Gradio UI
with gr.Blocks(title="Karlson Achegeba GPT", theme=gr.themes.Soft(primary_hue="blue")) as app:
    with gr.Column(elem_classes=["center-container"]):
        # Header
        gr.Markdown("""<div style='text-align: center'>
            <h1 style='color: #2563eb'>Karlson Achegeba GPT</h1>
            <p>My very first trained model</p>
        </div>""")
        
        # Chat Interface
        with gr.Group():
            prompt = gr.Textbox(
                placeholder="Type your message...", 
                lines=5,
                label="Your Message",
                elem_classes=["input-box"]
            )
            
            with gr.Row():
                submit = gr.Button("Generate", variant="primary")
                regenerate_btn = gr.Button("Regenerate", variant="secondary")
                clear_btn = gr.Button("Clear All", variant="secondary")
            
            output = gr.Textbox(
                label="AI Response", 
                interactive=False,
                lines=8,
                elem_classes=["output-box"]
            )

    # Event Handlers
    submit.click(fn=generate, inputs=prompt, outputs=output)
    regenerate_btn.click(fn=regenerate, outputs=output)
    clear_btn.click(fn=lambda: ("", ""), outputs=[prompt, output])

    # CSS Styling
    app.css = """
    .center-container { max-width: 800px; margin: auto; padding: 20px; }
    .input-box, .output-box { 
        border-radius: 8px; 
        border: 1px solid #e2e8f0;
        padding: 12px;
    }
    .action-button { min-width: 120px; }
    .group { 
        background: white; 
        border-radius: 12px; 
        padding: 20px; 
        box-shadow: 0 2px 8px rgba(0,0,0,0.05);
    }
    """

app.launch(share = "True")