ui/arena_tab.py

"""A/B Arena 탭 UI"""

import gradio as gr
import random
from typing import Dict, List, Tuple, Optional, Any

from models.model_registry import get_all_models, get_model_info, get_models_for_dropdown
from characters import get_character_loader, build_system_prompt
from scenarios import get_scenario_loader
from voting import get_vote_storage, get_elo_calculator
from utils import parse_thinking_response, format_thinking_for_display


def create_arena_tab(
    model_manager: Any = None,
    use_mock: bool = False,
):
    """A/B Arena 탭 생성"""

    # 데이터 로더
    char_loader = get_character_loader()
    scenario_loader = get_scenario_loader()
    vote_storage = get_vote_storage()
    elo_calculator = get_elo_calculator()

    # 모델 목록
    all_models = get_all_models()
    model_choices = [(f"{get_model_info(m).get('description', m)}", m) for m in all_models]

    # 캐릭터 목록
    characters = char_loader.get_character_names()

    # 시나리오 목록
    scenario_choices = scenario_loader.get_scenarios_for_dropdown()

    # ============================================================
    # UI 구성
    # ============================================================

    gr.Markdown("## A/B 테스트 아레나")
    gr.Markdown("두 모델의 응답을 비교하고 더 좋은 응답에 투표하세요.")

    # 설정 패널
    with gr.Row():
        with gr.Column(scale=1):
            character_dropdown = gr.Dropdown(
                choices=characters,
                value=characters[0] if characters else None,
                label="캐릭터 선택",
            )
        with gr.Column(scale=1):
            scenario_dropdown = gr.Dropdown(
                choices=scenario_choices,
                value=scenario_choices[0][1] if scenario_choices else None,
                label="시나리오 프리셋",
            )
        with gr.Column(scale=1):
            blind_mode = gr.Checkbox(
                value=True,
                label="블라인드 모드 (모델명 숨김)",
            )

    with gr.Row():
        with gr.Column(scale=2):
            model_a_dropdown = gr.Dropdown(
                choices=model_choices,
                value=all_models[0] if all_models else None,
                label="Model A",
            )
        with gr.Column(scale=2):
            model_b_dropdown = gr.Dropdown(
                choices=model_choices,
                value=all_models[1] if len(all_models) > 1 else None,
                label="Model B",
            )
        with gr.Column(scale=1):
            random_models_btn = gr.Button("랜덤 모델", size="sm")

    # 응답 영역
    with gr.Row():
        # Model A Response
        with gr.Column(scale=1):
            model_a_label = gr.Markdown("### Model A")
            with gr.Accordion("Thinking Process", open=False):
                thinking_a = gr.Markdown("*(응답 생성 후 표시됩니다)*")
            response_a = gr.Textbox(
                label="응답",
                lines=8,
                interactive=False,
            )
            metadata_a = gr.Markdown("")

        # Model B Response
        with gr.Column(scale=1):
            model_b_label = gr.Markdown("### Model B")
            with gr.Accordion("Thinking Process", open=False):
                thinking_b = gr.Markdown("*(응답 생성 후 표시됩니다)*")
            response_b = gr.Textbox(
                label="응답",
                lines=8,
                interactive=False,
            )
            metadata_b = gr.Markdown("")

    # 사용자 입력
    with gr.Row():
        user_input = gr.Textbox(
            label="팬 메시지",
            placeholder="아이돌에게 보낼 메시지를 입력하세요...",
            lines=2,
            scale=4,
        )
        with gr.Column(scale=1):
            random_scenario_btn = gr.Button("랜덤 시나리오")
            submit_btn = gr.Button("전송", variant="primary")

    # 투표 영역
    gr.Markdown("### 투표")
    with gr.Row():
        vote_a_btn = gr.Button("A가 더 좋음", variant="secondary")
        vote_tie_btn = gr.Button("비슷함", variant="secondary")
        vote_b_btn = gr.Button("B가 더 좋음", variant="secondary")
        vote_skip_btn = gr.Button("스킵", variant="secondary")

    vote_reason = gr.Textbox(
        label="투표 이유 (선택사항)",
        placeholder="왜 이 응답이 더 좋다고 생각하시나요?",
        lines=1,
    )

    vote_result = gr.Markdown("")

    # 상태 저장
    state = gr.State({
        "model_a": None,
        "model_b": None,
        "response_a": None,
        "response_b": None,
        "character": None,
        "user_input": None,
    })

    # ============================================================
    # 이벤트 핸들러
    # ============================================================

    def select_random_models():
        """랜덤 모델 선택"""
        if len(all_models) < 2:
            return all_models[0] if all_models else None, None
        selected = random.sample(all_models, 2)
        return selected[0], selected[1]

    def load_random_scenario(character: str):
        """랜덤 시나리오 로드"""
        scenario = scenario_loader.get_random_scenario()
        if scenario:
            user_msg = scenario_loader.format_user_input(scenario, character)
            return user_msg, scenario["id"]
        return "", None

    def load_scenario_input(scenario_id: str, character: str):
        """선택된 시나리오 로드"""
        scenario = scenario_loader.get_scenario(scenario_id)
        if scenario:
            return scenario_loader.format_user_input(scenario, character)
        return ""

    def generate_responses(
        model_a: str,
        model_b: str,
        character: str,
        user_msg: str,
        current_state: dict,
    ):
        """두 모델의 응답 생성"""
        if not model_a or not model_b:
            return (
                "*(모델을 선택해주세요)*", "", "",
                "*(모델을 선택해주세요)*", "", "",
                current_state,
            )

        if not user_msg.strip():
            return (
                "*(메시지를 입력해주세요)*", "", "",
                "*(메시지를 입력해주세요)*", "", "",
                current_state,
            )

        system_prompt = build_system_prompt(character)
        messages = [{"role": "user", "content": user_msg}]

        # Mock 모드 (실제 모델 없이 테스트)
        if use_mock or model_manager is None:
            response_a_full = f"<think>\n{character}의 입장에서 생각해보면... 이 메시지에 어떻게 반응해야 할까?\n</think>\n\n안녕! 반가워~ (Mock Response A)"
            response_b_full = f"<think>\n음... 이런 상황에서는...\n</think>\n\n헤이~ 뭐해? (Mock Response B)"
            meta_a = {"latency_s": 0.5, "output_tokens": 50}
            meta_b = {"latency_s": 0.6, "output_tokens": 55}
        else:
            # 실제 모델 추론
            try:
                response_a_full, meta_a = model_manager.generate_response(
                    model_a, messages, system_prompt
                )
            except Exception as e:
                response_a_full = f"*Error: {str(e)}*"
                meta_a = {"latency_s": 0, "output_tokens": 0}

            try:
                response_b_full, meta_b = model_manager.generate_response(
                    model_b, messages, system_prompt
                )
            except Exception as e:
                response_b_full = f"*Error: {str(e)}*"
                meta_b = {"latency_s": 0, "output_tokens": 0}

        # Thinking 파싱
        think_a, clean_a = parse_thinking_response(response_a_full)
        think_b, clean_b = parse_thinking_response(response_b_full)

        # 메타데이터 문자열
        meta_str_a = f"⏱️ {meta_a.get('latency_s', 0):.2f}s | {meta_a.get('output_tokens', 0)} tokens"
        meta_str_b = f"⏱️ {meta_b.get('latency_s', 0):.2f}s | {meta_b.get('output_tokens', 0)} tokens"

        # 상태 업데이트
        new_state = {
            "model_a": model_a,
            "model_b": model_b,
            "response_a": response_a_full,
            "response_b": response_b_full,
            "character": character,
            "user_input": user_msg,
        }

        return (
            format_thinking_for_display(think_a) if think_a else "*No thinking*",
            clean_a,
            meta_str_a,
            format_thinking_for_display(think_b) if think_b else "*No thinking*",
            clean_b,
            meta_str_b,
            new_state,
        )

    def handle_vote(vote_type: str, reason: str, current_state: dict):
        """투표 처리"""
        if not current_state.get("model_a") or not current_state.get("model_b"):
            return "먼저 응답을 생성해주세요."

        vote_data = {
            "model_a": current_state["model_a"],
            "model_b": current_state["model_b"],
            "response_a": current_state.get("response_a", ""),
            "response_b": current_state.get("response_b", ""),
            "character": current_state.get("character", ""),
            "user_input": current_state.get("user_input", ""),
            "vote": vote_type,
            "reason": reason,
        }

        vote_id = vote_storage.save_vote(vote_data)

        # ELO 업데이트
        if vote_type != "skip":
            new_a, new_b = elo_calculator.update_ratings(
                current_state["model_a"],
                current_state["model_b"],
                vote_type,
            )
            return f"투표 완료! (ID: {vote_id})\n\nELO 변경:\n- {current_state['model_a']}: {new_a:.0f}\n- {current_state['model_b']}: {new_b:.0f}"

        return f"스킵됨 (ID: {vote_id})"

    def update_model_labels(blind: bool, model_a: str, model_b: str):
        """블라인드 모드에 따라 레이블 업데이트"""
        if blind:
            return "### Model A", "### Model B"
        else:
            info_a = get_model_info(model_a)
            info_b = get_model_info(model_b)
            label_a = f"### {info_a.get('description', model_a)}" if info_a else f"### {model_a}"
            label_b = f"### {info_b.get('description', model_b)}" if info_b else f"### {model_b}"
            return label_a, label_b

    # ============================================================
    # 이벤트 바인딩
    # ============================================================

    random_models_btn.click(
        fn=select_random_models,
        outputs=[model_a_dropdown, model_b_dropdown],
    )

    random_scenario_btn.click(
        fn=load_random_scenario,
        inputs=[character_dropdown],
        outputs=[user_input, scenario_dropdown],
    )

    scenario_dropdown.change(
        fn=load_scenario_input,
        inputs=[scenario_dropdown, character_dropdown],
        outputs=[user_input],
    )

    submit_btn.click(
        fn=generate_responses,
        inputs=[model_a_dropdown, model_b_dropdown, character_dropdown, user_input, state],
        outputs=[thinking_a, response_a, metadata_a, thinking_b, response_b, metadata_b, state],
    )

    # 블라인드 모드 변경 시 레이블 업데이트
    blind_mode.change(
        fn=update_model_labels,
        inputs=[blind_mode, model_a_dropdown, model_b_dropdown],
        outputs=[model_a_label, model_b_label],
    )

    # 투표 버튼
    vote_a_btn.click(
        fn=lambda r, s: handle_vote("a", r, s),
        inputs=[vote_reason, state],
        outputs=[vote_result],
    )
    vote_b_btn.click(
        fn=lambda r, s: handle_vote("b", r, s),
        inputs=[vote_reason, state],
        outputs=[vote_result],
    )
    vote_tie_btn.click(
        fn=lambda r, s: handle_vote("tie", r, s),
        inputs=[vote_reason, state],
        outputs=[vote_result],
    )
    vote_skip_btn.click(
        fn=lambda r, s: handle_vote("skip", r, s),
        inputs=[vote_reason, state],
        outputs=[vote_result],
    )