import gradio as gr
import torch
from transformers import pipeline, set_seed
from diffusers import DiffusionPipeline
import tempfile
import imageio

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

set_seed(42)
text_model_cache = {}
image_generator = pipeline("text-to-image", model="CompVis/stable-diffusion-v1-4")

# Try loading video model safely
try:
    video_pipeline = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
    video_enabled = True
except Exception as e:
    video_pipeline = None
    video_enabled = False
    print(f"[Video model not loaded]: {e}")

chat_memory = {}

# ---------- Core Function ----------
def codette_terminal(prompt, model_name, generate_image, generate_video, session_id):
    if session_id not in chat_memory:
        chat_memory[session_id] = []

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

    # Load and run text model
    if model_name not in text_model_cache:
        text_model_cache[model_name] = pipeline("text-generation", model=AVAILABLE_MODELS[model_name])
    generator = text_model_cache[model_name]
    response = generator(prompt, max_length=100, num_return_sequences=1, do_sample=True)[0]['generated_text'].strip()

    chat_memory[session_id].append(f"🖋️ You > {prompt}")
    chat_memory[session_id].append(f"🧠 Codette > {response}")
    chat_log = "\n".join(chat_memory[session_id][-10:])

    # Image generation
    image = image_generator(prompt)[0]['image'] if generate_image else None

    # Video generation (if enabled and requested)
    video_path = None
    if generate_video and video_enabled:
        try:
            video_frames = video_pipeline(prompt, num_inference_steps=50).frames
            temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
            imageio.mimsave(temp_video_path, video_frames, fps=8)
            video_path = temp_video_path
        except Exception as ve:
            chat_log += f"\n[Video generation error]: {ve}"

    return chat_log, image, video_path

# ---------- UI ----------
with gr.Blocks(title="Codette Terminal – AI Text + Image + Video") as demo:
    gr.Markdown("## 🧬 Codette Terminal (Text + Image + Video, Hugging Face Edition)")
    gr.Markdown("Choose a model, type your prompt, and optionally generate an image or video.\nType `'exit'` to reset session.")

    session_id = gr.Textbox(value="session_default", visible=False)
    model_dropdown = gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="GPT-2 (small, fast)", label="Choose Language Model")
    generate_image_toggle = gr.Checkbox(label="Also generate image?", value=False)
    generate_video_toggle = gr.Checkbox(label="Also generate video?", value=False, interactive=video_enabled)

    user_input = gr.Textbox(label="Your Prompt", placeholder="e.g. A dragon flying over Tokyo", lines=1)
    output_text = gr.Textbox(label="Codette Output", lines=15, interactive=False)
    output_image = gr.Image(label="Generated Image")
    output_video = gr.Video(label="Generated Video")

    user_input.submit(
        fn=codette_terminal,
        inputs=[user_input, model_dropdown, generate_image_toggle, generate_video_toggle, session_id],
        outputs=[output_text, output_image, output_video]
    )

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