import os
import gradio as gr
from google import genai
import time
import requests
from huggingface_hub import HfApi, create_repo

# ================= تنظیمات =================
HF_TOKEN = os.environ.get("HF_TOKEN", "")
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")

if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY not set!")

client = genai.Client(api_key=GOOGLE_API_KEY)

USER_NAME = "Vgjkmhf"
MAX_RETRIES = 10


def extract_code(text):
    import re
    pattern = r"```(?:python)?\n(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    if matches:
        return matches[0].strip()
    return text.strip()


def generate_code(prompt, previous_code=None, error=None):
    if previous_code and error:
        message = f"""Fix this Python/Gradio code that has an error.

Error:
{error}

Code:
{previous_code}

Return ONLY the corrected Python code inside ```python``` block. No explanations."""
    else:
        message = f"""Write a complete Python Gradio app for: {prompt}

Requirements:
- Use gradio as the UI framework
- Include `demo.launch()` at the end
- Make it functional and complete
- Return ONLY the code inside ```python``` block"""

    response = client.models.generate_content(
        model="gemini-3.1-pro-preview",
        contents=message
    )
    return extract_code(response.text)


def create_space(space_name):
    api = HfApi(token=HF_TOKEN)
    try:
        create_repo(
            repo_id=f"{USER_NAME}/{space_name}",
            repo_type="space",
            space_sdk="gradio",
            private=False,
            exist_ok=True
        )
        return True, f"{USER_NAME}/{space_name}"
    except Exception as e:
        return False, str(e)


def upload_to_space(code, space_id):
    api = HfApi(token=HF_TOKEN)
    try:
        api.upload_file(
            path_or_fileobj=code.encode(),
            path_in_repo="app.py",
            repo_id=space_id,
            repo_type="space",
        )
        api.upload_file(
            path_or_fileobj=b"gradio\ngoogle-genai\nhuggingface_hub\n",
            path_in_repo="requirements.txt",
            repo_id=space_id,
            repo_type="space",
        )
        return True, ""
    except Exception as e:
        return False, str(e)


def get_space_status(space_id):
    try:
        url = f"https://huggingface.co/api/spaces/{space_id}"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
        r = requests.get(url, headers=headers, timeout=30)
        if r.status_code == 200:
            data = r.json()
            return data.get("stage", "unknown")
    except:
        pass
    return "unknown"


def get_space_logs(space_id):
    try:
        url = f"https://huggingface.co/api/spaces/{space_id}/logs"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
        r = requests.get(url, headers=headers, timeout=30)
        if r.status_code == 200:
            return r.text
    except:
        pass
    return ""


def build_and_deploy(user_request):
    if not user_request.strip():
        yield "⚠️ لطفاً یه درخواست بنویس!", ""
        return

    logs = ""
    code = ""
    space_name = f"app-{int(time.time())}"

    def log(msg):
        nonlocal logs
        logs += msg + "\n"
        return logs

    # مرحله ۱: تولید کد
    logs = log(f"📝 درخواست: {user_request}")
    logs = log("─" * 50)
    logs = log("🤖 در حال تولید کد با Gemini 3.1 Pro...")
    yield logs, ""

    try:
        code = generate_code(user_request)
        logs = log("✅ کد اولیه تولید شد!")
        yield logs, code
    except Exception as e:
        logs = log(f"❌ خطا در تولید کد: {e}")
        yield logs, ""
        return

    for attempt in range(1, MAX_RETRIES + 1):
        logs = log(f"\n{'═'*50}")
        logs = log(f"🔄 تلاش {attempt} از {MAX_RETRIES}")
        logs = log(f"{'═'*50}")
        yield logs, code

        # ساخت Space
        logs = log("📦 [1/4] در حال ساخت Space...")
        yield logs, code
        success, result = create_space(space_name)
        if not success:
            logs = log(f"❌ خطا در ساخت Space: {result}")
            space_name = f"app-{int(time.time())}-{attempt}"
            yield logs, code
            continue

        space_id = f"{USER_NAME}/{space_name}"
        logs = log(f"✅ Space ساخته شد: {space_id}")
        yield logs, code

        # آپلود کد
        logs = log("⬆️ [2/4] در حال آپلود کد...")
        yield logs, code
        success, err = upload_to_space(code, space_id)
        if not success:
            logs = log(f"❌ خطا در آپلود: {err}")
            yield logs, code
            continue
        logs = log("✅ کد آپلود شد!")
        yield logs, code

        # انتظار برای build
        logs = log("⏳ [3/4] در حال build و اجرا... (این ممکنه چند دقیقه طول بکشه)")
        yield logs, code

        start = time.time()
        timeout = 240
        while time.time() - start < timeout:
            stage = get_space_status(space_id)

            if stage == "running":
                logs = log("✅ [4/4] Space اجرا شد!")
                logs = log(f"\n🎉 موفقیت‌آمیز!")
                logs = log(f"🌐 لینک: https://huggingface.co/spaces/{space_id}")
                yield logs, code
                return

            elif stage == "error":
                logs = log("❌ [4/4] Space خطا داشت!")
                space_logs = get_space_logs(space_id)
                logs = log(f"📋 لاگ خطا:\n{space_logs[-1500:]}")
                logs = log("🔧 در حال رفع خطا با Gemini 3.1 Pro...")
                yield logs, code

                try:
                    new_code = generate_code(user_request, code, space_logs[-3000:])
                    code = new_code
                    logs = log("✅ کد اصلاح شد! در حال تلاش مجدد...")
                    yield logs, code
                except Exception as e:
                    logs = log(f"❌ خطا در اصلاح کد: {e}")
                    yield logs, code

                time.sleep(3)
                break

            elif stage == "building":
                elapsed = int(time.time() - start)
                logs = log(f"🔨 در حال build... ({elapsed} ثانیه گذشته)")
                yield logs, code
                time.sleep(10)

            else:
                logs = log(f"⏳ وضعیت: {stage}...")
                yield logs, code
                time.sleep(5)
        else:
            logs = log("⚠️ زمان انتظار تموم شد، دوباره تلاش می‌کنیم...")
            yield logs, code

    logs = log("\n" + "═"*50)
    logs = log("❌ حداکثر تعداد تلاش رسیده شد!")
    logs = log("═"*50)
    yield logs, code


CSS = """
body, .gradio-container {
    background: #0a0a0f !important;
    color: #e8e8f0 !important;
}
.gradio-container {
    max-width: 1000px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}
button.primary {
    background: linear-gradient(135deg, #7c6cfc, #fc6c8f) !important;
    border: none !important;
    border-radius: 10px !important;
    color: white !important;
    font-weight: 700 !important;
    width: 100% !important;
}
textarea {
    background: #111118 !important;
    border: 1px solid #2a2a3a !important;
    color: #e8e8f0 !important;
    border-radius: 10px !important;
}
"""

with gr.Blocks(css=CSS, title="AI Code Generator - Gemini 3.1 Pro") as demo:
    gr.HTML("""
    <div style="text-align:center; padding:30px 0 20px;">
        <h1 style="font-size:2.5rem; background:linear-gradient(135deg,#7c6cfc,#fc6c8f);
        -webkit-background-clip:text; -webkit-text-fill-color:transparent; margin:0;">
        🤖 AI Code Generator</h1>
        <p style="color:#7070a0; margin-top:8px;">Powered by Gemini 3.1 Pro | ساخت و اجرای خودکار</p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            request_input = gr.Textbox(
                label="درخواست شما",
                placeholder="مثلاً: یه ماشین حساب ساده با ۴ عملیات",
                lines=3,
                value="Simple calculator with 4 operations"
            )
            run_btn = gr.Button("🚀 تولید و دیپلوی", variant="primary")

        with gr.Column(scale=1):
            log_output = gr.Textbox(
                label="📊 لاگ زنده",
                lines=18,
                interactive=False,
            )

    code_output = gr.Code(
        label="💻 کد تولید شده",
        language="python",
        interactive=False,
        lines=20,
    )

    run_btn.click(
        build_and_deploy,
        inputs=[request_input],
        outputs=[log_output, code_output],
    )

demo.launch(server_name="0.0.0.0", server_port=7860, share=False)