import gradio as gr
import os
import re
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ.get("OPENROUTER_API_KEY"),
)

def clean_response(text: str) -> str:
    lines = []
    for line in text.splitlines():
        line = re.sub(r"^\s*[*#•\-]+\s*", "", line)
        line = re.sub(r"\*\*(.*?)\*\*", r"\1", line)
        # Only strip double underscores around a phrase. A single token like
        # __init__ is a Python dunder, which shares the exact shape of
        # __underline__ and cannot be told apart by position alone.
        line = re.sub(r"(?<!\w)__([^_\n]*?\s[^_\n]*?)__(?!\w)", r"\1", line)
        lines.append(line)
    return "\n".join(lines)

MAX_SENT_TURNS = 20


def chat(user_message: str, history: list):
    history = history or []

    # Cap what is sent, not what is shown: unbounded history eventually exceeds
    # the model's context window, and the failure then looks like a model error.
    recent = history[-(MAX_SENT_TURNS * 2):]
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for msg in recent:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})

    history = history + [{"role": "user", "content": user_message}]

    try:
        # openrouter/free automatically picks a working free model
        stream = client.chat.completions.create(
            model="openrouter/free",
            messages=messages,
            stream=True,
            temperature=0.7,
            max_tokens=4096,
        )

        response_text = ""
        history = history + [{"role": "assistant", "content": ""}]
        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                response_text += chunk.choices[0].delta.content
                history[-1]["content"] = response_text
                yield history

        history[-1]["content"] = clean_response(response_text)
        yield history

    except Exception as e:
        history.append({
            "role": "assistant",
            "content": f"Error: {str(e)}"
        })
        yield history

with gr.Blocks(title="OpenRouter Chat") as demo:
    gr.Markdown("## OpenRouter Chat")
    
    chatbot = gr.Chatbot(
        label="Conversation",
        height=600
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Ask something",
            placeholder="Type your prompt and press Enter",
            scale=4,
            container=False
        )
        clear = gr.Button("Clear", scale=1)

    msg.submit(chat, [msg, chatbot], chatbot).then(
        lambda: "", None, msg
    )
    clear.click(lambda: [], None, chatbot)

if __name__ == "__main__":
    demo.launch()