import gradio as gr
from transformers import pipeline, set_seed

# Preload models into a dictionary for selection
AVAILABLE_MODELS = {
    "GPT-2 (small, fast)": "gpt2",
    "Falcon (TII UAE)": "tiiuae/falcon-7b-instruct",
    "Mistral (OpenAccess)": "mistralai/Mistral-7B-v0.1"
}

set_seed(42)

# Cache loaded models to avoid reloading
model_cache = {}

def load_model(model_name):
    if model_name not in model_cache:
        model_id = AVAILABLE_MODELS[model_name]
        model_cache[model_name] = pipeline("text-generation", model=model_id)
    return model_cache[model_name]

# Chat history buffer
chat_memory = {}

# Terminal response function
def codette_terminal(user_input, model_name, session_id):
    if session_id not in chat_memory:
        chat_memory[session_id] = []

    if user_input.lower() in ['exit', 'quit']:
        chat_memory[session_id] = []
        return "🧠 Codette signing off... Session reset."

    generator = load_model(model_name)
    output = generator(user_input, max_length=100, num_return_sequences=1, do_sample=True)
    response = output[0]['generated_text'].strip()

    chat_memory[session_id].append(f"🖋️ You > {user_input}")
    chat_memory[session_id].append(f"🧠 Codette > {response}")
    return "\n".join(chat_memory[session_id][-10:])  # limit memory output

# UI
with gr.Blocks(title="Codette Terminal - Multi-Model AI") as demo:
    gr.Markdown("## 🧬 Codette Terminal (Multi-Model, Hugging Face Edition)")
    gr.Markdown("Type something below. Use the dropdown to switch models. Type `'exit'` to reset session.")

    session_id = gr.Textbox(label="Session ID (hidden)", value="session_default", visible=False)
    model_select = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Choose AI Model")
    user_input = gr.Textbox(label="Your Prompt", placeholder="Ask me anything...")
    output_box = gr.Textbox(label="Terminal Output", lines=20, interactive=False)

    user_input.submit(fn=codette_terminal, inputs=[user_input, model_select, session_id], outputs=output_box)

# For Hugging Face Spaces
if __name__ == "__main__":
    demo.launch()