import os
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel

# خواندن توکن از بخش Secrets تنظیمات اسپیس
HF_TOKEN = os.getenv("HF_TOKEN")

# شناسه‌های دقیق مدل‌ها
base_model_id = "meta-llama/Llama-3.2-1B-Instruct"
my_trained_model_id = "Alireza1913/Llama-3.2-1B-YouTube-Persona"

print("🤖 Step 1: Loading Tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(base_model_id, token=HF_TOKEN)

print("🤖 Step 2: Loading Base Llama Model...")
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    token=HF_TOKEN
)

print("🤖 Step 3: Merging Your Specific YouTube Persona Layers...")
# فراخوانی و ادغام مستقیم لایه‌های فاین‌تیون شده شما
model = PeftModel.from_pretrained(base_model, my_trained_model_id, token=HF_TOKEN)

print("🤖 Step 4: Initializing Text-Generation Pipeline...")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_youtube_idea(creator_name, video_topic):
    messages = [
        {"role": "system", "content": "You are an AI YouTube strategist. Expert in generating high-CTR titles and script hooks by perfectly mimicking top creators."},
        {"role": "user", "content": f"Creator: {creator_name}\nTopic: {video_topic}"}
    ]
    
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # اعمال تنظیماتی که قفل تولید متن مدل لاما را باز کردند
    outputs = generator(
        prompt, 
        max_new_tokens=150,
        min_new_tokens=40,
        temperature=0.8,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.2
    )
    
    reply = outputs['generated_text'][len(prompt):].strip()
    return reply if len(reply) > 0 else "Generation failed. Try another topic!"

# طراحی رابط کاربری گرافیکی شیک و مدرن با Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📹 YouTube Creator Mindset Engine (Llama-3.2-1B)")
    gr.Markdown("Select a top creator and input your topic to generate a specialized high-CTR title and viral hook instantly using your custom fine-tuned model.")
    
    with gr.Row():
        with gr.Column():
            creator = gr.Dropdown(["MrBeast", "Ali Abdaal", "MKBHD"], label="Choose Creator Style", value="MrBeast")
            topic = gr.Textbox(label="Video Topic / Concept", placeholder="e.g., Spending 24 hours in a luxury submarine...")
            btn = gr.Button("⚡ Generate Viral Strategy", variant="primary")
        
        with gr.Column():
            output = gr.TextArea(label="Generated Title & Hook (English)", lines=8)
            
    btn.click(generate_youtube_idea, inputs=[creator, topic], outputs=[output])

demo.launch()