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

# تنظیمات مدل‌ها
MODEL_CONFIG = {
    "codegen": {
        "name": "CodeGen-350M",
        "model_name": "Salesforce/codegen-350M-multi"
    },
    "codellama": {
        "name": "CodeLlama-7B", 
        "model_name": "codellama/CodeLlama-7b-hf"
    },
    "starcoder": {
        "name": "StarCoder2-3B",
        "model_name": "bigcode/starcoder2-3b"
    }
}

# بارگذاری مدل‌ها
try:
    print("🔄 در حال بارگذاری مدل CodeGen...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_CONFIG["codegen"]["model_name"])
    model = AutoModelForCausalLM.from_pretrained(MODEL_CONFIG["codegen"]["model_name"])
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    codegen_loaded = True
    print("✅ مدل CodeGen بارگذاری شد")
except Exception as e:
    print(f"❌ خطا در بارگذاری مدل CodeGen: {e}")
    codegen_loaded = False

# بارگذاری مدل‌های اختیاری دیگر
try:
    print("🔄 در حال بارگذاری مدل StarCoder2...")
    starcoder_pipeline = pipeline(
        "text-generation",
        model=MODEL_CONFIG["starcoder"]["model_name"],
        torch_dtype=torch.float16,
        device_map="auto"
    )
    starcoder_loaded = True
    print("✅ مدل StarCoder2 بارگذاری شد")
except Exception as e:
    print(f"⚠️ خطا در بارگذاری StarCoder2: {e}")
    starcoder_loaded = False

# دیکشنری برای ذخیره تاریخچه گفتگو
conversation_history = {}

def generate_with_starcoder(prompt, max_length=500, temperature=0.7):
    """تولید کد با استفاده از StarCoder2"""
    try:
        enhanced_prompt = f"# Write Python code for: {prompt}\n\n'''\nWrite clean, efficient Python code:\n{prompt}\n'''\n\ndef"
        
        response = starcoder_pipeline(
            enhanced_prompt,
            max_length=max_length,
            temperature=temperature,
            do_sample=True,
            num_return_sequences=1,
            pad_token_id=starcoder_pipeline.tokenizer.eos_token_id
        )
        
        generated_text = response[0]['generated_text']
        # استخراج فقط بخش کد
        if enhanced_prompt in generated_text:
            code = generated_text[len(enhanced_prompt):].strip()
        else:
            code = generated_text.strip()
            
        return code
    
    except Exception as e:
        return f"Error with StarCoder: {str(e)}"

def generate_with_codegen(prompt, user_id="default", max_length=300, temperature=0.7):
    """تولید کد با مدل CodeGen"""
    try:
        # ساخت prompt بهتر برای مدل
        enhanced_prompt = f"# Python code for: {prompt}\n'''\nWrite Python code to solve this problem:\n{prompt}\n'''\n\n"
        
        # اضافه کردن تاریخچه اگر وجود دارد
        if user_id in conversation_history and conversation_history[user_id]:
            history_text = "\n".join([f"Q: {q}\nA: {a}" for q, a in conversation_history[user_id][-3:]])
            enhanced_prompt = f"Previous conversation:\n{history_text}\n\nNew request: {prompt}\n\nCode:"
        
        # تولید کد با پارامترهای بهینه
        inputs = tokenizer.encode(enhanced_prompt, return_tensors="pt", max_length=512, truncation=True)
        
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=max_length,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                num_return_sequences=1,
                early_stopping=True
            )
        
        generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # استخراج فقط بخش کد
        if enhanced_prompt in generated_code:
            generated_code = generated_code[len(enhanced_prompt):].strip()
        
        return generated_code
    
    except Exception as e:
        return f"Error with CodeGen: {str(e)}"

def chat_with_ai(prompt, user_id="default", language="python", model_choice="codegen", max_length=500, temperature=0.7):
    """تابع اصلی چت با قابلیت انتخاب مدل"""
    
    if not prompt.strip():
        return "لطفاً یک درخواست برنامه‌نویسی وارد کنید."
    
    # تشخیص نوع درخواست
    prompt_lower = prompt.lower()
    
    if any(word in prompt_lower for word in ['سلام', 'hi', 'hello', 'سلامتی']):
        return "سلام 👋 به دستیار هوشمند برنامه‌نویسی خوش آمدید! میتوانید از من بخواهید کد بنویسم، باگ‌ها را رفع کنم، یا مفاهیم را توضیح دهم."
    
    elif any(word in prompt_lower for word in ['تشکر', 'ممنون', 'thanks']):
        return "خوشحالم که مفید بودم! اگر سوال دیگری دارید بپرسید."
    
    elif any(word in prompt_lower for word in ['خداحافظ', 'bye', 'خدانگهدار']):
        return "خدانگهدار! اگر باز هم نیاز به کمک داشتید در خدمتم."
    
    else:
        # تولید کد بر اساس مدل انتخاب شده
        try:
            if model_choice == "starcoder" and starcoder_loaded:
                code = generate_with_starcoder(prompt, max_length, temperature)
            elif model_choice == "codegen" and codegen_loaded:
                code = generate_with_codegen(prompt, user_id, max_length, temperature)
            else:
                # فال‌بک به مدل دیگر
                if codegen_loaded:
                    code = generate_with_codegen(prompt, user_id, max_length, temperature)
                elif starcoder_loaded:
                    code = generate_with_starcoder(prompt, max_length, temperature)
                else:
                    return "❌ هیچ مدلی در دسترس نیست. لطفاً مدل‌ها را بررسی کنید."
            
            # ذخیره در تاریخچه
            if user_id not in conversation_history:
                conversation_history[user_id] = []
            conversation_history[user_id].append((prompt, code))
            
            # فرمت کردن خروجی
            if code and not code.startswith("Error"):
                response = f"**🤖 مدل استفاده شده: {MODEL_CONFIG.get(model_choice, {}).get('name', model_choice)}**\n\n"
                response += f"**کد تولید شده:**\n```{language}\n{code}\n```\n\n"
                response += "💡 *نکته: همیشه کد را تست کنید و برای پروژه‌های مهم بررسی‌های بیشتری انجام دهید.*"
                return response
            else:
                return f"❌ خطا در تولید کد: {code}"
                
        except Exception as e:
            return f"خطا در پردازش: {str(e)}"

def clear_history(user_id="default"):
    """پاک کردن تاریخچه گفتگو"""
    if user_id in conversation_history:
        conversation_history[user_id] = []
    return "✅ تاریخچه گفتگو پاک شد."

def get_model_status():
    """دریافت وضعیت مدل‌ها"""
    status = {
        "CodeGen-350M": "✅ در دسترس" if codegen_loaded else "❌ غیرفعال",
        "StarCoder2-3B": "✅ در دسترس" if starcoder_loaded else "❌ غیرفعال",
        "مدل پیش‌فرض": "CodeGen-350M"
    }
    return status

# رابط Gradio
with gr.Blocks(theme=gr.themes.Soft(), title="دستیار هوشمند برنامه‌نویسی - Hugging Face") as demo:
    
    gr.Markdown("""
    # 🤖 دستیار هوشمند برنامه‌نویسی
    **نسخه Hugging Face - بدون نیاز به API خارجی**
    
    🔥 **پشتیبانی از مدل‌های محلی Hugging Face**
    """)
    
    with gr.Row():
        with gr.Column():
            user_id = gr.Textbox(label="شناسه کاربر", value="user_1")
            prompt_input = gr.Textbox(
                label="درخواست برنامه‌نویسی خود را وارد کنید",
                placeholder="مثال: تابعی بنویس که اعداد اول را پیدا کند...",
                lines=3
            )
            
            with gr.Row():
                model_choice = gr.Dropdown(
                    choices=["codegen", "starcoder"],
                    label="انتخاب مدل AI",
                    value="codegen",
                    info="CodeGen: سریع | StarCoder: قدرتمند"
                )
                
                language_select = gr.Dropdown(
                    choices=["python", "javascript", "html", "css", "sql"],
                    label="زبان برنامه‌نویسی",
                    value="python"
                )
            
            with gr.Row():
                temperature = gr.Slider(0.1, 1.0, value=0.7, label="میزان خلاقیت (Temperature)")
                max_length = gr.Slider(100, 1000, value=500, step=50, label="حداکثر طول خروجی")
            
            with gr.Row():
                submit_btn = gr.Button("🚀 تولید کد", variant="primary")
                clear_btn = gr.Button("🧹 پاک کردن تاریخچه", variant="secondary")
                status_btn = gr.Button("📊 وضعیت مدل‌ها", variant="secondary")
            
        with gr.Column():
            output = gr.Markdown(label="پاسخ دستیار")
            status_output = gr.JSON(label="وضعیت سیستم")
    
    # نمونه درخواست‌های پیش‌فرض
    gr.Examples(
        examples=[
            ["تابعی بنویس که عدد n را بگیرد و فاکتوریل آن را برگرداند"],
            ["یک برنامه بنویس که لیستی از اعداد را بگیرد و ماکزیمم را پیدا کند"],
            ["کلاس ساده برای مدیریت دانشجو در پایتون بنویس"],
            ["کد HTML برای یک صفحه لاگین ساده بنویس"]
        ],
        inputs=prompt_input
    )
    
    # رویدادها
    submit_btn.click(
        fn=chat_with_ai,
        inputs=[prompt_input, user_id, language_select, model_choice, max_length, temperature],
        outputs=output
    )
    
    clear_btn.click(
        fn=clear_history,
        inputs=user_id,
        outputs=output
    )
    
    status_btn.click(
        fn=get_model_status,
        inputs=[],
        outputs=status_output
    )
    
    # امکان ارسال با Enter
    prompt_input.submit(
        fn=chat_with_ai,
        inputs=[prompt_input, user_id, language_select, model_choice, max_length, temperature],
        outputs=output
    )

if __name__ == "__main__":
    print("🚀 در حال راه‌اندازی دستیار هوشمند برنامه‌نویسی...")
    print(f"📊 وضعیت مدل‌ها: {get_model_status()}")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )