"""bulk 분할 계획 -- 상한 안에서 문단 > 문장 > 어절 경계만 자르고, 구분자를 원장으로 남겨 접합 때 복원.

왜 이 모듈인가: 옛 `_split_text` 는 500번째 **문자**에서 무조건 잘랐다. 절단면이 공백이면
`.rstrip()` + `"".join` 이 공백을 삼키고(`미국 제너럴다이내믹스`→`미국제너럴다이내믹스`), 어절
중간이면 모델이 뒤 조각을 완성하거나(`들이`→`이들이`) 앞 조각을 버렸고(`이뤄졌다. 지난`→
`이뤄졌다.난`), `\\n`→`∥PARA∥` 왕복은 마커 한가운데를 잘라 리터럴을 출력에 남겼다. 상세는
`accounts/chosun/body-proofread/03-evaluation/insights/lab_notes_archive.md#260910_long_article_regression_set`.

왜 greedy(상한까지 채우고 가장 가까운 경계) 가 아니라 DP 인가:
- greedy 는 꼬리 chunk 를 수십 자로 남긴다 (1022자 → 500/500/22). 조각만 받은 모델은 헛손질한다.
- 따옴표가 열린 채 끝나는 문장 경계는 약하다 -- 뒤 chunk 가 닫는 따옴표로 시작하면 모델이 "고친다".
  후보 경계마다 종류별 비용을 두고 chunk 수·짧은 chunk 벌점까지 합쳐 전체 최소를 고르면 둘 다 풀린다.

상한(`max_len`)은 파인튜닝 모델이 잘 도는 구간이라 **절대 넘기지 않는다**. 상한 이하 텍스트는 손대지
않고 그대로 1개 chunk 로 돌려준다 (기존 평가셋 결과가 바이트 단위로 보존되는 조건).
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass

WS_RUN = re.compile(r"\s+")
SENTENCE_END = ".!?…。！？"
# 문장 끝 부호 뒤에 붙을 수 있는 닫는 기호 -- `했다.”`, `있다)`, `한다.’` 도 문장 경계로 본다
TRAILING_CLOSERS = "”’\"')]」』》〉"

# 비용은 "선호 순서와 교환 비율" 이다. v2 (2026-09-10 ablation 반영): 문단 정렬이 호출 1회보다 우선한다 --
# v1(chunk 1.0 > sent 0.15) 은 호출을 아끼려 문장 경계로 ~450자까지 채웠는데, 그 큰 다문단 chunk 가
# 문단 단위 chunk 보다 recall 을 5점 깎았다(long-article 매칭 ΔF1 -1.4 vs greedy +0.4). 어절 절단·꼬리·
# 따옴표 안 절단 회피는 그대로다.
KIND_COST: dict[str, float] = {
    "para": 0.0,  # 문단 경계 (공백 런에 개행 포함) -- 공짜. 문단이 상한에 들어가는 한 문단째로 자른다
    "sent": 0.5,  # 문장 끝 부호 + 공백 -- 호출 하나(0.1)보다 비싸다: 문단을 쪼개느니 chunk 를 하나 더 쓴다
    "quoted_sent": 1.0,  # 따옴표/괄호가 열린 채 만난 문장 끝 -- 뒤 chunk 가 닫는 기호로 시작한다
    "word": 3.0,  # 어절 사이 공백 -- 문장 한가운데
    "char": 5.0,  # 공백 없는 런이 상한보다 길 때만 쓰는 강제 절단
}
CHUNK_COST = 0.1  # chunk 하나 = LLM 호출 하나. 같은 경계 품질이면 호출이 적은 쪽 (tie-breaker 수준)
SHORT_PENALTY = 3.0  # `min_len` 미만 chunk 벌점 계수 (부족 비율의 제곱에 곱한다) -- 꼬리 방지
DEFAULT_MIN_LEN = 150


@dataclass(frozen=True)
class Boundary:
    """chunk 사이 경계. `text[start:end]` 가 구분자(공백 런, 강제 절단이면 빈 문자열)."""

    start: int
    end: int
    kind: str


@dataclass(frozen=True)
class ChunkPlan:
    """`join_chunks(plan.chunks, plan)` 이 원문을 그대로 돌려주는 분할 계획."""

    chunks: list[str]
    seps: list[str]  # chunks 와 같은 길이. seps[i] 는 chunks[i] 뒤 구분자, 마지막은 trailing 공백
    kinds: list[str]  # 내부 경계 종류 (len == len(chunks) - 1)
    lead: str  # 원문 leading 공백

    @property
    def is_split(self) -> bool:
        return len(self.chunks) > 1


STRATEGIES = ("dp", "greedy")


def plan_chunks(
    text: str, max_len: int, min_len: int = DEFAULT_MIN_LEN, strategy: str = "dp"
) -> ChunkPlan:
    """`text` 를 `max_len` 이하 chunk 들로 나누는 계획을 세운다. 상한 이하면 그대로 1개.

    `strategy="greedy"` 는 "상한까지 채우고 그 안에서 가장 뒤의 좋은 경계" 라는 단순 사다리 —
    DP 와의 ablation 용이지 운영값이 아니다 (꼬리 chunk · 따옴표 안 절단을 못 막는다).
    """
    if max_len <= 0:
        raise ValueError(f"max_len 은 양수여야 한다: {max_len}")
    if strategy not in STRATEGIES:
        raise ValueError(f"strategy 는 {STRATEGIES} 중 하나여야 한다: {strategy!r}")
    if len(text) <= max_len:
        return ChunkPlan(chunks=[text], seps=[""], kinds=[], lead="")

    lead_len = len(text) - len(text.lstrip())
    body = text.strip()
    lead, trail = text[:lead_len], text[lead_len + len(body) :]
    if len(body) <= max_len:
        return ChunkPlan(chunks=[body], seps=[trail], kinds=[], lead=lead)

    boundaries = _candidate_boundaries(body)
    boundaries = _add_forced_cuts(body, boundaries, max_len)
    if strategy == "greedy":
        cut_ids = _solve_greedy(boundaries, len(body), max_len)
    else:
        cut_ids = _solve(body, boundaries, max_len, min_len)

    chunks, seps, kinds = [], [], []
    pos = 0
    for b in (boundaries[i] for i in cut_ids):
        chunks.append(body[pos : b.start])
        seps.append(body[b.start : b.end])
        kinds.append(b.kind)
        pos = b.end
    chunks.append(body[pos:])
    seps.append(trail)
    return ChunkPlan(chunks=chunks, seps=seps, kinds=kinds, lead=lead)


def join_chunks(outputs: list[str], plan: ChunkPlan) -> str:
    """LLM 을 거친 chunk 출력들을 원장의 구분자로 다시 잇는다. 출력 가장자리 공백은 모델 잡음이라 버린다."""
    if len(outputs) != len(plan.chunks):
        raise ValueError(f"chunk 수 불일치: outputs={len(outputs)} plan={len(plan.chunks)}")
    if not plan.is_split:
        return outputs[0]
    return plan.lead + "".join(
        out.strip() + sep for out, sep in zip(outputs, plan.seps, strict=True)
    )


# ── 출력 위생 ────────────────────────────────────────────────────────────

# 토큰 앞뒤의 공백·개행까지 흡수한다 -- 모델이 토큰과 개행을 같이 내면 개행이 겹친다 (실측: 문단 +1)
SEPARATOR_LITERAL = re.compile(r"\s*<paragraph_separator>\s*")
LITERAL = "<paragraph_separator>"
ANCHOR_LEN = 20


def normalize_separator_literals(text: str, source: str | None = None) -> str:
    """모델이 지어낸 `<paragraph_separator>` 리터럴을 그 자리의 원래 구분자로 되돌린다.

    step1~3 프롬프트는 이 토큰으로 문단 경계를 표시하라고 하지만 엔진은 주입도 제거도 하지 않는다
    (계약 불일치 -- insight 260910 참조). 그래서 모델이 가끔 토큰을 그대로 출력에 남기고, 문단 정렬
    chunk(DP v2) 에서 빈도가 올라간다(실기사 6건 run 당 3건). 계약을 (A) 주입/복원으로 살리든 (B)
    프롬프트에서 빼든, 최종 텍스트에 리터럴이 남는 건 어느 쪽에서도 옳지 않다 -- 이 가드는 그 결정과
    독립이다.

    무엇으로 바꾸나: `source`(그 chunk 의 입력)가 있으면 리터럴 직전 20자를 source 에서 찾아 그 뒤에 오는
    원래 공백 런을 쓴다 -- 원문이 개행이면 개행, 문장 사이 공백이면 공백. 모델이 문장 사이에 토큰을
    지어냈을 때 무조건 개행으로 바꾸면 문단이 늘어난다(실측: 기사 2건에서 +3·+2 문단). 앵커를 못 찾으면
    개행(토큰의 계약상 의미).
    """
    if LITERAL not in text:
        return text
    if source is None:
        return SEPARATOR_LITERAL.sub("\n", text)
    parts = SEPARATOR_LITERAL.split(text)
    out = parts[0]
    for nxt in parts[1:]:
        out += _original_separator(source, out) + nxt
    return out


def _original_separator(source: str, left: str) -> str:
    anchor = left[-ANCHOR_LEN:].rstrip()
    if anchor:
        i = source.rfind(anchor)
        if i >= 0:
            m = WS_RUN.match(source, i + len(anchor))
            if m:
                return "\n" if ("\n" in m.group() or "\r" in m.group()) else " "
            return " "  # 원문엔 공백조차 없던 자리 -- 붙여 쓰진 않고 최소 구분자
    return "\n"


# ── 후보 경계 ────────────────────────────────────────────────────────────


def _candidate_boundaries(body: str) -> list[Boundary]:
    open_depth = _open_quote_depth(body)
    out: list[Boundary] = []
    for m in WS_RUN.finditer(body):
        start, end = m.span()
        if start == 0 or end == len(body):
            continue  # strip() 했으므로 없지만 방어
        if "\n" in m.group() or "\r" in m.group():
            kind = "para"
        elif _ends_sentence(body, start, end):
            kind = "quoted_sent" if open_depth[start] > 0 else "sent"
        else:
            kind = "word"
        out.append(Boundary(start, end, kind))
    return out


def _ends_sentence(body: str, ws_start: int, ws_end: int) -> bool:
    i = ws_start - 1
    while i >= 0 and body[i] in TRAILING_CLOSERS:
        i -= 1
    if i < 0 or body[i] not in SENTENCE_END:
        return False
    # `3. 4.` 같은 번호 매김 -- 숫자 사이의 마침표는 문장 끝이 아니다
    numbered = (
        body[i] == "." and i > 0 and body[i - 1].isdigit() and body[ws_end : ws_end + 1].isdigit()
    )
    return not numbered


def _open_quote_depth(body: str) -> list[int]:
    """각 위치 직전까지 열려 있는 따옴표/괄호 수. 곧은따옴표(")는 홀수 개면 열린 것으로 본다."""
    pairs = {"“": 1, "”": -1, "‘": 1, "’": -1, "(": 1, ")": -1, "「": 1, "」": -1}
    depth: list[int] = [0] * (len(body) + 1)
    curly = 0
    straight = 0
    for i, ch in enumerate(body):
        depth[i] = max(curly, 0) + straight
        if ch == '"':
            straight ^= 1
        else:
            curly += pairs.get(ch, 0)
    depth[len(body)] = max(curly, 0) + straight
    return depth


def _add_forced_cuts(body: str, boundaries: list[Boundary], max_len: int) -> list[Boundary]:
    """공백 없는 런이 `max_len` 보다 길면 그 안에 강제 절단점을 넣는다 (URL 덩어리 같은 예외)."""
    anchors = [0] + [b.end for b in boundaries]
    starts = [b.start for b in boundaries] + [len(body)]
    forced: list[Boundary] = []
    for prev_end, next_start in zip(anchors, starts, strict=True):
        pos = prev_end + max_len
        while pos < next_start:
            forced.append(Boundary(pos, pos, "char"))
            pos += max_len
    return sorted(boundaries + forced, key=lambda b: b.start)


# ── DP ───────────────────────────────────────────────────────────────────


def _solve(body: str, boundaries: list[Boundary], max_len: int, min_len: int) -> list[int]:
    """전체 비용 최소인 경계 부분집합(인덱스, 오름차순)을 고른다.

    노드 0 = 본문 시작, 노드 k(1..K) = boundaries[k-1], 노드 K+1 = 본문 끝.
    chunk(i→j) 비용 = CHUNK_COST + 짧은 chunk 벌점 + (j 가 내부 경계면) 경계 종류 비용.
    """
    n_b = len(boundaries)
    node_end = [0] + [b.end for b in boundaries]  # chunk 시작 위치 (노드별)
    node_start = [b.start for b in boundaries] + [len(body)]  # chunk 끝 위치 (노드 1..K+1)
    best = [math.inf] * (n_b + 2)
    back = [-1] * (n_b + 2)
    best[0] = 0.0
    for j in range(1, n_b + 2):
        end_pos = node_start[j - 1]
        edge_cost = KIND_COST[boundaries[j - 1].kind] if j <= n_b else 0.0
        for i in range(j - 1, -1, -1):
            length = end_pos - node_end[i]
            if length > max_len:
                break
            if length <= 0 or best[i] == math.inf:
                continue
            cost = best[i] + CHUNK_COST + _short_penalty(length, min_len) + edge_cost
            if cost < best[j]:
                best[j], back[j] = cost, i
    if best[n_b + 1] == math.inf:
        raise RuntimeError("분할 계획 실패 -- 강제 절단점이 부족하다 (버그)")
    cut_ids: list[int] = []
    j = n_b + 1
    while back[j] > 0:
        j = back[j]
        cut_ids.append(j - 1)
    return sorted(cut_ids)


GREEDY_LADDER = (("para",), ("sent", "quoted_sent"), ("word",), ("char",))


def _solve_greedy(boundaries: list[Boundary], n: int, max_len: int) -> list[int]:
    """regex 사다리 재현: 창(pos, pos+max_len] 안에서 문단 > 문장 > 어절 > 강제 순으로, 가장 뒤 경계."""
    cut_ids: list[int] = []
    pos = 0
    while n - pos > max_len:
        window = [i for i, b in enumerate(boundaries) if pos < b.start <= pos + max_len]
        pick = next(
            (
                cands[-1]
                for kinds in GREEDY_LADDER
                if (cands := [i for i in window if boundaries[i].kind in kinds])
            ),
            None,
        )
        if pick is None:
            raise RuntimeError("greedy 분할 실패 -- 강제 절단점이 부족하다 (버그)")
        cut_ids.append(pick)
        pos = boundaries[pick].end
    return cut_ids


def _short_penalty(length: int, min_len: int) -> float:
    if min_len <= 0 or length >= min_len:
        return 0.0
    deficit = 1.0 - length / min_len
    return SHORT_PENALTY * deficit * deficit