import gradio as gr
from .case_loader import load_case, run_case

# 시연용 케이스 이름 하드코딩
CASE_NAMES = [
    "폐공장에서 NPC와 대화하는 장면",
    "마을 대장장이와 무기 수리에 대해 대화하는 장면",
    "숲속 은둔자와 희귀 약초에 대해 대화하는 장면",
    "항구 관리관과 출항 허가에 대해 대화하는 장면",
    "마법사 견습생과 고대 주문서에 대해 대화하는 장면"
]

def format_case_info(case: dict) -> dict:
    """케이스 정보를 보기 좋게 정리해서 반환"""
    inp = case["input"]
    tags = inp.get("tags", {})
    context_lines = [f"{h['role'].upper()}: {h['text']}" for h in inp.get("context", [])]

    return {
        "npc_id": inp.get("npc_id", ""),
        "npc_location": inp.get("npc_location", ""),
        "quest_stage": tags.get("quest_stage", ""),
        "relationship": tags.get("relationship", ""),
        "trust": tags.get("trust", ""),
        "npc_mood": tags.get("npc_mood", ""),
        "player_reputation": tags.get("player_reputation", ""),
        "style": tags.get("style", ""),
        "lore": inp.get("lore", ""),
        "description": inp.get("description", ""),
        "player_state": inp.get("player_state", {}),
        "context": "\n".join(context_lines),
        "player_utterance": inp.get("player_utterance", "")
    }

def build_ui():
    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple")) as demo:
        gr.Markdown("""
        # 👾 PersonaChatEngine HF-Serve
        **게임 내 NPC 메인 모델 추론 서버**  
        Qwen 3B 기반 LoRA 파인튜닝 모델을 사용하여 NPC 대사를 생성합니다.
        """)

        with gr.Row():
            gr.Button("📄 상세 문서 보기",
                      link="https://huggingface.co/spaces/m97j/PersonaChatEngine_HF-serve/blob/main/README.md")
            gr.Button("💻 Colab 노트북 열기",
                      link="https://colab.research.google.com/drive/1_-qH8kdoU2Jj58TdaSnswHex-BFefInq?usp=sharing#scrollTo=cFJGv8BJ8oPD")

        gr.Markdown("### 🎯 테스트 케이스 기반 간단 실행")
        gr.Markdown("⚠️ 추론에는 수 초 ~ 최대 1분 정도 소요될 수 있습니다. 잠시만 기다려주세요.")

        with gr.Row():
            case_dropdown = gr.Dropdown(choices=CASE_NAMES, label="테스트 케이스 선택", value=CASE_NAMES[0])
            load_btn = gr.Button("케이스 불러오기")

        # 케이스 정보 표시 영역
        with gr.Row():
            with gr.Column():
                npc_id = gr.Textbox(label="NPC ID", interactive=False)
                npc_loc = gr.Textbox(label="NPC Location", interactive=False)
                quest_stage = gr.Textbox(label="Quest Stage", interactive=False)
                relationship = gr.Textbox(label="Relationship", interactive=False)
                trust = gr.Textbox(label="Trust", interactive=False)
                npc_mood = gr.Textbox(label="NPC Mood", interactive=False)
                player_rep = gr.Textbox(label="Player Reputation", interactive=False)
                style = gr.Textbox(label="Style", interactive=False)

            with gr.Column():
                lore = gr.Textbox(label="Lore", lines=3, interactive=False)
                desc = gr.Textbox(label="Description", lines=3, interactive=False)
                player_state = gr.JSON(label="Player State")
                context = gr.Textbox(label="Context", lines=6, interactive=False)

        # Player Utterance는 별도 입력창
        player_input = gr.Textbox(label="Player Utterance", lines=2)

        run_btn = gr.Button("🚀 Run Inference", variant="primary")
        npc_resp = gr.Textbox(label="NPC Response")
        deltas = gr.JSON(label="Deltas")
        flags = gr.JSON(label="Flags Probabilities")

        # 케이스 불러오기 동작
        def on_load_case(name):
            idx = CASE_NAMES.index(name)
            case = load_case(idx)  # TEST_CASES 직접 접근 대신 load_case 사용
            info = format_case_info(case)
            return (
                info["npc_id"], info["npc_location"], info["quest_stage"],
                info["relationship"], info["trust"], info["npc_mood"],
                info["player_reputation"], info["style"], info["lore"],
                info["description"], info["player_state"], info["context"],
                info["player_utterance"]
            )

        load_btn.click(
            fn=on_load_case,
            inputs=[case_dropdown],
            outputs=[
                npc_id, npc_loc, quest_stage, relationship, trust,
                npc_mood, player_rep, style, lore, desc, player_state, context,
                player_input
            ]
        )

        # 추론 실행
        run_btn.click(
            fn=lambda name, utt: run_case(CASE_NAMES.index(name), utt),
            inputs=[case_dropdown, player_input],
            outputs=[npc_resp, deltas, flags]
        )

        gr.Markdown("""
        ---
        ⚠️ **실제 게임 파이프라인 테스트**는 [ai-server Swagger](https://huggingface.co/spaces/m97j/PersonaChatEngine_ai_server)에서 진행하세요.
        """)

    return demo