import io
import os
import tempfile
import uuid

import gradio as gr
from PIL import Image

from crew2 import (
    run_pipeline,
    generate_image,
    generate_caption,
    generate_voice,
    transcribe_audio,
)

# ---------------------------------------------------------------------------
# Custom CSS
# ---------------------------------------------------------------------------
_CSS = """
/* ── Global ─────────────────────────────────────────────── */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');

* { font-family: 'Inter', system-ui, -apple-system, sans-serif; }

.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }

/* ── Header ─────────────────────────────────────────────── */
.header-box {
    background: linear-gradient(135deg, #0f0c29, #302b63, #24243e);
    border-radius: 24px;
    padding: 2.5rem 2rem;
    margin-bottom: 2rem;
    text-align: center;
    border: 1px solid rgba(255,255,255,0.06);
    box-shadow: 0 20px 60px rgba(0,0,0,0.5);
}
.header-box h1 {
    font-size: 2.4rem;
    font-weight: 800;
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 50%, #4facfe 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 0.4rem;
    letter-spacing: -0.02em;
}
.header-box p {
    font-size: 1.05rem;
    color: rgba(255,255,255,0.55);
    font-weight: 300;
    margin: 0;
}

/* ── Card containers ────────────────────────────────────── */
.card {
    background: rgba(255,255,255,0.04);
    backdrop-filter: blur(16px);
    -webkit-backdrop-filter: blur(16px);
    border: 1px solid rgba(255,255,255,0.07);
    border-radius: 18px;
    padding: 1.8rem 1.5rem;
    margin-bottom: 1.5rem;
    box-shadow: 0 8px 32px rgba(0,0,0,0.25);
    transition: border-color 0.3s;
}
.card:hover { border-color: rgba(255,255,255,0.12); }

.card-label {
    font-size: 0.7rem;
    font-weight: 600;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    color: rgba(255,255,255,0.3);
    margin-bottom: 0.6rem;
}

/* ── Input overrides ────────────────────────────────────── */
input, textarea, .gr-box {
    border-radius: 12px !important;
    background: rgba(255,255,255,0.05) !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    color: #fff !important;
    transition: border-color 0.2s, box-shadow 0.2s;
}
input:focus, textarea:focus, .gr-box:focus-within {
    border-color: #4facfe !important;
    box-shadow: 0 0 0 3px rgba(79,172,254,0.15) !important;
}

/* ── Button ─────────────────────────────────────────────── */
.submit-btn {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%) !important;
    border: none !important;
    border-radius: 14px !important;
    color: #fff !important;
    font-weight: 700 !important;
    font-size: 1.05rem !important;
    padding: 0.9rem 2rem !important;
    transition: transform 0.2s, box-shadow 0.2s !important;
    box-shadow: 0 6px 24px rgba(245,87,108,0.3) !important;
    cursor: pointer;
}
.submit-btn:hover {
    transform: translateY(-2px) scale(1.01);
    box-shadow: 0 10px 36px rgba(245,87,108,0.4) !important;
}
.submit-btn:active { transform: translateY(0); }

/* ── Progress ───────────────────────────────────────────── */
.progress-text {
    font-size: 0.95rem;
    font-weight: 500;
    color: rgba(255,255,255,0.75);
    margin-top: 0.5rem;
    min-height: 1.6em;
}
.gr-progress-bar {
    --progress-bar-color: linear-gradient(90deg, #f093fb, #f5576c, #4facfe) !important;
    height: 6px !important;
    border-radius: 4px !important;
}

/* ── Output image / audio ───────────────────────────────── */
.output-image img {
    border-radius: 16px !important;
    box-shadow: 0 12px 48px rgba(0,0,0,0.4);
    width: 100%;
    height: auto;
    max-height: 500px;
    object-fit: cover;
}
audio { width: 100%; border-radius: 40px; }

/* ── Detail accordions ──────────────────────────────────── */
details.card-detail {
    background: rgba(255,255,255,0.03);
    border: 1px solid rgba(255,255,255,0.06);
    border-radius: 14px;
    padding: 0.8rem 1.2rem;
    margin-bottom: 0.8rem;
    transition: border-color 0.2s;
    cursor: pointer;
}
details.card-detail[open] { border-color: rgba(79,172,254,0.25); }
details.card-detail summary {
    font-weight: 600;
    font-size: 0.9rem;
    color: rgba(255,255,255,0.7);
    padding: 0.2rem 0;
    outline: none;
}
details.card-detail summary::-webkit-details-marker { color: rgba(255,255,255,0.3); }
details.card-detail .detail-body {
    margin-top: 0.8rem;
    font-size: 0.85rem;
    line-height: 1.6;
    color: rgba(255,255,255,0.55);
    max-height: 300px;
    overflow-y: auto;
    white-space: pre-wrap;
}

/* ── Status badges ──────────────────────────────────────── */
.badge {
    display: inline-block;
    font-size: 0.65rem;
    font-weight: 700;
    letter-spacing: 0.05em;
    padding: 0.2rem 0.6rem;
    border-radius: 20px;
    text-transform: uppercase;
}
.badge-success { background: rgba(52,211,153,0.15); color: #34d399; }
.badge-idle   { background: rgba(255,255,255,0.06); color: rgba(255,255,255,0.3); }

/* ── Responsive ─────────────────────────────────────────── */
@media (max-width: 640px) {
    .header-box h1 { font-size: 1.6rem; }
    .card { padding: 1.2rem 1rem; }
}
"""


# ---------------------------------------------------------------------------
# Pipeline runner — exposed as a module-level function for tests
# ---------------------------------------------------------------------------
def process(text_input: str | None, audio_input: str | None) -> tuple:
    """Run the full pipeline (text or audio). Returns (image, audio_path, caption, voice_style)."""
    session_id = str(uuid.uuid4())

    if audio_input is not None:
        text = transcribe_audio(audio_input, session_id=session_id)
    else:
        text = text_input

    if not text:
        raise ValueError("Please provide text or audio input.")

    results = run_pipeline(text, session_id=session_id)
    prompt = results.get("prompt")
    corroborate = results.get("result_corroborate", "")
    opposite = results.get("result_opposite", "")

    if not prompt:
        raise RuntimeError("Failed to generate an image prompt.")

    image_bytes = generate_image(prompt)
    image = Image.open(io.BytesIO(image_bytes))

    caption_result = generate_caption(corroborate, opposite, prompt, session_id=session_id)
    caption = caption_result.get("caption") or "(not found)"
    voice_style = caption_result.get("voice_style") or "(not found)"

    audio_bytes = generate_voice(voice_style, caption, session_id=session_id)
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp.write(audio_bytes)
    audio_path = tmp.name
    tmp.close()

    return image, audio_path, caption, voice_style


# ---------------------------------------------------------------------------
# Build demo
# ---------------------------------------------------------------------------
with gr.Blocks(
    title="CrewAI Research → Image → Voice",
    fill_height=True,
) as demo:

    # ---------- header ----------
    gr.HTML(
        """
<div class="header-box">
    <h1>🧠 CrewAI Research Pipeline</h1>
    <p>Enter a statement → AI researches both sides → generates a Flux LoRA image &amp; VoxCPM voiceover</p>
</div>"""
    )

    # ---------- input row ----------
    with gr.Column(elem_classes="card"):
        gr.HTML('<div class="card-label">Input</div>')
        with gr.Row(equal_height=True):
            with gr.Column(scale=3):
                text_input = gr.Textbox(
                    label="Statement",
                    placeholder='e.g. "AI will replace humans at work"',
                    lines=2,
                    elem_classes="input-field",
                )
            with gr.Column(scale=2):
                audio_input = gr.Audio(
                    label="Or speak your statement",
                    type="filepath",
                    sources=["microphone", "upload"],
                    waveform_options=gr.WaveformOptions(
                        waveform_color="rgba(79,172,254,0.3)",
                        waveform_progress_color="rgba(79,172,254,0.8)",
                    ),
                )
        with gr.Row():
            submit_btn = gr.Button(
                "🚀 Run Pipeline",
                variant="primary",
                size="lg",
                elem_classes="submit-btn",
            )

    # ---------- progress ----------
    progress_text = gr.HTML(
        '<div class="progress-text">⏳ Ready — enter a statement and press Run</div>'
    )

    # ---------- results ----------
    with gr.Row():
        with gr.Column(scale=3):
            # Image
            with gr.Column(elem_classes="card"):
                gr.HTML('<div class="card-label">Generated Image — Flux LoRA</div>')
                image_output = gr.Image(
                    show_label=False,
                    height=420,
                    elem_classes="output-image",
                )

            # Audio
            with gr.Column(elem_classes="card"):
                gr.HTML('<div class="card-label">Voiceover — VoxCPM</div>')
                audio_output = gr.Audio(
                    show_label=False,
                    type="filepath",
                    elem_classes="output-audio",
                )

        with gr.Column(scale=2):
            # Caption + style
            with gr.Column(elem_classes="card"):
                gr.HTML('<div class="card-label">Narration</div>')
                caption_output = gr.Textbox(
                    label="Caption (30 s)",
                    lines=4,
                    max_lines=6,
                    placeholder="The spoken narration will appear here ...",
                )
                style_output = gr.Textbox(
                    label="Voice Style",
                    lines=1,
                    placeholder="e.g. warm thoughtful professor",
                )

            # Research details
            with gr.Column(elem_classes="card"):
                gr.HTML('<div class="card-label">📚 Research Details</div>')

                with gr.Row():
                    corroborate_badge = gr.HTML(
                        '<span class="badge badge-idle">⏳ awaiting run</span>'
                    )
                    opposite_badge = gr.HTML(
                        '<span class="badge badge-idle">⏳ awaiting run</span>'
                    )

                corroborate_html = gr.HTML(
                    '<details class="card-detail"><summary>📖 Corroborative Evidence</summary>'
                    '<div class="detail-body">Run the pipeline to see findings …</div></details>'
                )
                opposite_html = gr.HTML(
                    '<details class="card-detail"><summary>📖 Opposing / Complementary</summary>'
                    '<div class="detail-body">Run the pipeline to see findings …</div></details>'
                )
                prompt_html = gr.HTML(
                    '<details class="card-detail"><summary>🖼️ Image Prompt</summary>'
                    '<div class="detail-body">Run the pipeline to see the prompt …</div></details>'
                )

    # ---------- event wiring ----------
    def on_submit(text, audio, progress: gr.Progress = gr.Progress(), request: gr.Request = None):
        session_id = str(uuid.uuid4())
        try:
            yield from _generate(text, audio, session_id, progress)
        except gr.Error:
            raise
        except Exception as exc:
            raise gr.Error(f"Pipeline failed: {exc}")

    def _generate(text, audio, session_id, progress):
        # ---- run pipeline ----
        if audio is not None:
            progress(0.02, desc="🎤 Transcribing audio …")
            text = transcribe_audio(audio, session_id=session_id)
        else:
            text = text

        if not text:
            raise gr.Error("Please enter a statement or record audio.")

        progress(0.05, desc="📡 Crew 1 — Corroborative research …")
        progress(0.18, desc="🔄 Crew 2 — Opposite research …")
        progress(0.32, desc="🧠 Crew 3 — Synthesizing image prompt …")

        results = run_pipeline(text, session_id=session_id)
        prompt = results["prompt"]
        corroborate = results["result_corroborate"]
        opposite = results["result_opposite"]

        if not prompt:
            raise gr.Error("Failed to generate an image prompt.")

        # ---- partial update: research badges + details ----
        yield (
            None, None, None, None,  # image, audio, caption, style
            _badge("done", "corroborated"),
            _badge("done", "opposed"),
            _detail("📖 Corroborative Evidence", _shorten(corroborate)),
            _detail("📖 Opposing / Complementary", _shorten(opposite)),
            _detail("🖼️ Image Prompt", prompt),
            _progress("🎨 Generating image …"),
        )

        # ---- image ----
        progress(0.50, desc="🎨 Generating image with Flux LoRA …")
        try:
            image_bytes = generate_image(prompt)
            image = Image.open(io.BytesIO(image_bytes))
        except Exception as exc:
            raise gr.Error(f"Image generation failed: {exc}")

        yield (
            image, None, None, None,
            _badge("done", "corroborated"),
            _badge("done", "opposed"),
            _detail("📖 Corroborative Evidence", _shorten(corroborate)),
            _detail("📖 Opposing / Complementary", _shorten(opposite)),
            _detail("🖼️ Image Prompt", prompt),
            _progress("📝 Generating narration …"),
        )

        # ---- caption ----
        progress(0.70, desc="📝 Generating narration via Gemma 4 …")
        try:
            caption_result = generate_caption(corroborate, opposite, prompt, session_id=session_id)
        except Exception as exc:
            raise gr.Error(f"Caption generation failed: {exc}")

        caption = caption_result.get("caption") or "(not found)"
        voice_style = caption_result.get("voice_style") or "(not found)"

        yield (
            image, None, caption, voice_style,
            _badge("done", "corroborated"),
            _badge("done", "opposed"),
            _detail("📖 Corroborative Evidence", _shorten(corroborate)),
            _detail("📖 Opposing / Complementary", _shorten(opposite)),
            _detail("🖼️ Image Prompt", prompt),
            _progress("🔊 Synthesizing voiceover …"),
        )

        # ---- voice ----
        progress(0.85, desc="🔊 Synthesizing voiceover with VoxCPM …")
        try:
            audio_bytes = generate_voice(voice_style, caption, session_id=session_id)
            tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
            tmp.write(audio_bytes)
            audio_path = tmp.name
            tmp.close()
        except Exception as exc:
            raise gr.Error(f"Voice synthesis failed: {exc}")

        yield (
            image, audio_path, caption, voice_style,
            _badge("done", "corroborated"),
            _badge("done", "opposed"),
            _detail("📖 Corroborative Evidence", _shorten(corroborate)),
            _detail("📖 Opposing / Complementary", _shorten(opposite)),
            _detail("🖼️ Image Prompt", prompt),
            _progress("✅ Complete!"),
        )

    # helpers -------------------------------------------------
    def _shorten(text: str, max_chars: int = 2000) -> str:
        """Truncate long text for the accordion detail view."""
        return text if len(text) <= max_chars else text[:max_chars] + "\n\n… (truncated)"

    def _badge(state: str, label: str) -> str:
        cls = "badge-success" if state == "done" else "badge-idle"
        icon = "✅" if state == "done" else "⏳"
        return f'<span class="badge {cls}">{icon} {label}</span>'

    def _detail(summary: str, body: str) -> str:
        safe = body.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        return (
            f"<details class='card-detail' open>"
            f"<summary>{summary}</summary>"
            f"<div class='detail-body'>{safe}</div>"
            f"</details>"
        )

    def _progress(msg: str) -> str:
        return f'<div class="progress-text">{msg}</div>'

    # ---------- wire submit ----------
    submit_event = submit_btn.click(
        fn=on_submit,
        inputs=[text_input, audio_input],
        outputs=[
            image_output,
            audio_output,
            caption_output,
            style_output,
            corroborate_badge,
            opposite_badge,
            corroborate_html,
            opposite_html,
            prompt_html,
            progress_text,
        ],
        queue=True,
    )

    # ---------- examples ----------
    gr.Markdown(
        """
---
<div style="text-align:center;color:rgba(255,255,255,0.35);font-size:0.85rem;padding:1rem 0;">
  Powered by
  <strong>CrewAI</strong> ·
  <strong>Gemma 4 26B</strong> (vLLM / H200) ·
  <strong>Flux LoRA</strong> (A100) ·
  <strong>VoxCPM</strong> (T4) ·
  <strong>Cohere Transcribe</strong> (T4)
</div>"""
    )

# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import os
    is_hf_space = os.environ.get("SPACE_ID") is not None
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=not is_hf_space,
        theme=gr.themes.Soft(
            primary_hue="violet",
            secondary_hue="blue",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter"),
        ),
        css=_CSS,
    )