import os
import gradio as gr
from transformers import pipeline

# Initialize the pipeline with settings that work well on Hugging Face
def initialize_model():
    # HF-specific configuration for memory efficiency
    pipe = pipeline(
        "text-generation",
        model="abhinand/tamil-llama-7b-instruct-v0.2",
        device_map="auto",
        torch_dtype="auto",
        model_kwargs={"load_in_8bit": True}  # 8-bit quantization for HF Spaces
    )
    return pipe

# Generate response
def generate_response(pipe, user_input, chat_history):
    # Format messages for the model
    messages = []
    for human, bot in chat_history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": bot})
    
    # Add the current message
    messages.append({"role": "user", "content": user_input})
    
    try:
        # Generate response with settings suitable for Spaces
        response = pipe(
            messages,
            max_length=256,  # Shorter responses to save compute
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1
        )
        
        # Extract the generated text
        generated_text = response[0]['generated_text']
        
        # Extract only the assistant's response
        for msg in generated_text:
            if isinstance(msg, dict) and msg.get("role") == "assistant":
                return msg.get("content", "")
        
        # Fallback if no assistant response is found
        if isinstance(generated_text, str):
            return generated_text
        return "சரியான பதிலைக் கண்டுபிடிக்க முடியவில்லை." # Could not find proper response
        
    except Exception as e:
        print(f"Error generating response: {e}")
        return f"பிழை ஏற்பட்டது. மீண்டும் முயற்சிக்கவும்." # Error occurred, please try again

# Create the Gradio interface
def create_chatbot_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# தமிழ் உரையாடல் பொத்தான் (Tamil Chatbot)")
        
        chatbot = gr.Chatbot(label="உரையாடல் (Conversation)")
        msg = gr.Textbox(label="உங்கள் செய்தி (Your Message)", placeholder="இங்கே தட்டச்சு செய்யவும்...")
        clear = gr.Button("அழி (Clear)")
        
        # Initialize model only once when first needed
        model = gr.State(None)
        
        def load_model_if_needed(model_state):
            if model_state is None:
                return initialize_model()
            return model_state
        
        def respond(message, chat_history, model_state):
            # Load model if not already loaded
            if model_state is None:
                model_state = initialize_model()
                
            bot_message = generate_response(model_state, message, chat_history)
            chat_history.append((message, bot_message))
            return "", chat_history, model_state
        
        msg.submit(respond, [msg, chatbot, model], [msg, chatbot, model])
        clear.click(lambda: None, None, chatbot, queue=False)
    
    return demo

# Create and launch the demo
demo = create_chatbot_interface()

# This is the key part for Hugging Face Spaces
if __name__ == "__main__":
    demo.launch()