import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

# File to store model links
MODEL_FILE = "model_links.txt"

def load_model_links():
    # """Load model links from file"""
    # if not os.path.exists(MODEL_FILE):
    #     # Create default file with some example models
    #     with open(MODEL_FILE, "w") as f:
    #         f.write("facebook/opt-125m\n")
    #         f.write("facebook/opt-350m\n")
    
    with open(MODEL_FILE, "r") as f:
        return [line.strip() for line in f.readlines() if line.strip()]

class ModelManager:
    def __init__(self):
        self.current_model = None
        self.current_tokenizer = None
        self.current_model_name = None
    
    def load_model(self, model_name):
        """Load model and free previous model's memory"""
        if self.current_model is not None:
            del self.current_model
            del self.current_tokenizer
            torch.cuda.empty_cache()
        
        self.current_tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.current_model = AutoModelForCausalLM.from_pretrained(model_name)
        self.current_model_name = model_name
        return f"Loaded model: {model_name}"
    
    def generate_response(self, system_message, user_message):
        """Generate response from the model"""
        if self.current_model is None:
            return "Please select and load a model first."
        
        # Combine system and user messages
        prompt = f"{system_message}\n\nUser: {user_message}\n\nAssistant:"
        
        # Generate response
        inputs = self.current_tokenizer(prompt, return_tensors="pt", padding=True)
        outputs = self.current_model.generate(
            inputs.input_ids,
            max_length=200,
            num_return_sequences=1,
            temperature=0.7,
            pad_token_id=self.current_tokenizer.eos_token_id
        )
        
        response = self.current_tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Extract only the assistant's response
        response = response.split("Assistant:")[-1].strip()
        return response

# Initialize model manager
model_manager = ModelManager()

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Chat Interface with Model Selection")
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input components
            model_dropdown = gr.Dropdown(
                choices=load_model_links(),
                label="Select Model",
                info="Choose a model from the list"
            )
            load_button = gr.Button("Load Selected Model")
            system_msg = gr.Textbox(
                label="System Message",
                placeholder="Enter system message here...",
                lines=3
            )
            user_msg = gr.Textbox(
                label="User Message",
                placeholder="Enter your message here...",
                lines=3
            )
            submit_button = gr.Button("Generate Response")
        
        with gr.Column(scale=1):
            # Output components
            model_status = gr.Textbox(label="Model Status")
            chat_output = gr.Textbox(
                label="Assistant Response",
                lines=10,
                interactive=False
            )
    
    # Event handlers
    load_button.click(
        fn=model_manager.load_model,
        inputs=[model_dropdown],
        outputs=[model_status]
    )
    
    submit_button.click(
        fn=model_manager.generate_response,
        inputs=[system_msg, user_msg],
        outputs=[chat_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()