"""Palette-K-Midm live demo — 회의록 → 결정사항 추출 (PALETTE-BENCH-KO #1 model).

HF Space (ZeroGPU). Deploy: upload this folder as a Gradio Space, hardware = ZeroGPU.
"""
import json
import re

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "imcapsule/palette-k-midm"  # TODO: replace ORG on drop day

tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto"
)

PROMPT = """다음 회의록에서 결정사항을 추출하라.

규칙:
- 실제로 확정된 결정만 포함한다. 보류/검토중/논의만 된 항목은 제외한다.
- 각 결정에 담당자(owner)와 기한(due_date, YYYY-MM-DD)을 붙인다. 기한이 명시되지 않았으면 null.
- 상대 날짜("다음 주 금요일까지")는 회의 날짜 기준으로 환산한다.
- JSON만 출력한다: {{"decisions": [{{"decision": "...", "owner": "...", "due_date": "..."}}]}}

회의록:
{minutes}"""

EXAMPLE = """[주간 운영회의] 2026-08-17 (월) 14:00 · 참석: 김부장, 이과장, 박대리, 최주임

김부장: 3분기 예산안은 지난주 논의대로 확정합니다. 이과장이 이번 주 금요일까지 품의서 올려주세요.
이과장: 네. 그리고 신규 거래처 계약 건은 법무 검토가 아직이라 다음 회의로 넘기겠습니다.
박대리: 사무실 이전 견적 3곳 받았습니다. 비교표는 제가 정리 중입니다.
김부장: 비교표는 정리되는 대로 공유만 해주세요. 아, 그리고 전사 보안교육은 9월 5일로 확정. 최주임이 공지 담당.
최주임: 알겠습니다."""


def _extract_json(s: str):
    start = s.find("{")
    while start != -1:
        depth, in_str, esc = 0, False, False
        for i in range(start, len(s)):
            c = s[i]
            if in_str:
                if esc:
                    esc = False
                elif c == "\\":
                    esc = True
                elif c == '"':
                    in_str = False
                continue
            if c == '"':
                in_str = True
            elif c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(s[start : i + 1])
                    except Exception:
                        break
        start = s.find("{", start + 1)
    return None


@spaces.GPU(duration=60)
def extract(minutes: str):
    if not minutes.strip():
        return "회의록을 입력하세요.", ""
    msgs = [{"role": "user", "content": PROMPT.format(minutes=minutes.strip())}]
    ids = tok.apply_chat_template(
        msgs, return_tensors="pt", add_generation_prompt=True
    ).to(model.device)
    out = model.generate(ids, max_new_tokens=800, do_sample=False)
    text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True)
    parsed = _extract_json(text)
    if parsed and "decisions" in parsed:
        rows = "\n".join(
            f"| {d.get('decision','')} | {d.get('owner','')} | {d.get('due_date') or '—'} |"
            for d in parsed["decisions"]
        )
        table = "| 결정사항 | 담당자 | 기한 |\n|---|---|---|\n" + rows
        return table, text
    return "JSON 파싱 실패 — 원문 출력을 확인하세요.", text


with gr.Blocks(title="Palette-K-Midm · 회의록→결정 추출") as demo:
    gr.Markdown(
        "# Palette-K-Midm — 회의록에서 결정사항 추출\n"
        "PALETTE-BENCH-KO **1위 (0.866)** 모델의 라이브 데모. "
        "보류·논의중 항목 제외, 상대 날짜 환산, 기한 없으면 null — "
        "전 모델 공통 난제 과제입니다.\n\n"
        "_정직 고지: 1위 격차는 0.010(n=32)이며, 이 과제 자체는 최고 모델도 0.582 수준의 난제입니다._"
    )
    inp = gr.Textbox(lines=12, label="회의록", value=EXAMPLE)
    btn = gr.Button("결정사항 추출", variant="primary")
    out_table = gr.Markdown(label="추출 결과")
    out_raw = gr.Textbox(lines=6, label="모델 원문 출력")
    btn.click(extract, inputs=inp, outputs=[out_table, out_raw])
    gr.Markdown(
        "[벤치마크](https://github.com/palette-lab/palette-bench-ko) · "
        "[모델](https://huggingface.co/imcapsule/palette-k-midm) · "
        "[기술 보고서](https://pltt.xyz/paper) · Apache-2.0, base: KT Mi:dm-2.0 (MIT)"
    )

demo.launch()