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

# 🔹 Download and load model locally instead of using API
MODEL_NAME = "mistralai/Mistral-7B-Instruct"
MODEL_PATH = "./mistral-7b"

def download_model():
    """
    Download model locally if not already present.
    """
    if not os.path.exists(MODEL_PATH):
        print("Downloading model... This may take a while.")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
        tokenizer.save_pretrained(MODEL_PATH)
        model.save_pretrained(MODEL_PATH)
    else:
        print("Model already downloaded.")

download_model()

# 🔹 Load model from local storage
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH, 
    torch_dtype=torch.float16, 
    device_map="auto"
)

def generate_email_response(prompt, tone="professional, direct, concise"):
    """
    Generates an email response based on the provided prompt using Mistral 7B hosted locally.
    """
    input_text = f"[INST] You are an AI trained to respond like a business executive with a {tone} tone. {prompt} [/INST]"
    inputs = tokenizer(input_text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    outputs = model.generate(**inputs, max_length=200)
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def gradio_interface(prompt, tone):
    return generate_email_response(prompt, tone)

iface = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.Textbox(label="Email Prompt"), gr.Textbox(label="Tone", value="professional, direct, concise")],
    outputs=gr.Textbox(label="Generated Response"),
    title="AI Email Responder",
    description="Enter an email prompt and get an AI-generated response."
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", share=True)