"""
AxionStudio — AI Studio by AxionLab Co.
Chat · Image · Video generation powered by HF Inference Providers.
"""

import os
import tempfile
import gradio as gr
from huggingface_hub import InferenceClient

# ─── Auth ───────────────────────────────────────────────────────────────────
AXION_PASSWORD = os.environ.get("AXION_PASSWORD", "")

def check_login(username, password):
    """Aceita qualquer username, só verifica a senha."""
    return password == AXION_PASSWORD

# ─── Clients ────────────────────────────────────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN", "")

chat_client = InferenceClient(provider="auto", api_key=HF_TOKEN)
image_client = InferenceClient(provider="fal-ai", api_key=HF_TOKEN)
video_client = InferenceClient(provider="fal-ai", api_key=HF_TOKEN)

# ─── Model Choices ──────────────────────────────────────────────────────────
CHAT_MODELS = {
    "Qwen3 235B (Melhor)": "Qwen/Qwen3-235B-A22B",
    "DeepSeek R1 (Raciocínio)": "deepseek-ai/DeepSeek-R1-0528",
    "Llama 4 Maverick": "meta-llama/Llama-4-Maverick-17B-128E-Instruct",
}

IMAGE_MODELS = {
    "FLUX.1 Schnell (Rápido)": "black-forest-labs/FLUX.1-schnell",
    "FLUX.1 Dev (Qualidade)": "black-forest-labs/FLUX.1-dev",
}

VIDEO_MODELS = {
    "Wan 2.1 14B (Qualidade)": "Wan-AI/Wan2.1-T2V-14B",
    "LTX Video (Rápido)": "Lightricks/LTX-Video-0.9.8-13B-distilled",
}

SYSTEM_PROMPT = """Você é o AxionStudio, um assistente AI avançado criado pela AxionLab Co. 
Você é especialista em programação, raciocínio lógico, matemática, ciência e conversação geral.
Responda de forma clara, precisa e útil. Use markdown quando apropriado.
Quando o usuário pedir código, forneça código completo e funcional com explicações.
Responda no idioma em que o usuário falar."""

# ─── Chat Logic ─────────────────────────────────────────────────────────────
def chat_respond(message, history, model_name, temperature, max_tokens, system_prompt):
    model_id = CHAT_MODELS.get(model_name, "Qwen/Qwen3-235B-A22B")
    
    messages = [{"role": "system", "content": system_prompt or SYSTEM_PROMPT}]
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": message})

    partial = ""
    try:
        stream = chat_client.chat_completion(
            model=model_id,
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            stream=True,
        )
        for chunk in stream:
            token = chunk.choices[0].delta.content or ""
            partial += token
            yield partial
    except Exception as e:
        yield f"❌ Erro: {str(e)}\n\nTenta outro modelo ou tenta novamente."

# ─── Image Logic ────────────────────────────────────────────────────────────
def generate_image(prompt, model_name, width, height, steps, guidance, seed):
    if not prompt or not prompt.strip():
        gr.Warning("Escreve um prompt primeiro!")
        return None
    
    model_id = IMAGE_MODELS.get(model_name, "black-forest-labs/FLUX.1-schnell")
    
    try:
        kwargs = {}
        if seed and int(seed) >= 0:
            kwargs["seed"] = int(seed)
            
        image = image_client.text_to_image(
            prompt,
            model=model_id,
            width=int(width),
            height=int(height),
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            **kwargs,
        )
        return image
    except Exception as e:
        gr.Warning(f"Erro na geração: {str(e)}")
        return None

# ─── Video Logic ────────────────────────────────────────────────────────────
def generate_video(prompt, model_name, steps):
    if not prompt or not prompt.strip():
        gr.Warning("Escreve um prompt primeiro!")
        return None

    model_id = VIDEO_MODELS.get(model_name, "Wan-AI/Wan2.1-T2V-14B")
    
    try:
        video_bytes = video_client.text_to_video(
            prompt,
            model=model_id,
            num_inference_steps=int(steps),
        )
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        tmp.write(video_bytes)
        tmp.close()
        return tmp.name
    except Exception as e:
        gr.Warning(f"Erro na geração: {str(e)}")
        return None

# ─── Custom CSS ─────────────────────────────────────────────────────────────
CSS = """
/* Global */
.gradio-container {
    max-width: 1400px !important;
    margin: 0 auto !important;
}

/* Hide footer */
footer { display: none !important; }
.built-with { display: none !important; }

/* Header */
.axion-header {
    text-align: center;
    padding: 20px 0 10px 0;
}
.axion-header h1 {
    font-size: 2.2rem;
    background: linear-gradient(135deg, #6366f1, #a855f7, #06b6d4);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-weight: 800;
    margin: 0;
    letter-spacing: -0.5px;
}
.axion-header p {
    color: #94a3b8;
    font-size: 0.95rem;
    margin-top: 4px;
}

/* Tabs styling */
.tab-nav {
    border-bottom: 2px solid rgba(99, 102, 241, 0.3) !important;
}
.tab-nav button {
    font-size: 1rem !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
    transition: all 0.2s ease !important;
}
.tab-nav button.selected {
    color: #6366f1 !important;
    border-bottom-color: #6366f1 !important;
}

/* Chat area */
#chatbot {
    border-radius: 16px !important;
    min-height: 500px !important;
}

/* Image/Video output cards */
.output-card {
    border-radius: 16px !important;
    overflow: hidden;
}

/* Generate buttons */
.generate-btn {
    background: linear-gradient(135deg, #6366f1, #a855f7) !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.05rem !important;
    padding: 12px 32px !important;
    border-radius: 12px !important;
    border: none !important;
    transition: all 0.3s ease !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
}
.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(99, 102, 241, 0.4) !important;
}

/* Settings panel */
.settings-panel {
    background: rgba(99, 102, 241, 0.04);
    border: 1px solid rgba(99, 102, 241, 0.12);
    border-radius: 14px;
    padding: 16px;
}

/* Prompt textbox */
.prompt-box textarea {
    border-radius: 12px !important;
    font-size: 0.95rem !important;
}

/* Image gallery feel */
.image-output img {
    border-radius: 12px !important;
}

/* Info badges */
.info-badge {
    display: inline-block;
    background: rgba(99, 102, 241, 0.1);
    color: #6366f1;
    padding: 4px 12px;
    border-radius: 20px;
    font-size: 0.8rem;
    font-weight: 600;
    margin: 2px;
}

/* Responsive */
@media (max-width: 768px) {
    .axion-header h1 { font-size: 1.6rem; }
    .tab-nav button { padding: 8px 14px !important; font-size: 0.9rem !important; }
}
"""

# ─── Build App ──────────────────────────────────────────────────────────────
with gr.Blocks(fill_height=True, title="AxionStudio") as demo:
    # ── Header
    gr.HTML("""
    <div class="axion-header">
        <h1>⚡ AxionStudio</h1>
        <p>AI Studio by AxionLab Co. — Chat · Imagens · Vídeos</p>
    </div>
    """)

    with gr.Tabs() as tabs:
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # TAB 1: CHAT
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        with gr.Tab("💬 Chat", id="chat"):
            gr.ChatInterface(
                fn=chat_respond,
                save_history=True,
                editable=True,
                fill_height=True,
                chatbot=gr.Chatbot(
                    elem_id="chatbot",
                    placeholder="""<div style="display:flex;flex-direction:column;align-items:center;justify-content:center;height:100%;padding:40px">
                        <h2 style="margin:0;color:#6366f1;font-weight:700">👋 Olá!</h2>
                        <p style="color:#94a3b8;text-align:center;max-width:400px;margin-top:8px">
                            Sou o AxionStudio. Posso ajudar com programação, raciocínio, matemática, 
                            escrita criativa e muito mais. Pergunta-me qualquer coisa!
                        </p>
                    </div>""",
                    render_markdown=True,
                ),
                textbox=gr.Textbox(
                    placeholder="Escreve a tua mensagem...",
                    show_label=False,
                    container=False,
                    scale=7,
                    autofocus=True,
                    elem_classes=["prompt-box"],
                ),
                additional_inputs=[
                    gr.Dropdown(
                        choices=list(CHAT_MODELS.keys()),
                        value="Qwen3 235B (Melhor)",
                        label="🤖 Modelo",
                    ),
                    gr.Slider(
                        minimum=0.0, maximum=2.0, value=0.7, step=0.05,
                        label="🌡️ Temperatura",
                        info="Menor = mais preciso, Maior = mais criativo",
                    ),
                    gr.Slider(
                        minimum=256, maximum=16384, value=4096, step=256,
                        label="📏 Tokens Máximos",
                    ),
                    gr.Textbox(
                        value=SYSTEM_PROMPT,
                        label="📋 System Prompt",
                        lines=4,
                        info="Define a personalidade do assistente",
                    ),
                ],
                additional_inputs_accordion=gr.Accordion("⚙️ Configurações do Chat", open=False),
                examples=[
                    ["Explica-me como funcionam as redes neurais de forma simples"],
                    ["Escreve um script Python que faz web scraping com BeautifulSoup"],
                    ["Resolve: Se x² + 5x + 6 = 0, qual o valor de x?"],
                    ["Cria uma API REST completa em FastAPI com autenticação JWT"],
                    ["Escreve um poema sobre inteligência artificial"],
                ],
                example_labels=[
                    "🧠 Redes Neurais",
                    "🐍 Web Scraping",
                    "📐 Matemática",
                    "🚀 FastAPI + JWT",
                    "✍️ Poema sobre IA",
                ],
            )

        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # TAB 2: IMAGE GENERATION
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        with gr.Tab("🎨 Imagens", id="images"):
            gr.HTML("""
            <div style="text-align:center;padding:10px 0">
                <p style="color:#94a3b8;font-size:0.9rem">
                    Gera imagens com os modelos FLUX — descreve o que imaginas!
                </p>
            </div>
            """)

            with gr.Row(equal_height=True):
                with gr.Column(scale=2, min_width=320):
                    img_prompt = gr.Textbox(
                        label="✨ Prompt",
                        placeholder="Ex: A futuristic city floating in the clouds, cinematic lighting, 8k, photorealistic",
                        lines=4,
                        elem_classes=["prompt-box"],
                    )

                    with gr.Accordion("⚙️ Configurações", open=True):
                        img_model = gr.Dropdown(
                            choices=list(IMAGE_MODELS.keys()),
                            value="FLUX.1 Schnell (Rápido)",
                            label="🎨 Modelo",
                        )
                        with gr.Row():
                            img_width = gr.Slider(256, 1536, value=1024, step=64, label="Largura")
                            img_height = gr.Slider(256, 1536, value=1024, step=64, label="Altura")
                        with gr.Row():
                            img_steps = gr.Slider(1, 50, value=4, step=1, label="Steps", info="Schnell: 4 | Dev: 20-30")
                            img_guidance = gr.Slider(1.0, 20.0, value=3.5, step=0.5, label="Guidance Scale")
                        img_seed = gr.Number(value=-1, label="🎲 Seed (-1 = aleatório)", precision=0)

                    img_btn = gr.Button("✨ Gerar Imagem", variant="primary", size="lg", elem_classes=["generate-btn"])

                    gr.HTML("<p style='color:#94a3b8;font-size:0.85rem;margin-top:12px'>💡 Ideias rápidas:</p>")
                    with gr.Row():
                        for prompt_text, label in [
                            ("A majestic dragon made of crystal, perched on a mountain peak, sunset, fantasy art, highly detailed", "🐉 Dragão"),
                            ("Portrait of a cyberpunk samurai, neon city background, rain, cinematic lighting, 8k", "🤖 Samurai"),
                            ("A cozy coffee shop interior, warm lighting, plants, watercolor painting style", "☕ Café"),
                        ]:
                            gr.Button(label, size="sm").click(fn=lambda p=prompt_text: p, outputs=img_prompt)

                with gr.Column(scale=3, min_width=400):
                    img_output = gr.Image(label="Imagem Gerada", type="pil", elem_classes=["image-output", "output-card"], height=600)

            img_btn.click(fn=generate_image, inputs=[img_prompt, img_model, img_width, img_height, img_steps, img_guidance, img_seed], outputs=img_output, show_progress="full")

        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # TAB 3: VIDEO GENERATION
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        with gr.Tab("🎬 Vídeos", id="videos"):
            gr.HTML("""
            <div style="text-align:center;padding:10px 0">
                <p style="color:#94a3b8;font-size:0.9rem">
                    Gera vídeos a partir de texto — pode demorar 1-3 minutos.
                </p>
            </div>
            """)

            with gr.Row(equal_height=True):
                with gr.Column(scale=2, min_width=320):
                    vid_prompt = gr.Textbox(
                        label="🎬 Prompt",
                        placeholder="Ex: A cinematic shot of a rocket launching at night, slow motion, dramatic clouds, 4K",
                        lines=4,
                        elem_classes=["prompt-box"],
                    )

                    with gr.Accordion("⚙️ Configurações", open=True):
                        vid_model = gr.Dropdown(choices=list(VIDEO_MODELS.keys()), value="Wan 2.1 14B (Qualidade)", label="🎬 Modelo")
                        vid_steps = gr.Slider(10, 50, value=30, step=5, label="Steps", info="Mais steps = melhor qualidade, mais lento")

                    vid_btn = gr.Button("🎬 Gerar Vídeo", variant="primary", size="lg", elem_classes=["generate-btn"])

                    gr.HTML("""
                    <div style="margin-top:16px;padding:12px;background:rgba(245,158,11,0.08);border:1px solid rgba(245,158,11,0.2);border-radius:10px">
                        <p style="color:#f59e0b;font-size:0.85rem;margin:0">
                            ⏱️ A geração de vídeo pode demorar <b>1-3 minutos</b>. Sê paciente!
                        </p>
                    </div>
                    """)

                    gr.HTML("<p style='color:#94a3b8;font-size:0.85rem;margin-top:12px'>💡 Ideias rápidas:</p>")
                    with gr.Row():
                        for prompt_text, label in [
                            ("Ocean waves crashing on a rocky shore at golden hour, cinematic, slow motion, 4K", "🌊 Oceano"),
                            ("A cat sitting on a windowsill watching rain fall outside, cozy atmosphere, warm lighting", "🐱 Gato"),
                            ("Northern lights dancing over a snowy mountain landscape, timelapse, 4K", "🌌 Aurora"),
                        ]:
                            gr.Button(label, size="sm").click(fn=lambda p=prompt_text: p, outputs=vid_prompt)

                with gr.Column(scale=3, min_width=400):
                    vid_output = gr.Video(label="Vídeo Gerado", autoplay=True, elem_classes=["output-card"], height=500)

            vid_btn.click(fn=generate_video, inputs=[vid_prompt, vid_model, vid_steps], outputs=vid_output, show_progress="full")

        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        # TAB 4: ABOUT
        # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
        with gr.Tab("ℹ️ Sobre", id="about"):
            gr.HTML("""
            <div style="max-width:700px;margin:30px auto;padding:30px">
                <div style="text-align:center;margin-bottom:30px">
                    <h2 style="background:linear-gradient(135deg,#6366f1,#a855f7);-webkit-background-clip:text;-webkit-text-fill-color:transparent;font-size:2rem;font-weight:800">
                        ⚡ AxionStudio
                    </h2>
                    <p style="color:#94a3b8;font-size:1.1rem">AI Studio by AxionLab Co.</p>
                </div>
                
                <div style="display:grid;grid-template-columns:repeat(auto-fit,minmax(200px,1fr));gap:16px;margin:24px 0">
                    <div style="background:rgba(99,102,241,0.06);border:1px solid rgba(99,102,241,0.15);border-radius:14px;padding:20px;text-align:center">
                        <div style="font-size:2rem;margin-bottom:8px">💬</div>
                        <h3 style="color:#6366f1;margin:0 0 4px 0;font-size:1rem">Chat AI</h3>
                        <p style="color:#94a3b8;font-size:0.8rem;margin:0">Qwen3 235B, DeepSeek R1, Llama 4</p>
                    </div>
                    <div style="background:rgba(168,85,247,0.06);border:1px solid rgba(168,85,247,0.15);border-radius:14px;padding:20px;text-align:center">
                        <div style="font-size:2rem;margin-bottom:8px">🎨</div>
                        <h3 style="color:#a855f7;margin:0 0 4px 0;font-size:1rem">Imagens</h3>
                        <p style="color:#94a3b8;font-size:0.8rem;margin:0">FLUX.1 Schnell & Dev</p>
                    </div>
                    <div style="background:rgba(6,182,212,0.06);border:1px solid rgba(6,182,212,0.15);border-radius:14px;padding:20px;text-align:center">
                        <div style="font-size:2rem;margin-bottom:8px">🎬</div>
                        <h3 style="color:#06b6d4;margin:0 0 4px 0;font-size:1rem">Vídeos</h3>
                        <p style="color:#94a3b8;font-size:0.8rem;margin:0">Wan 2.1 & LTX Video</p>
                    </div>
                </div>
                
                <div style="background:rgba(255,255,255,0.03);border:1px solid rgba(255,255,255,0.08);border-radius:14px;padding:20px;margin-top:20px">
                    <h3 style="color:#e2e8f0;margin:0 0 12px 0;font-size:1rem">🚀 Modelos Disponíveis</h3>
                    <table style="width:100%;border-collapse:collapse;font-size:0.85rem">
                        <tr style="border-bottom:1px solid rgba(255,255,255,0.06)">
                            <td style="padding:8px;color:#94a3b8;font-weight:600">Chat</td>
                            <td style="padding:8px;color:#e2e8f0">Qwen3-235B-A22B · DeepSeek-R1 · Llama 4 Maverick</td>
                        </tr>
                        <tr style="border-bottom:1px solid rgba(255,255,255,0.06)">
                            <td style="padding:8px;color:#94a3b8;font-weight:600">Imagem</td>
                            <td style="padding:8px;color:#e2e8f0">FLUX.1-schnell · FLUX.1-dev</td>
                        </tr>
                        <tr>
                            <td style="padding:8px;color:#94a3b8;font-weight:600">Vídeo</td>
                            <td style="padding:8px;color:#e2e8f0">Wan2.1-T2V-14B · LTX-Video-13B</td>
                        </tr>
                    </table>
                </div>
                
                <div style="text-align:center;margin-top:30px">
                    <p style="color:#64748b;font-size:0.8rem">
                        Feito com ❤️ pela <a href="https://huggingface.co/AxionLab-Co" target="_blank" style="color:#6366f1;text-decoration:none;font-weight:600">AxionLab Co.</a>
                        <br>Powered by Hugging Face Inference Providers
                    </p>
                </div>
            </div>
            """)

# ─── Theme ──────────────────────────────────────────────────────────────────
theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="purple",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
)

# ─── Launch (always runs — no __main__ guard for HF Spaces) ────────────────
demo.launch(
    theme=theme,
    css=CSS,
    auth=check_login,
    auth_message="🔒 Insere a senha para aceder ao AxionStudio (username pode ser qualquer coisa)",
    share=False,
)