"""Han2Han demo: Hangul and Hanja transcription with cadazar/han2han-rl.

Each message is converted on its own. Ratings and corrections of replies are
appended to a JSONL file that a CommitScheduler pushes to the dataset repo
named by the FEEDBACK_REPO variable; inputs nobody rates or corrects are not
stored.
"""

import json
import logging
import os
import re
import threading
import uuid
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr
import torch
from huggingface_hub import CommitScheduler
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

logging.basicConfig(level=logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)
logger = logging.getLogger("han2han-space")

MODEL_REPO = "cadazar/han2han-rl"
MODEL_REVISION = "49395d62d5f5da084b587d31d5bffdf3410004a3"
MAX_INPUT_CHARS = 2000
CHUNK_CHARS = 300

TO_HANJA = "한글을 한자로 (Hangul to Hanja)"
TO_HANGUL = "한자를 한글로 (Hanja to Hangul)"
TASK_PROMPTS = {
    TO_HANJA: "한글을 한자로 전사하시오:",
    TO_HANGUL: "한자를 한글로 전사하시오:",
}

EXAMPLE_HANGUL = (
    "단순한 스케치나 카메라적 전사가 아니라 작가의 예술적 정서의 전체적인 안계로 "
    "이루어지는 가장 사실적인 표현을 말하는 것이다."
)
EXAMPLE_MIXED = (
    "單純한 스케치나 카메라的 轉寫가 아니라 作家의 藝術的 情緖의 全體的인 眼界로 "
    "이루어지는 가장 寫實的인 表現을 말하는 것이다."
)

MODEL_LINK = f"[`{MODEL_REPO}`](https://huggingface.co/{MODEL_REPO})"
UI_TEXT = {
    "en": {
        "intro": (
            "# Han2Han\n"
            f"Convert Korean text between Hangul and mixed Hangul-Hanja script with {MODEL_LINK}. "
            "Each message is converted on its own; the model never sees earlier messages. "
            "Restoring Hanja is the hard direction: check the output against a source before "
            "quoting it."
        ),
        "direction": "Direction",
        "clear": "Clear chat",
        "chat_placeholder": "Type or paste Korean text below.",
        "feedback": (
            "Rate a reply with the thumbs, or edit it to submit a correction. Rated and "
            "corrected replies are stored together with their input text in a private "
            "dataset and may be used to improve the model; nothing else you enter is stored."
        ),
        "empty_input": "Enter some text to convert.",
        "too_long": "Please keep inputs under {limit} characters ({given} given).",
        "user_edit": "Edit a reply to correct it; send a new message to convert other text.",
    },
    "ko": {
        "intro": (
            "# Han2Han\n"
            f"{MODEL_LINK} 모델로 한국어 텍스트를 한글 전용과 국한문 혼용 사이에서 변환합니다. "
            "메시지는 하나씩 따로 변환되며, 모델은 이전 메시지를 보지 않습니다. 한자 복원은 더 "
            "어려운 방향이니, 결과를 인용하기 전에 원문과 대조해 확인하세요."
        ),
        "direction": "변환 방향",
        "clear": "대화 지우기",
        "chat_placeholder": "아래에 한국어 텍스트를 입력하거나 붙여 넣으세요.",
        "feedback": (
            "답변 아래의 엄지 버튼으로 평가하거나, 답변을 직접 고쳐 교정본을 보낼 수 있습니다. "
            "평가하거나 교정한 답변은 입력 텍스트와 함께 비공개 데이터셋에 저장되며, 모델 개선에 "
            "쓰일 수 있습니다. 그 밖에 입력한 내용은 저장되지 않습니다."
        ),
        "empty_input": "변환할 텍스트를 입력하세요.",
        "too_long": "입력은 {limit}자 이내로 해 주세요. (현재 {given}자)",
        "user_edit": "교정하려면 답변을 고쳐 주세요. 다른 텍스트를 변환하려면 새 메시지를 보내세요.",
    },
}

# cpu-basic has 2 vCPUs; torch would otherwise size its pool to the host's cores
torch.set_num_threads(int(os.environ.get("TORCH_NUM_THREADS", "2")))

tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, revision=MODEL_REVISION, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_REPO, revision=MODEL_REVISION, trust_remote_code=True
).eval()

FEEDBACK_REPO = os.environ.get("FEEDBACK_REPO")
FEEDBACK_DIR = Path("feedback")
FEEDBACK_DIR.mkdir(exist_ok=True)
# one file per process, so a restart never rewrites rows already pushed
FEEDBACK_FILE = FEEDBACK_DIR / f"feedback-{uuid.uuid4().hex}.jsonl"

if FEEDBACK_REPO:
    scheduler = CommitScheduler(
        repo_id=FEEDBACK_REPO,
        repo_type="dataset",
        folder_path=FEEDBACK_DIR,
        path_in_repo="data",
        every=5,
        private=True,
    )
    feedback_lock = scheduler.lock
elif os.environ.get("SPACE_ID"):
    raise RuntimeError("FEEDBACK_REPO is not set: ratings and corrections would be lost on restart.")
else:
    logger.warning("FEEDBACK_REPO is not set: feedback stays in %s and is not pushed.", FEEDBACK_FILE)
    feedback_lock = threading.Lock()


def log_feedback(row: dict) -> None:
    row = {
        "time": datetime.now(timezone.utc).isoformat(),
        "model": MODEL_REPO,
        "revision": MODEL_REVISION,
        **row,
    }
    with feedback_lock:
        with FEEDBACK_FILE.open("a", encoding="utf-8") as f:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    logger.info("feedback logged: %s %s", row["event"], row["id"])


def message_text(content) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, dict) and content.get("type") == "text":
        return content["text"]
    if isinstance(content, list):
        return "".join(message_text(part) for part in content)
    raise TypeError(f"unexpected chat message content: {content!r}")


def split_line(line: str) -> list[str]:
    """Pack the sentences of one line into pieces of about CHUNK_CHARS characters."""
    pieces, current = [], ""
    for sentence in re.split(r"(?<=[.?!。])\s+", line.strip()):
        if current and len(current) + 1 + len(sentence) > CHUNK_CHARS:
            pieces.append(current)
            current = sentence
        else:
            current = f"{current} {sentence}" if current else sentence
    if current:
        pieces.append(current)
    return pieces


def stream_piece(task_prompt: str, text: str):
    """Yield the reply to one piece of text as it is generated (greedy decoding)."""
    messages = [
        {"role": "system", "content": task_prompt},
        {"role": "user", "content": text},
    ]
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    )
    budget = 2 * len(tokenizer(text).input_ids) + 8
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    errors = []

    def generate():
        try:
            with torch.no_grad():
                model.generate(**inputs, max_new_tokens=budget, do_sample=False, streamer=streamer)
        except Exception as e:
            logger.exception("generation failed")
            errors.append(e)
            streamer.end()

    thread = threading.Thread(target=generate)
    thread.start()
    for new_text in streamer:
        yield new_text
    thread.join()
    if errors:
        raise gr.Error("Generation failed; see the Space logs.") from errors[0]


def stream_convert(direction: str, text: str):
    """Yield the converted text so far, line by line and piece by piece."""
    task_prompt = TASK_PROMPTS[direction]
    done_lines = []
    for line in text.split("\n"):
        line_out = ""
        for piece in split_line(line):
            prefix = f"{line_out} " if line_out else ""
            piece_out = ""
            for new_text in stream_piece(task_prompt, piece):
                piece_out += new_text
                yield "\n".join(done_lines + [prefix + piece_out])
            line_out = prefix + piece_out.strip()
        done_lines.append(line_out)
    yield "\n".join(done_lines)


def respond(message: str, history: list, direction: str, records: dict, language: str):
    """Convert the message in the chosen direction and stream the reply into the chat."""
    text = message.strip()
    if not text:
        raise gr.Error(UI_TEXT[language]["empty_input"])
    if len(text) > MAX_INPUT_CHARS:
        raise gr.Error(UI_TEXT[language]["too_long"].format(limit=MAX_INPUT_CHARS, given=len(text)))
    history = history + [
        {"role": "user", "content": text},
        {"role": "assistant", "content": ""},
    ]
    output = ""
    for output in stream_convert(direction, text):
        history[-1]["content"] = output
        yield history, "", records
    records = {
        **records,
        len(history) - 1: {"id": uuid.uuid4().hex, "direction": direction, "input": text, "output": output},
    }
    yield history, "", records


def message_index(index) -> int:
    return index[0] if isinstance(index, (list, tuple)) else index


def on_like(data: gr.LikeData, records: dict) -> None:
    record = records.get(message_index(data.index))
    if record is None:
        logger.warning("rating for message %s, which has no record", data.index)
        return
    log_feedback({"event": "rating", "liked": data.liked, **record})


def on_edit(data: gr.EditData, history: list, records: dict, language: str):
    index = message_index(data.index)
    record = records.get(index)
    if record is None:
        history[index]["content"] = message_text(data.previous_value)
        gr.Info(UI_TEXT[language]["user_edit"])
        return history, records
    correction = message_text(data.value)
    history[index]["content"] = correction
    log_feedback({"event": "correction", "correction": correction, **record})
    records = {**records, index: {**record, "corrected": True}}
    return history, records


def detect_language(request: gr.Request) -> str:
    accept = request.headers.get("accept-language") or ""
    return "ko" if accept.lower().startswith("ko") else "en"


def apply_language(language: str):
    text = UI_TEXT[language]
    return (
        gr.update(value=text["intro"]),
        gr.update(label=text["direction"]),
        gr.update(placeholder=text["chat_placeholder"]),
        gr.update(value=text["clear"]),
        gr.update(value=text["feedback"]),
    )


def clear_chat():
    return [], "", {}


def reset_records() -> dict:
    return {}


with gr.Blocks(title="Han2Han") as demo:
    with gr.Row(equal_height=False):
        intro = gr.Markdown(UI_TEXT["en"]["intro"])
        language = gr.Radio(
            [("English", "en"), ("한국어", "ko")],
            value="en",
            show_label=False,
            container=False,
            scale=0,
            min_width=180,
        )
    direction = gr.Radio(list(TASK_PROMPTS), value=TO_HANJA, label=UI_TEXT["en"]["direction"])
    chatbot = gr.Chatbot(
        height=480,
        editable="all",
        placeholder=UI_TEXT["en"]["chat_placeholder"],
    )
    textbox = gr.Textbox(
        placeholder="학문을 닦고 덕을 쌓다",
        show_label=False,
        submit_btn=True,
        lines=2,
        max_lines=12,
    )
    clear_button = gr.Button(UI_TEXT["en"]["clear"], size="sm")
    gr.Examples(
        [[TO_HANJA, EXAMPLE_HANGUL], [TO_HANGUL, EXAMPLE_MIXED]],
        inputs=[direction, textbox],
        label="예시 / Examples",
    )
    feedback_note = gr.Markdown(UI_TEXT["en"]["feedback"])
    records = gr.State({})

    demo.load(detect_language, None, language, api_visibility="private")
    language.change(
        apply_language,
        language,
        [intro, direction, chatbot, clear_button, feedback_note],
        api_visibility="private",
    )
    submit_event = textbox.submit(
        respond, [textbox, chatbot, direction, records, language], [chatbot, textbox, records]
    )
    clear_button.click(
        clear_chat, None, [chatbot, textbox, records], cancels=[submit_event], api_visibility="private"
    )
    chatbot.clear(reset_records, None, records, cancels=[submit_event], api_visibility="private")
    chatbot.like(on_like, [records], None, api_visibility="private")
    chatbot.edit(on_edit, [chatbot, records, language], [chatbot, records], api_visibility="private")


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()