"""MoxhiMT-30 Demo Space - Chinese to Vietnamese web-novel translation.

Uses CTranslate2 INT8 runtime on CPU with batched inference.
"""
import os
import re
import time
from pathlib import Path

import ctranslate2
import gradio as gr
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

MODEL_ID = "DanVP/MoxhiMT-60"
CT2_SUBDIR = "ct2-int8"
MAX_INPUT_TOKENS = 256
MAX_OUTPUT_TOKENS = 256
CT2_THREADS = min(2, max(1, os.cpu_count() or 1))
CT2_MAX_BATCH_SIZE = 8
DEFAULT_CHUNK_TOKENS = 128
SENTENCE_RE = re.compile(r"[^。！？!?；;\n]+[。！？!?；;]*[”’\"']*")
VI_SENTENCE_RE = re.compile(r".+?(?:[.!?…]+[”’\"']*|$)", re.S)
VI_BREAK_CHARS = set(".!?…。！？;；,:，、”’\"'")
HEADING_RE = re.compile(r"^第[0-9零〇一二三四五六七八九十百千万两]+[章节回卷部篇]")

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

print(f"Loading {MODEL_ID} ({CT2_SUBDIR})...")
MODEL_PATH = Path(snapshot_download(
    MODEL_ID,
    allow_patterns=[
        "config.json", "source.spm", "target.spm", "vocab.json",
        "tokenizer_config.json", "special_tokens_map.json",
        f"{CT2_SUBDIR}/*",
    ],
    token=os.environ.get("HF_TOKEN"),
))

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
translator = ctranslate2.Translator(
    str(MODEL_PATH / CT2_SUBDIR),
    device="cpu",
    compute_type="int8_float32",
    intra_threads=CT2_THREADS,
    inter_threads=1,
)
print(f"Loaded CTranslate2 CPU runtime. threads={CT2_THREADS}, batch={CT2_MAX_BATCH_SIZE}")


def source_token_ids(text, truncation=True):
    token_ids = tokenizer(text, truncation=truncation,
                          max_length=MAX_INPUT_TOKENS)["input_ids"]
    if tokenizer.pad_token_id is not None:
        token_ids = [t for t in token_ids if t != tokenizer.pad_token_id]
    return token_ids


def source_tokens(text):
    return tokenizer.convert_ids_to_tokens(source_token_ids(text, truncation=True))


def decode_tokens(tokens):
    return tokenizer.decode(tokenizer.convert_tokens_to_ids(tokens),
                            skip_special_tokens=True).strip()


def source_token_count(text):
    return len(source_token_ids(text, truncation=False))


def char_chunks(text, max_tokens=MAX_INPUT_TOKENS):
    chunks, current = [], ""
    for ch in text:
        cand = current + ch
        if current and source_token_count(cand) > max_tokens:
            chunks.append(current); current = ch
        else:
            current = cand
    if current: chunks.append(current)
    return chunks


def sentence_chunks(line, max_tokens=MAX_INPUT_TOKENS):
    if source_token_count(line) <= max_tokens: return [line]
    pieces = [m.group(0) for m in SENTENCE_RE.finditer(line)]
    if not pieces: return char_chunks(line, max_tokens)
    chunks, current = [], ""
    for p in pieces:
        if source_token_count(p) > max_tokens:
            if current: chunks.append(current); current = ""
            chunks.extend(char_chunks(p, max_tokens)); continue
        cand = current + p
        if current and source_token_count(cand) > max_tokens:
            chunks.append(current); current = p
        else:
            current = cand
    if current: chunks.append(current)
    return chunks


def source_sentence_pieces(line):
    pieces = [m.group(0).strip() for m in SENTENCE_RE.finditer(line or "")]
    pieces = [p for p in pieces if p]
    return pieces or ([line.strip()] if line and line.strip() else [])


def is_heading_line(line):
    stripped = (line or "").strip()
    if not stripped or "\n" in stripped:
        return False
    if len(stripped) > 24:
        return False
    if re.search(r"[。！？!?；;，,：:\"“”]", stripped):
        return False
    return bool(HEADING_RE.match(stripped))


def chunk_text_for_context(text, max_tokens=DEFAULT_CHUNK_TOKENS):
    chunks, current = [], ""
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        if is_heading_line(stripped):
            if current:
                chunks.append(current)
                current = ""
            chunks.append(stripped)
            continue

        for piece in source_sentence_pieces(stripped):
            if source_token_count(piece) > max_tokens:
                if current:
                    chunks.append(current)
                    current = ""
                chunks.extend(char_chunks(piece, max_tokens))
                continue
            cand = f"{current}\n{piece}" if current else piece
            if current and source_token_count(cand) > max_tokens:
                chunks.append(current)
                current = piece
            else:
                current = cand
    if current:
        chunks.append(current)
    return chunks


def split_vi_sentences(text):
    units = [re.sub(r"\s+", " ", m.group(0)).strip()
             for m in VI_SENTENCE_RE.finditer(text or "")]
    return [u for u in units if u]


def restore_line_breaks_by_sentences(source_text, target_text):
    source_lines = source_text.splitlines()
    nonblank = [line.strip() for line in source_lines if line.strip()]
    if len(nonblank) <= 1:
        return target_text.strip(), "single-line"

    target_units = split_vi_sentences(target_text)
    if len(target_units) < len(nonblank):
        return None, "too-few-target-sentences"

    total_source_chars = max(sum(len(line) for line in nonblank), 1)
    total_target_units = len(target_units)
    remaining_lines = len(nonblank)
    source_seen = 0
    target_cursor = 0
    restored = []

    for line in source_lines:
        stripped = line.strip()
        if not stripped:
            restored.append("")
            continue

        source_seen += len(stripped)
        remaining_lines -= 1
        ideal_end = round(source_seen / total_source_chars * total_target_units)
        min_end = target_cursor + 1
        max_end = total_target_units - remaining_lines
        end = min(max(ideal_end, min_end), max_end)
        restored.append(" ".join(target_units[target_cursor:end]).strip())
        target_cursor = end

    if target_cursor < total_target_units:
        for i in range(len(restored) - 1, -1, -1):
            if restored[i].strip():
                restored[i] = (restored[i] + " " + " ".join(target_units[target_cursor:])).strip()
                break
    return "\n".join(restored).strip(), "sentence-proportional"


def find_nearest_break(text, desired, min_pos, max_pos):
    desired = min(max(desired, min_pos), max_pos)
    window_start = max(min_pos, desired - 80)
    window_end = min(max_pos, desired + 80)
    best_idx, best_score = desired, float("inf")

    for idx in range(window_start, window_end + 1):
        prev = text[idx - 1] if idx > 0 else ""
        cur = text[idx] if idx < len(text) else ""
        if prev in VI_BREAK_CHARS or cur.isspace():
            score = abs(idx - desired)
            if prev in ".!?…。！？":
                score -= 8
            if score < best_score:
                best_idx, best_score = idx, score
    return best_idx


def restore_line_breaks_by_chars(source_text, target_text):
    source_lines = source_text.splitlines()
    nonblank = [line.strip() for line in source_lines if line.strip()]
    if len(nonblank) <= 1:
        return target_text.strip(), "single-line"

    compact_target = re.sub(r"\s+", " ", target_text or "").strip()
    if not compact_target:
        return "", "char-proportional"
    total_source_chars = max(sum(len(line) for line in nonblank), 1)
    target_len = len(compact_target)
    source_seen = 0
    last_pos = 0
    parts = []

    for line in nonblank[:-1]:
        source_seen += len(line)
        desired = round(source_seen / total_source_chars * target_len)
        remaining_parts = len(nonblank) - len(parts) - 1
        min_pos = min(last_pos + 1, target_len)
        max_pos = min(target_len, max(min_pos, target_len - remaining_parts))
        pos = find_nearest_break(compact_target, desired, min_pos, max_pos)
        parts.append(compact_target[last_pos:pos].strip())
        last_pos = pos
    parts.append(compact_target[last_pos:].strip())

    restored = []
    part_cursor = 0
    for line in source_lines:
        if line.strip():
            restored.append(parts[part_cursor] if part_cursor < len(parts) else "")
            part_cursor += 1
        else:
            restored.append("")
    return "\n".join(restored).strip(), "char-proportional"


def restore_line_breaks(source_text, target_text):
    source_lines = source_text.splitlines()
    first_nonblank = next((i for i, line in enumerate(source_lines) if line.strip()), None)
    target_lines = [line.strip() for line in (target_text or "").splitlines()]
    target_lines = [line for line in target_lines if line]

    if (first_nonblank is not None and len(target_lines) > 1 and
            is_heading_line(source_lines[first_nonblank])):
        body_start = first_nonblank + 1
        blank_after_heading = 0
        while body_start < len(source_lines) and not source_lines[body_start].strip():
            blank_after_heading += 1
            body_start += 1
        body_source = "\n".join(source_lines[body_start:])
        body_target = "\n".join(target_lines[1:])
        restored_body, body_mode = restore_line_breaks(body_source, body_target)
        restored = [""] * first_nonblank
        restored.append(target_lines[0])
        restored.extend([""] * blank_after_heading)
        if restored_body:
            restored.append(restored_body)
        return "\n".join(restored).strip(), f"heading+{body_mode}"

    restored, mode = restore_line_breaks_by_sentences(source_text, target_text)
    if restored is not None:
        return restored, mode
    return restore_line_breaks_by_chars(source_text, target_text)


def split_text(text):
    plan, chunks = [], []
    for line in text.splitlines():
        if not line.strip(): plan.append(None); continue
        lc = sentence_chunks(line.strip())
        idxs = []
        for c in lc: idxs.append(len(chunks)); chunks.append(c)
        plan.append(idxs)
    return plan, chunks


def translate_batch_texts(chunks, max_length=MAX_OUTPUT_TOKENS, beam_size=1):
    results = translator.translate_batch(
        [source_tokens(c) for c in chunks],
        max_batch_size=CT2_MAX_BATCH_SIZE,
        batch_type="examples",
        beam_size=max(1, int(beam_size)),
        max_decoding_length=int(max_length),
        no_repeat_ngram_size=3,
    )
    return [decode_tokens(r.hypotheses[0]) for r in results]


def translate_per_line(text, max_length=MAX_OUTPUT_TOKENS, beam_size=1):
    plan, chunks = split_text(text)
    if not chunks:
        return "", 0, "none"
    translated = translate_batch_texts(chunks, max_length=max_length, beam_size=beam_size)
    lines = []
    for idxs in plan:
        if idxs is None: lines.append("")
        else:
            lines.append(" ".join(translated[i] for i in idxs if translated[i]))
    return "\n".join(lines), len(chunks), "source-lines"


def translate_per_chunk(text, max_length=MAX_OUTPUT_TOKENS, beam_size=1,
                        chunk_tokens=DEFAULT_CHUNK_TOKENS):
    chunk_tokens = min(256, max(64, int(chunk_tokens)))
    chunks = chunk_text_for_context(text, max_tokens=chunk_tokens)
    if not chunks:
        return "", 0, "none"
    translated = translate_batch_texts(chunks, max_length=max_length, beam_size=beam_size)
    parts = []
    for source_chunk, translated_chunk in zip(chunks, translated):
        translated_chunk = (translated_chunk or "").strip()
        if not translated_chunk:
            continue
        if is_heading_line(source_chunk):
            parts.append(f"\n{translated_chunk}\n")
        else:
            parts.append(translated_chunk)
    raw_output = " ".join(parts)
    raw_output = re.sub(r"[ \t]+", " ", raw_output)
    raw_output = re.sub(r" *\n *", "\n", raw_output).strip()
    restored, restore_mode = restore_line_breaks(text, raw_output)
    return restored, len(chunks), restore_mode


def translate_with_metadata(text, mode="Per line", max_length=MAX_OUTPUT_TOKENS,
                            beam_size=1, chunk_tokens=DEFAULT_CHUNK_TOKENS):
    text = (text or "").strip("\r\n")
    if not text:
        return "", {"elapsed": 0.0, "chunks": 0, "source_chars": 0, "target_chars": 0}
    started = time.perf_counter()
    if mode == "Per chunk":
        output, chunks, restore_mode = translate_per_chunk(
            text, max_length=max_length, beam_size=beam_size,
            chunk_tokens=chunk_tokens,
        )
    else:
        output, chunks, restore_mode = translate_per_line(
            text, max_length=max_length, beam_size=beam_size,
        )
    return output, {
        "elapsed": time.perf_counter() - started,
        "chunks": chunks,
        "source_chars": len(text),
        "target_chars": len(output),
        "mode": mode,
        "restore_mode": restore_mode,
        "chunk_tokens": int(chunk_tokens),
    }


def format_stats(stats):
    elapsed = max(float(stats["elapsed"]), 1e-9)
    sc, tc, ch = int(stats["source_chars"]), int(stats["target_chars"]), int(stats["chunks"])
    mode = stats.get("mode", "Per line")
    restore = stats.get("restore_mode", "source-lines")
    chunk_tokens = int(stats.get("chunk_tokens", DEFAULT_CHUNK_TOKENS))
    mode_text = "per-chunk" if mode == "Per chunk" else "per-line"
    extra = f" · chunk_tok={chunk_tokens} · restore={restore}" if mode == "Per chunk" else ""
    return (f"{elapsed*1000:.0f}ms · mode={mode_text}{extra} · chunks={ch} · {sc}→{tc} chars · "
            f"{sc/elapsed:.1f} source chars/s · "
            f"CT2 int8 CPU · threads={CT2_THREADS} · batch={CT2_MAX_BATCH_SIZE}")


def translate_ui(text, mode="Per line", max_length=MAX_OUTPUT_TOKENS,
                 beam_size=1, chunk_tokens=DEFAULT_CHUNK_TOKENS):
    out, stats = translate_with_metadata(text, mode=mode, max_length=max_length,
                                          beam_size=beam_size,
                                          chunk_tokens=chunk_tokens)
    return out, format_stats(stats)


# Warm up
_ = translate_with_metadata("他抬头看向远处的山门。", max_length=128)

examples = [
    ["他抬头看向远处的山门。"],
    ["第一章 春"],
    ["她知道一定是修仙的压力太大，让张羽精神紧绷。"],
    ["第一章 春\n他抬头看向远处的山门。\n她知道一定是修仙的压力太大，让张羽精神紧绷。"],
]

with gr.Blocks(title="MoxhiMT 60 zh-vi") as demo:
    gr.Markdown("# MoxhiMT 60 — Chinese → Vietnamese MT\n"
                "Chinese web-novel translation demo (~57M params, CT2 int8).")
    with gr.Row():
        source = gr.Textbox(label="Chinese source", lines=16,
                            placeholder="Nhập tiếng Trung...")
        target = gr.Textbox(label="Vietnamese translation", lines=16,
                            show_copy_button=True)
    with gr.Row():
        mode = gr.Radio(["Per line", "Per chunk"], value="Per line",
                        label="Translation mode")
        max_length = gr.Slider(64, 512, value=MAX_OUTPUT_TOKENS, step=32,
                               label="Max output length")
        beam_size = gr.Slider(1, 5, value=4, step=1, label="Beam size")
    chunk_tokens = gr.Slider(64, 256, value=DEFAULT_CHUNK_TOKENS, step=32,
                             label="Chunk size (source tokens, used by Per chunk)")
    button = gr.Button("Translate")
    stats = gr.Textbox(label="Runtime stats", lines=1, interactive=False)
    button.click(translate_ui, inputs=[source, mode, max_length, beam_size, chunk_tokens],
                 outputs=[target, stats])
    gr.Examples(examples=examples, inputs=[source])

if __name__ == "__main__":
    demo.launch()