import base64
import json
import logging

import gradio as gr

from src.config import settings
from src.gemini_image import GeminiImageClient
from src.gemini_llm import GeminiLLMClient
from src.image_generator import ImageGenerator
from src.llm import LLMClient
from src.stable_diffusion import DiffusionClient
from src.story_generator import StoryGenerator

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

_NUM = settings.STORYBOARD_NUM_SCENES

_EXAMPLES = [
    (
        "Một con thỏ đã thách thức một con rùa chạy đua. Trong cuộc đua, thỏ chủ quan "
        "và ngủ quên giữa đường còn rùa vẫn kiên trì đến cuối nên thỏ đã thua cuộc."
    ),
    (
        "Cậu bé chăn cừu vì buồn chán nên đã nhiều lần nói dối dân làng rằng có sói đến. "
        "Đến khi sói xuất hiện thật, không ai tin lời cậu nữa và bầy cừu bị ăn thịt hết."
    ),
    (
        "Hai con dê cùng đi qua một chiếc cầu hẹp từ hai phía ngược nhau. Không con nào "
        "chịu nhường đường cho con nào, cuối cùng cả hai đều rơi xuống suối."
    ),
]

_INFO_NOTE_HTML = """
<div class="info-card">
  <div class="info-card-title">⚠️ Lưu ý sử dụng</div>
  <ul class="info-card-list">
    <li>Đây là ứng dụng demo phục vụ mục đích học tập và minh hoạ.</li>
    <li>Nội dung truyện và hình ảnh do AI tạo sinh có thể không chính xác hoặc không phù hợp với mọi đối tượng.</li>
    <li>Thời gian tạo thường khoảng <b>3–5 phút</b>, tuỳ tải hệ thống và giới hạn token/API.</li>
  </ul>
</div>
"""

_USAGE_MARKDOWN = """
1. **Nhập tóm tắt truyện** bằng tiếng Việt, nêu rõ nhân vật và diễn biến chính.  
2. **Tuỳ chọn:** thêm **HF Token** cá nhân nếu gặp giới hạn tốc độ.  
3. Bấm **Tạo truyện tranh** và chờ hệ thống tạo 4 ô truyện.  
4. Mở phần **Storyboard & Prompts** nếu muốn xem storyboard JSON và prompt chi tiết.

**Mẹo:** Tóm tắt càng rõ ràng và có diễn biến cụ thể thì kết quả càng sát ý hơn.
"""

_llm_client: LLMClient | GeminiLLMClient | None = None
_story_gen: StoryGenerator | None = None
_img_gen: ImageGenerator | None = None
_active_token: str = ""
_active_backend: str = ""  # "gemini" | "hf"

_RATE_LIMIT_HINT = (
    "💡 Hãy thử lại sau hoặc dán HF Token / Gemini API Key cá nhân vào rồi bấm lại."
)


def _resolve_token(runtime_token: str = "") -> str:
    token = (runtime_token or "").strip()
    return token if token else settings.HF_TOKEN


def _resolve_gemini_key(runtime_key: str = "") -> str:
    key = (runtime_key or "").strip()
    return key if key else settings.GEMINI_API_KEY


def _is_limit_error(exc: Exception | str) -> bool:
    msg = str(exc).lower()
    keys = (
        "429",
        "402",
        "rate limit",
        "rate-limit",
        "quota",
        "payment required",
        "too many requests",
        "credits",
        "pre-paid",
        "limit exceeded",
    )
    return any(key in msg for key in keys)


def _friendly_api_status(exc: Exception | str, stage: str) -> str:
    detail = str(exc)[:300]
    if _is_limit_error(exc):
        return f"⚠️ API đang bị giới hạn ở bước {stage}.\n{_RATE_LIMIT_HINT}\n\n🔍 Chi tiết: {detail}"
    return f"⚠️ Lỗi ở bước **{stage}**.\n\n🔍 Chi tiết lỗi: `{detail}`\n\n{_RATE_LIMIT_HINT}"


def _init_pipeline(runtime_token: str = "", runtime_gemini_key: str = "") -> None:
    global _llm_client, _story_gen, _img_gen, _active_token, _active_backend

    gemini_key = _resolve_gemini_key(runtime_gemini_key)
    hf_token = _resolve_token(runtime_token)

    # Determine which backend to use
    backend = "gemini" if gemini_key else "hf"
    cache_key = gemini_key if backend == "gemini" else hf_token

    if _llm_client is not None and cache_key == _active_token and backend == _active_backend:
        return

    if _llm_client is not None:
        logger.info("Backend or token changed — reinitializing pipeline (backend=%s).", backend)

    if backend == "gemini":
        logger.info("Initializing Gemini pipeline (text=%s, image=%s)…",
                    settings.GEMINI_TEXT_MODEL, settings.GEMINI_IMAGE_MODEL)
        _llm_client = GeminiLLMClient(
            api_key=gemini_key,
            model_id=settings.GEMINI_TEXT_MODEL,
        )
        _llm_client.load_model()

        img_client = GeminiImageClient(
            api_key=gemini_key,
            model_id=settings.GEMINI_IMAGE_MODEL,
        )
        img_client.load_model()
    else:
        logger.info("Initializing HuggingFace pipeline…")
        _llm_client = LLMClient(hf_token=hf_token, cache_dir=settings.CACHE_DIR)
        _llm_client.load_model()

        img_client = DiffusionClient(
            model_id=settings.IMAGE_MODEL_ID,
            hf_token=hf_token,
            provider=settings.IMAGE_PROVIDER,
        )
        img_client.load_model()

    _story_gen = StoryGenerator(_llm_client)
    _img_gen = ImageGenerator(img_client)
    _active_token = cache_key
    _active_backend = backend
    logger.info("Pipeline ready (backend=%s).", backend)


_STORY_PLACEHOLDER = "*Chưa chạy pipeline — kết quả storyboard sẽ hiển thị ở đây.*"

def _out(panels: list, story_md: str, status: str):
    padded = list(panels) + [None] * (_NUM - len(panels))
    story = story_md or _STORY_PLACEHOLDER
    status_html = f'<div class="status-bar">{status}</div>'
    return (*padded[:_NUM], story, status_html)


def _img_b64(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def _build_story_md(
    scenes: list[dict[str, str]],
    image_prompts: list[str] | None = None,
) -> str:
    parts: list[str] = [
        "## 🧾 Storyboard JSON",
        "```json",
        json.dumps({"scenes": scenes}, ensure_ascii=False, indent=2),
        "```",
        "",
        "## 🎬 Parsed Scenes",
    ]

    for idx, scene in enumerate(scenes, start=1):
        parts.extend(
            [
                f"### Cảnh {idx}",
                f"- **Caption (VI):** {scene['caption_vi']}",
                f"- **Prompt (EN):** {scene['prompt_en']}",
                "",
            ]
        )

    return "\n".join(parts).strip()


def inference(user_story: str, runtime_token: str = "", runtime_gemini_key: str = ""):
    """Generate a 4-panel comic from a Vietnamese story summary."""
    if not user_story or not user_story.strip():
        yield _out([], "", "⚠️ Vui lòng nhập tóm tắt truyện.")
        return

    try:
        _init_pipeline(runtime_token, runtime_gemini_key)
    except Exception as exc:
        logger.exception("Pipeline initialization failed.")
        yield _out([], "", _friendly_api_status(exc, "khởi tạo pipeline"))
        return

    yield _out([], "", "⏳ Đang tạo storyboard…")

    try:
        scenes = _story_gen.generate_storyboard(
            user_story=user_story,
            num_scenes=_NUM,
            max_new_tokens=settings.LLM_MAX_TOKENS,
            temperature=settings.LLM_TEMPERATURE,
        )
    except ValueError as exc:
        logger.exception("Storyboard generation failed.")
        yield _out(
            [],
            f"## ⚠️ Lỗi storyboard\n```text\n{exc}\n```",
            "⚠️ LLM trả về dữ liệu không hợp lệ. Vui lòng thử lại với tóm tắt rõ ràng hơn.",
        )
        return
    except Exception as exc:
        logger.exception("Unexpected error during storyboard generation.")
        yield _out([], "", _friendly_api_status(exc, "tạo storyboard"))
        return

    image_prompts: list[str] = []
    story_md = _build_story_md(scenes)
    yield _out([], story_md, f"⏳ Đang tạo hình cho {_NUM} cảnh…")

    panels: list = []
    for idx, scene in enumerate(scenes, start=1):
        full_prompt = f"{scene['prompt_en']} Style: {settings.COMIC_STYLE}."
        image_prompts.append(full_prompt)
        story_md = _build_story_md(scenes, image_prompts=image_prompts)

        yield _out(panels, story_md, f"🎨 Đang vẽ cảnh {idx}/{_NUM}…")

        try:
            panel = _img_gen.generate_image(
                prompt=full_prompt,
                paragraph=scene["caption_vi"],
            )
        except Exception as exc:
            logger.exception("Image generation failed at panel %d.", idx)
            panels.append(None)
            yield _out(panels, story_md, _friendly_api_status(exc, f"vẽ cảnh {idx}/{_NUM}"))
            continue

        if panel is not None:
            panels.append(panel)
            yield _out(panels, story_md, f"✅ Xong cảnh {idx}/{_NUM}")
        else:
            panels.append(None)
            yield _out(
                panels,
                story_md,
                f"⚠️ Cảnh {idx}/{_NUM} thất bại. Có thể do rate-limit/quota API hoặc lỗi tạo ảnh. {_RATE_LIMIT_HINT}",
            )

    ok = sum(1 for p in panels if p is not None)
    yield _out(panels, story_md, f"🎉 Hoàn tất! {ok}/{_NUM} panel.")


_CSS = """
:root,#gradio-app{
  --app-max-width:1440px; --panel-gap:4px;
  --r-lg:16px; --r-md:12px; --r-sm:10px;
  --bg-page:#f7faff; --bg-page-2:#eef4fb; --bg-card:#ffffff; --bg-card-soft:#fdfefe;
  --bg-note:#fff7e6; --bg-note-border:#f2d38b; --bg-status:#edf4ff; --bg-preview:#ffffff; --bg-debug:#fcfdff;
  --text-main:#0d0d0d; --text-sub:#0d0d0d; --text-muted:#0d0d0d; --text-note:#5a4200;
  --border-soft:#dde7f2; --border-mid:#cad8e8;
  --accent:#4c76ff; --accent-2:#6c84ff; --accent-soft:#eaf1ff; --accent-soft-border:#cfe0ff;
}

#gradio-app{
  --background-fill-primary:var(--bg-page);
  --background-fill-secondary:var(--bg-page-2);
  --block-background-fill:var(--bg-card);
  --block-border-color:var(--border-soft);
  --block-border-width:1px;
}

html,body{
  margin:0; padding:0;
  background:linear-gradient(180deg,var(--bg-page) 0%,var(--bg-page-2) 100%)!important;
  color:var(--text-main);
}

#gradio-app,.gradio-container,.main,.contain,.wrap,.column,.row,.block,.html-container,.gr-group{
  background:transparent!important;
}

.gradio-container{
  max-width:var(--app-max-width)!important;
  margin:0 auto!important;
  padding:10px 14px 8px!important;
}

/* Header */
.header-row{
  gap:12px!important; padding:2px 0 6px!important; align-items:center!important;
}
.header-row img{
  height:112px; width:auto; object-fit:contain;
}
.header-meta{
  display:flex; flex-direction:column; justify-content:center; min-width:0;
}
.header-title{
  margin:0; font-size:1.7rem; line-height:1.15; font-weight:800;
  color:#0d0d0d; letter-spacing:-.01em;
}
.header-sub{
  margin:4px 0 0; font-size:.96rem; line-height:1.35;
  color:#0d0d0d; font-weight:700;
}

/* Info note */
.info-card{
  background:var(--bg-note)!important;
  border:1px solid var(--bg-note-border)!important;
  border-radius:var(--r-md);
  padding:10px 12px;
  color:var(--text-note)!important;
  margin:0 0 10px;
}
.info-card-title,.info-card-list,.info-card-list li{
  color:var(--text-note)!important;
}
.info-card-title{
  margin-bottom:6px; font-size:1rem; font-weight:700;
}
.info-card-list{
  margin:0; padding-left:18px; font-size:.92rem; line-height:1.55;
}
.info-card-list li+li{ margin-top:4px; }

/* Layout */
.main-layout{
  gap:10px!important; align-items:stretch!important;
}
.control-stack,.preview-col{
  gap:8px!important;
}

/* Cards */
.control-card,.preview-card{
  background:#e8f0fe!important;
  border:1px solid var(--border-soft)!important;
  border-radius:var(--r-lg)!important;
  box-shadow:0 4px 18px rgba(25,45,80,.04)!important;
  padding:10px!important;
}
.control-card,.control-card *,.preview-card,.preview-card *{
  color:#0d0d0d!important;
}
.control-card a,.preview-card a,.help-markdown a,.story-markdown a{
  color:var(--accent)!important;
}

.section-help,.token-help p,.help-markdown p,.help-markdown li{
  font-size:.92rem; line-height:1.58; color:#0d0d0d!important;
}
.section-help{
  margin:0; padding:10px 12px 6px!important;
}

/* Inputs */
textarea,input{
  background:#ffffff!important;
  color:#0d0d0d!important;
  font-weight:600!important;
  border:1px solid var(--border-mid)!important;
  box-shadow:none!important;
}
.control-card textarea{
  min-height:140px!important;
}
textarea:focus,input:focus{
  border-color:var(--accent)!important;
  box-shadow:0 0 0 3px rgba(76,118,255,.10)!important;
}
::placeholder{
  color:#5a6b7d!important; opacity:1!important;
}
label,.gr-label,.block-title{
  color:#0d0d0d!important; font-size:.95rem!important;
}
.gr-info{
  color:#0d0d0d!important; font-size:.9rem!important;
}
[data-testid="block-info"]{
  display:none!important;
}

/* Button */
.gen-btn{
  width:100%!important; height:46px!important;
  font-size:16px!important; font-weight:700!important;
  border-radius:12px!important;
  border:1px solid var(--accent)!important;
  background:linear-gradient(180deg,#5f88ff 0%,var(--accent) 100%)!important;
  color:#fff!important; cursor:pointer!important;
  box-shadow:0 8px 18px rgba(76,118,255,.16);
}
.gen-btn:hover{
  filter:brightness(.985);
}

/* Status */
.status-bar{
  background:var(--bg-status)!important;
  border:1px solid var(--accent-soft-border)!important;
  border-radius:var(--r-sm);
  padding:9px 11px!important;
  font-size:.92rem!important;
  line-height:1.5!important;
  color:#0d0d0d!important;
}
.status-bar p,.status-bar .prose p{
  color:#0d0d0d!important;
}
.control-stack .html-container:has(.status-bar){
  background:transparent!important; border:none!important; box-shadow:none!important; padding:0!important;
}

/* Accordions */
.gradio-accordion,
.gradio-accordion > div,
.gradio-accordion details{
  background:#e8f0fe!important;
  border:1px solid var(--border-soft)!important;
  border-radius:var(--r-md)!important;
  box-shadow:none!important;
  overflow:hidden!important;
}

.gradio-accordion summary,
.gradio-accordion button{
  background:#e8f0fe!important;
  color:#0d0d0d!important;
  font-weight:700!important;
  font-size:.98rem!important;
  box-shadow:none!important;
  outline:none!important;
  border:none!important;
}

.gradio-accordion summary:hover,
.gradio-accordion button:hover{
  background:#f8fbff!important;
}

.gradio-accordion svg,
.gradio-accordion path{
  stroke:#0d0d0d!important;
  fill:#0d0d0d!important;
}

/* Examples */
.examples-wrap{
  gap:5px!important;
}
.example-btn{
  font-size:.88rem!important;
  padding:8px 10px!important;
  border-radius:10px!important;
  white-space:normal!important;
  text-align:left!important;
  line-height:1.38!important;
  background:#e8f0fe!important;
  border:1px solid var(--border-soft)!important;
  color:#0d0d0d!important;
}
.example-btn:hover{
  background:#f8fbff!important;
  border-color:#d6e4ff!important;
}

/* Preview */
.preview-header{
  display:flex; align-items:center; justify-content:space-between;
  gap:8px; margin:0 0 6px;
}
.preview-title{
  margin:0; font-size:1.06rem; font-weight:700; color:#0d0d0d!important;
}
.preview-sub{
  margin:2px 0 0; font-size:.88rem; color:#0d0d0d!important;
}
.comic-grid,.panel-row{
  gap:var(--panel-gap)!important;
}
.panel-row{
  margin:0!important; padding:0!important;
}
.panel-img,.panel-img>div,.panel-img .wrap{
  margin:0!important; padding:0!important;
  overflow:hidden; border-radius:12px;
  background:var(--bg-preview)!important;
  border:1px solid var(--border-soft)!important;
}
.panel-img img{
  display:block; border-radius:12px; background:var(--bg-preview)!important;
}

/* Storyboard bottom block */
.story-panel,
.story-panel > div,
.story-panel details{
  background:#e8f0fe!important;
  border:1px solid var(--border-soft)!important;
  border-radius:var(--r-lg)!important;
  box-shadow:0 4px 18px rgba(25,45,80,.04)!important;
  overflow:hidden!important;
}

.story-panel summary,
.story-panel button{
  background:#e8f0fe!important;
  color:#0d0d0d!important;
  font-weight:700!important;
  font-size:.98rem!important;
  border:none!important;
  box-shadow:none!important;
  outline:none!important;
}

.story-panel summary:hover,
.story-panel button:hover{
  background:#f8fbff!important;
}

.story-panel .label-wrap,
.story-panel .icon-wrap{
  background:transparent!important;
}

.story-hint{
  margin:0 0 4px;
  padding:8px 12px 0;
  font-size:.9rem;
  color:#0d0d0d!important;
}

.story-markdown{
  background:#e8f0fe!important;
  border-top:1px solid var(--border-soft)!important;
  border-radius:0 0 14px 14px!important;
  padding:12px 14px!important;
  min-height:2em;
}

.story-markdown h2,.story-markdown h3{
  margin:.95rem 0 .45rem!important;
  line-height:1.28!important;
  font-size:1.08rem!important;
  color:#0d0d0d!important;
}
.story-markdown h2:first-child{
  margin-top:0!important;
}
.story-markdown p,.story-markdown li{
  font-size:.94rem!important;
  line-height:1.64!important;
  color:#0d0d0d!important;
}
.story-markdown ul,.story-markdown ol{
  padding-left:1.1rem!important;
}
.story-markdown pre{
  margin:.55rem 0 .85rem!important;
  padding:12px 14px!important;
  overflow-x:auto!important;
  border-radius:12px!important;
  background:#1f2937!important;
  color:#f3f4f6!important;
  border:1px solid rgba(255,255,255,.08)!important;
}
.story-markdown pre,.story-markdown pre *,.story-markdown pre code{
  color:#f3f4f6!important;
}
.story-markdown code{
  font-size:.88rem!important;
}
.story-markdown :not(pre) > code{
  color:#0d0d0d!important;
  background:#eef2f7!important;
  padding:.15em .4em!important;
  border-radius:4px!important;
}
.story-markdown strong{
  color:#0d0d0d!important;
}

/* Footer */
.footer-text{
  text-align:center;
  font-size:.94rem;
  color:#0d0d0d!important;
  padding:10px 0 6px;
  margin-top:10px;
  border-top:1px solid var(--border-mid);
}
.footer-text a{
  color:var(--accent)!important;
  text-decoration:none;
  font-weight:600;
}
.footer-text a:hover{
  text-decoration:underline;
}

@media (max-width:980px){
  .gradio-container{ padding:6px!important; }
  .header-row{ gap:8px!important; padding:2px 0!important; }
  .header-row img{ height:60px; }
  .header-title{ font-size:1.4rem; }
  .header-sub{ font-size:.9rem; }
}
"""

_theme = gr.themes.Base().set(
    background_fill_primary="#f7faff",
    background_fill_secondary="#eef4fb",
    block_background_fill="#e8f0fe",
    block_border_color="#dde7f2",
    block_border_width="1px",
    input_background_fill="#ffffff",
)
with gr.Blocks(title="AI Comic Generation", fill_width=True, fill_height=True, theme=_theme) as demo:
    with gr.Row(elem_classes="header-row"):
        logo_b64 = _img_b64("static/aivn_logo.png")
        gr.HTML(f'<img src="data:image/png;base64,{logo_b64}" alt="AIVN">')
        gr.HTML(
            '<div class="header-meta">'
            '<p class="header-title">🎨 AI Comic Generation</p>'
            '<p class="header-sub">AIO2025 — LLM + FLUX.1-schnell API</p>'
            '</div>'
        )

    gr.HTML(_INFO_NOTE_HTML)

    with gr.Row(equal_height=False, elem_classes="main-layout"):
        with gr.Column(scale=4, min_width=320, elem_classes="control-stack"):
            with gr.Group(elem_classes="control-card"):
                gr.Markdown("### 💡 Tóm tắt truyện")
                prompt_input = gr.Textbox(
                    label="",
                    info="Nhập tóm tắt bằng tiếng Việt, nêu rõ nhân vật và các sự kiện chính.",
                    placeholder="Ví dụ: Một con thỏ lười biếng học được bài học về sự chăm chỉ từ một con rùa…",
                    lines=5,
                    max_lines=8,
                )
                gr.Markdown(
                    '<p class="section-help">Tóm tắt càng rõ ràng và có diễn biến cụ thể thì kết quả càng sát ý hơn.</p>'
                )

            with gr.Accordion("📌 Ví dụ nhanh", open=False):
                with gr.Column(elem_classes="examples-wrap"):
                    for example in _EXAMPLES:
                        gr.Button(
                            example[:92] + "…" if len(example) > 92 else example,
                            size="sm",
                            elem_classes="example-btn",
                        ).click(fn=lambda text=example: text, inputs=None, outputs=prompt_input)

            with gr.Accordion("🔑 API Keys (tuỳ chọn)", open=False):
                gr.Markdown(
                    "**Gemini API Key** *(ưu tiên)*: Nếu bạn cung cấp key này, pipeline sẽ dùng "
                    "Google Gemini để sinh text và Imagen để vẽ hình. "
                    "Vào **Google AI Studio → API Keys** để lấy key.\n\n"
                    "**HF Token** *(dự phòng)*: Dùng khi không có Gemini key và bị rate-limit trên "
                    "HuggingFace. Vào **HuggingFace → Settings → Access Tokens** để lấy token. "
                    "Cả hai key **chỉ dùng trong phiên hiện tại và không được lưu trữ**.",
                    elem_classes="token-help",
                )
                gemini_key_input = gr.Textbox(
                    label="Gemini API Key",
                    info="Ưu tiên — dùng Gemini + Imagen thay vì HuggingFace.",
                    type="password",
                    placeholder="AIza...",
                    lines=1,
                    max_lines=1,
                )
                token_input = gr.Textbox(
                    label="Hugging Face Token",
                    info="Dự phòng — chỉ dùng khi không có Gemini key và gặp lỗi rate-limit.",
                    type="password",
                    placeholder="hf_...",
                    lines=1,
                    max_lines=1,
                )

            generate_btn = gr.Button("Tạo truyện tranh 🎨", elem_classes="gen-btn")
            status = gr.HTML('<div class="status-bar">Sẵn sàng.</div>')

            with gr.Accordion("❓ Hướng dẫn sử dụng", open=False):
                gr.Markdown(_USAGE_MARKDOWN, elem_classes="help-markdown")

        with gr.Column(scale=7, min_width=520, elem_classes="preview-col"):
            with gr.Group(elem_classes="preview-card"):
                gr.HTML(
                    '<div class="preview-header"><div>'
                    '<p class="preview-title">🖼️ Comic Preview</p>'
                    '<p class="preview-sub">4 ô truyện sẽ được tạo và cập nhật lần lượt theo thời gian thực.</p>'
                    '</div></div>'
                )

                with gr.Column(elem_classes="comic-grid"):
                    panel_imgs = []
                    for row_start in (0, 2):
                        with gr.Row(equal_height=True, elem_classes="panel-row"):
                            for _ in range(row_start, row_start + 2):
                                panel_imgs.append(
                                    gr.Image(
                                        label="",
                                        show_label=False,
                                        type="pil",
                                        elem_classes="panel-img",
                                        height=300,
                                    )
                                )

    with gr.Accordion("📝 Storyboard & Prompts", open=False, elem_classes="story-panel"):
        gr.Markdown(
            "Mở phần này để xem storyboard JSON, caption đã parse và image prompt cho từng cảnh.",
            elem_classes="story-hint",
        )
        story_output = gr.Markdown(
            value="*Chưa chạy pipeline — kết quả storyboard sẽ hiển thị ở đây.*",
            elem_classes="story-markdown",
        )
        
    gr.HTML(
        '<div class="footer-text">'
        'Created by <a href="https://vlai.aivietnam.edu.vn/" target="_blank">VLAI</a>'
        ' &bull; <a href="https://aivietnam.edu.vn/" target="_blank">AI VIETNAM</a>'
        '</div>'
    )

    generate_btn.click(
        fn=inference,
        inputs=[prompt_input, token_input, gemini_key_input],
        outputs=[*panel_imgs, story_output, status],
    )

if __name__ == "__main__":
    demo.launch(
        allowed_paths=["static/aivn_logo.png"],
        css=_CSS
    )