"""v24 파이프라인의 `tool_calling_judge` 스텝 + corrections 공용 헬퍼.

HF Space 구현(`huggingface/pipelines.py:1349` 근방, `huggingface/pipelines_v24.py`)을
포팅하되 한 가지는 의도적으로 바꿨다. HF 는 "Upstage chat-completions 엔드포인트가
`tools=[...]` 를 네이티브로 지원하지 않는다"는 이유로 2-라운드 JSON 스키마 흉내를
썼다(라운드1: PMI 조회 계획을 JSON 으로 뱉게 하고, 엔진이 로컬에서 조회해 라운드2에
결과를 다시 먹인다). 이 마이그레이션에서 `providers/base.py` 에 실제 `tools`/
`tool_choice` 를 추가했으므로, 여기서는 진짜 tool-calling 루프(`max_tool_iterations`
회 반복 -> 도구 호출 -> 로컬 dispatch -> 결과를 `role: tool` 메시지로 피드백)를 쓴다.
판정 로직(PMI/빈도 조회로 keep/revert 결정)과 프롬프트 조립(원문/교정후/적용된 교정
목록)은 HF 구현을 그대로 따른다.

PMI 조회 캐시는 HF 처럼 파일 경로를 키로 쓰지만, HF 의 모듈 전역 `_ACTIVE_PMI`
포인터는 갖지 않는다 -- 동시 처리 중인 여러 요청이 서로 다른 활성 경로를 덮어쓸 수
있는 경합 지점이었다. 대신 `dispatch_tool` 이 매 호출마다 `pmi_path` 를 명시적으로
받는다.
"""

from __future__ import annotations

import gzip
import json
import logging
import re
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from solar_eval.providers.base import BaseProvider

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# PMI lookup -- 경로별 캐시. 로드 실패는 조용히 넘기지 않고 즉시 raise 한다
# (§5-C 가 정한 원칙: 모르는/없는 리소스는 명확히 실패시킨다).
# ---------------------------------------------------------------------------

_BIGRAM_CACHE: dict[Path, dict[str, dict[str, tuple[float, int]]]] = {}
_UNIGRAM_CACHE: dict[Path, dict[str, int]] = {}


def resolve_pmi_path(project_dir: Path | str | None, pmi_lookup: str) -> Path:
    """파이프라인 config 의 `pmi_lookup` 값을 프로젝트 디렉토리 기준 절대경로로 푼다.

    Raises:
        ValueError: `project_dir` 이 파이프라인에 주입되지 않았는데 스텝이
            `pmi_lookup` 을 요구하는 경우 (조용한 스킵 금지).
        FileNotFoundError: 풀어낸 경로에 파일이 없는 경우.
    """
    if project_dir is None:
        raise ValueError(
            f"tool_calling_judge 스텝이 pmi_lookup={pmi_lookup!r} 을 요구하지만 "
            "이 파이프라인에는 project_dir 이 주입되지 않았다 "
            "(registry.create_pipeline 호출부의 project_dir 인자를 확인할 것)."
        )
    path = (Path(project_dir) / pmi_lookup).resolve()
    if not path.exists():
        raise FileNotFoundError(
            f"PMI lookup 파일을 찾을 수 없다: {path} "
            f"(project_dir={project_dir}, pmi_lookup={pmi_lookup!r})"
        )
    return path


def load_pmi_lookup(path: Path) -> None:
    """`build_pmi_top_for_demo.py` 산출물(gzip TSV)에서 PMI/unigram 테이블을 로드한다.

    파일 포맷::
        # SECTION: bigrams
        left<TAB>right<TAB>count<TAB>pmi
        ...
        # SECTION: unigrams
        token<TAB>count
        ...

    같은 경로로 두 번째 호출하면 캐시를 재사용한다(idempotent).
    """
    if path in _BIGRAM_CACHE:
        return
    if not path.exists():
        raise FileNotFoundError(f"PMI lookup 파일을 찾을 수 없다: {path}")

    bigram: dict[str, dict[str, tuple[float, int]]] = {}
    unigram: dict[str, int] = {}
    section: str | None = None
    with gzip.open(path, "rt", encoding="utf-8") as f:
        for line in f:
            line = line.rstrip("\n")
            if not line:
                continue
            if line.startswith("# SECTION: bigrams"):
                section = "bigrams"
                continue
            if line.startswith("# SECTION: unigrams"):
                section = "unigrams"
                continue
            if line.startswith("#"):
                continue
            parts = line.split("\t")
            if section == "bigrams" and len(parts) >= 4:
                try:
                    left, right, count, pmi = parts[0], parts[1], int(parts[2]), float(parts[3])
                except ValueError:
                    continue
                bigram.setdefault(left, {})[right] = (pmi, count)
            elif section == "unigrams" and len(parts) >= 2:
                try:
                    unigram[parts[0]] = int(parts[1])
                except ValueError:
                    continue

    _BIGRAM_CACHE[path] = bigram
    _UNIGRAM_CACHE[path] = unigram
    logger.info(
        "[judge.pmi] loaded %s: %d unigrams, %d bigram-left-tokens",
        path.name,
        len(unigram),
        len(bigram),
    )


def lookup_pmi(path: Path, left: str, right: str) -> dict[str, Any]:
    if not left or not right:
        return {"pmi": None, "count": 0}
    inner = _BIGRAM_CACHE.get(path, {}).get(left, {})
    if right not in inner:
        return {"pmi": None, "count": 0}
    pmi, count = inner[right]
    return {"pmi": round(pmi, 4), "count": count}


def lookup_pmi_joined(path: Path, joined: str) -> dict[str, Any]:
    """붙여쓴 문자열을 모든 지점에서 좌/우로 나눠 PMI 가 가장 높은 조합을 찾는다."""
    if not joined or len(joined) < 4:
        return {"pmi": None, "count": 0, "left": None, "right": None}
    bigram = _BIGRAM_CACHE.get(path, {})
    best: tuple[float, int, str, str] | None = None
    for i in range(2, len(joined) - 1):
        left, right = joined[:i], joined[i:]
        inner = bigram.get(left, {})
        if right in inner:
            pmi, count = inner[right]
            if best is None or pmi > best[0]:
                best = (pmi, count, left, right)
    if best is None:
        return {"pmi": None, "count": 0, "left": None, "right": None}
    pmi, count, left, right = best
    return {"pmi": round(pmi, 4), "count": count, "left": left, "right": right}


def lookup_unigram_freq(path: Path, token: str) -> dict[str, Any]:
    if not token:
        return {"count": 0}
    return {"count": _UNIGRAM_CACHE.get(path, {}).get(token, 0)}


def dispatch_tool(tool: str, args: dict[str, Any], *, pmi_path: Path) -> dict[str, Any]:
    """도구 호출 하나를 실행한다.

    여기서는 예외를 삼킨다 -- 대화 상대는 LLM 이고, 잘못된 tool 인자는 raise 대상인
    "우리 설정 실수"가 아니라 모델이 만들어낸 신뢰할 수 없는 외부 입력이다.
    `{"ok": False, ...}` 를 돌려주면 모델이 다음 턴에서 스스로 고칠 기회를 얻는다.
    """
    if not isinstance(args, dict):
        return {"ok": False, "error": "args must be an object"}
    try:
        if tool == "lookup_pmi":
            return {
                "ok": True,
                "result": lookup_pmi(
                    pmi_path, str(args.get("left", "")), str(args.get("right", ""))
                ),
            }
        if tool == "lookup_pmi_joined":
            return {"ok": True, "result": lookup_pmi_joined(pmi_path, str(args.get("joined", "")))}
        if tool == "lookup_unigram_freq":
            return {"ok": True, "result": lookup_unigram_freq(pmi_path, str(args.get("token", "")))}
        return {"ok": False, "error": f"unknown tool: {tool}"}
    except Exception as exc:  # noqa: BLE001 -- 모델 입력발 예외는 도구 결과로 흡수
        logger.warning("[judge.dispatch_tool] %s raised: %s", tool, exc)
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}


JUDGE_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "lookup_pmi",
            "description": (
                "두 어절(왼쪽/오른쪽)의 PMI(점별 상호정보량)와 코퍼스 등장 빈도를 조회한다. "
                "붙여쓰기/띄어쓰기 교정이 타당한지 판단할 때 참고한다."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "left": {"type": "string", "description": "왼쪽 어절"},
                    "right": {"type": "string", "description": "오른쪽 어절"},
                },
                "required": ["left", "right"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "lookup_pmi_joined",
            "description": "붙여쓴 문자열을 나눠볼 수 있는 모든 지점에서 PMI 를 조회해 가장 높은 조합을 반환한다.",
            "parameters": {
                "type": "object",
                "properties": {"joined": {"type": "string", "description": "붙여쓴 문자열"}},
                "required": ["joined"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "lookup_unigram_freq",
            "description": "단일 토큰의 코퍼스 등장 빈도를 조회한다.",
            "parameters": {
                "type": "object",
                "properties": {"token": {"type": "string", "description": "조회할 토큰"}},
                "required": ["token"],
            },
        },
    },
]

_DECISIONS_SCHEMA: dict[str, Any] = {
    "type": "json_schema",
    "json_schema": {
        "name": "judge_decisions",
        "strict": True,
        "schema": {
            "type": "object",
            "properties": {
                "decisions": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "correction_index": {"type": "integer"},
                            "verdict": {"type": "string", "enum": ["keep", "revert"]},
                            "reason": {"type": "string"},
                        },
                        "required": ["correction_index", "verdict", "reason"],
                        "additionalProperties": False,
                    },
                }
            },
            "required": ["decisions"],
            "additionalProperties": False,
        },
    },
}


# ---------------------------------------------------------------------------
# corrections JSON 파싱/적용 -- self-consistency(multi_step.py)와
# apply_judge_decisions(rules.py) 양쪽이 공유한다.
# ---------------------------------------------------------------------------


def _parse_json_robust(content: str) -> dict[str, Any] | None:
    """임의의 LLM 응답에서 JSON object 를 파싱한다. raw -> ```json``` 블록 -> 마지막 중괄호 균형 순."""
    if not content:
        return None
    try:
        data = json.loads(content)
        if isinstance(data, dict):
            return data
    except (json.JSONDecodeError, AttributeError):
        pass

    for m in reversed(list(re.finditer(r"```(?:json)?\s*(.*?)```", content, re.DOTALL))):
        try:
            data = json.loads(m.group(1).strip())
            if isinstance(data, dict):
                return data
        except (json.JSONDecodeError, AttributeError):
            continue

    last_open = content.rfind("{")
    while last_open != -1:
        depth = 0
        for i in range(last_open, len(content)):
            if content[i] == "{":
                depth += 1
            elif content[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(content[last_open : i + 1])
                        if isinstance(data, dict):
                            return data
                    except (json.JSONDecodeError, AttributeError):
                        pass
                    break
        last_open = content.rfind("{", 0, last_open)
    return None


def extract_corrections(raw_json: str) -> list[dict[str, Any]]:
    """basic_correction 류 스텝의 raw JSON 응답에서 `corrections` 배열을 뽑는다."""
    data = _parse_json_robust(raw_json) or {}
    arr = data.get("corrections", [])
    if not isinstance(arr, list):
        return []
    return [
        c
        for c in arr
        if isinstance(c, dict)
        and isinstance(c.get("original"), str)
        and isinstance(c.get("corrected"), str)
    ]


def apply_corrections_to_text(text: str, corrections: list[dict[str, Any]]) -> str:
    """(original -> corrected) 치환을 왼쪽부터, 긴 것부터, 겹치지 않게 적용한다."""
    if not corrections:
        return text
    seen: set[tuple[str, str]] = set()
    clean: list[dict[str, str]] = []
    for c in corrections:
        orig, new = c.get("original"), c.get("corrected")
        if not isinstance(orig, str) or not isinstance(new, str):
            continue
        if not orig or orig == new:
            continue
        key = (orig, new)
        if key in seen:
            continue
        seen.add(key)
        clean.append({"original": orig, "corrected": new})
    clean.sort(key=lambda c: -len(c["original"]))

    edits: list[tuple[int, int, str]] = []
    for c in clean:
        idx = text.find(c["original"])
        while idx != -1:
            end = idx + len(c["original"])
            if not any(not (end <= s or idx >= e) for s, e, _ in edits):
                edits.append((idx, end, c["corrected"]))
                break
            idx = text.find(c["original"], idx + 1)
    edits.sort(key=lambda e: e[0], reverse=True)
    for start, end, corrected in edits:
        text = text[:start] + corrected + text[end:]
    return text


def intersect_corrections(runs: list[list[dict[str, Any]]], threshold: int) -> list[dict[str, Any]]:
    """여러 run 의 corrections 목록 중 >= threshold 개의 run 에 등장한 것만 남긴다."""
    counts: dict[tuple[str, str], int] = {}
    metas: dict[tuple[str, str], dict[str, Any]] = {}
    for run_corrs in runs:
        seen_in_run: set[tuple[str, str]] = set()
        for c in run_corrs:
            key = (c["original"], c["corrected"])
            if key in seen_in_run:
                continue
            seen_in_run.add(key)
            counts[key] = counts.get(key, 0) + 1
            metas.setdefault(
                key,
                {
                    "original": c["original"],
                    "corrected": c["corrected"],
                    "type": c.get("type", ""),
                    "reason": c.get("reason", ""),
                },
            )
    return [metas[k] for k, n in counts.items() if n >= threshold]


def _accumulate_usage(total: dict[str, int], usage: dict[str, Any]) -> None:
    for k in total:
        total[k] += usage.get(k, 0)


async def run_self_consistency(
    bulk: str,
    user_content: str,
    step: dict[str, Any],
    *,
    provider: BaseProvider,
    step_model: str,
    step_re: str | None,
    step_max_tokens: int,
    step_response_format: dict[str, Any] | None,
    step_prompt: str,
) -> tuple[str, dict[str, int], list[dict[str, Any]], list[dict[str, Any]]]:
    """basic_correction 의 self-consistency: bulk 를 N번 반복 실행하고 교집합만 남긴다.

    HF 구현(`huggingface/pipelines_v24.py::run_self_consistency_for_bulk`) 포팅.
    각 run 의 결과 텍스트 자체는 버린다 -- 최종 텍스트는 항상 교집합 corrections 를
    원본 bulk 위에 다시 적용해서(`apply_corrections_to_text`) 만든다. 이게
    self-consistency 의 핵심이다: 각 run 이 개별적으로 만든 "output" 문자열이
    아니라, 여러 run 이 공통으로 동의한 교정만 신뢰한다.

    `user_content` 는 호출자(`MultiStepPipeline`)가 `user_template`/
    `include_original` 규칙에 따라 이미 조립해 넘긴다 -- 이 함수는 프롬프트 조립
    방식을 모른다(순수하게 반복 호출 + corrections 파싱/교집합/적용만 담당).

    Returns:
        (최종 bulk 텍스트, 이 bulk 가 쓴 usage 합계, run별 {temperature, corrections}
        리스트(교집합 이전 -- 가시성용), 교집합을 통과한 corrections 리스트).
    """
    repeat = int(step.get("repeat", 1) or 1)
    temps = list(step.get("repeat_temperatures") or [])
    while len(temps) < repeat:
        temps.append(0.0 if not temps else temps[-1])
    threshold = int(step.get("intersection_threshold", repeat) or repeat)

    usage_total = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    runs: list[dict[str, Any]] = []

    for i in range(repeat):
        msgs: list[dict[str, str]] = []
        if step_prompt:
            msgs.append({"role": "system", "content": step_prompt})
        msgs.append({"role": "user", "content": user_content})
        result = await provider.chat_completion(
            messages=msgs,
            model=step_model,
            temperature=temps[i],
            max_tokens=step_max_tokens,
            reasoning_effort=step_re,
            response_format=step_response_format,
        )
        corrections = extract_corrections(result.get("content", ""))
        runs.append({"temperature": temps[i], "corrections": corrections})
        _accumulate_usage(usage_total, result.get("usage", {}))

    survivors = intersect_corrections([r["corrections"] for r in runs], threshold)
    final_text = apply_corrections_to_text(bulk, survivors)
    return final_text, usage_total, runs, survivors


# ---------------------------------------------------------------------------
# tool_calling_judge 스텝 본체
# ---------------------------------------------------------------------------


async def run_tool_calling_judge(
    step: dict[str, Any],
    *,
    text: str,
    original_text: str,
    corrections: list[dict[str, Any]],
    provider: BaseProvider,
    model: str,
    reasoning_effort: str | None,
    max_tokens: int,
    prompts: dict[str, str],
    project_dir: Path | str | None,
) -> dict[str, Any]:
    """corrections 목록을 판정한다 (실제 tool-calling 루프, 최대 `max_tool_iterations` 회).

    Returns:
        {"decisions": [...], "tool_call_log": [...], "usage": {...}}.
        `decisions` 항목: {"correction_index", "verdict": "keep"|"revert", "reason"}.
        text 는 이 함수가 직접 바꾸지 않는다 -- 되돌리기는
        `pipelines/rules.py::apply_judge_decisions` 가 별도 rule 스텝에서 한다.
    """
    step_name = step.get("name", "tool_calling_judge")
    empty_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    if not corrections:
        return {"decisions": [], "tool_call_log": [], "usage": dict(empty_usage)}

    pmi_path: Path | None = None
    pmi_lookup = step.get("pmi_lookup")
    if pmi_lookup:
        pmi_path = resolve_pmi_path(project_dir, pmi_lookup)
        load_pmi_lookup(pmi_path)

    judge_prompt = prompts.get(step.get("prompt_key", "judge"), "")
    max_iters = int(step.get("max_tool_iterations", 4))

    corrs_for_prompt = [
        {
            "i": i,
            "original": c.get("original", ""),
            "corrected": c.get("corrected", ""),
            "type": c.get("type", ""),
        }
        for i, c in enumerate(corrections)
    ]
    user_content = (
        f"<원문>\n{original_text}\n</원문>\n\n"
        f"<교정후>\n{text}\n</교정후>\n\n"
        f"<적용된_교정_목록_JSON>\n{json.dumps(corrs_for_prompt, ensure_ascii=False)}\n"
        f"</적용된_교정_목록_JSON>\n\n"
        "각 교정에 대해 필요하면 PMI/빈도 조회 도구를 사용한 뒤, 모든 correction_index 에 "
        "대한 keep/revert 판정을 근거와 함께 최종적으로 출력하라."
    )
    messages: list[dict[str, Any]] = [
        {"role": "system", "content": judge_prompt},
        {"role": "user", "content": user_content},
    ]

    usage_total = dict(empty_usage)
    tool_call_log: list[dict[str, Any]] = []
    tools = JUDGE_TOOLS if pmi_path is not None else None

    # pmi_lookup 이 없는 스텝(도구가 없는 judge)은 루프를 아예 돌지 않는다 --
    # tools=None 상태로 chat_completion 을 불러봐야 어차피 tool_calls 가 비어
    # break 될 뿐이라, 그 한 번의 호출은 최종 결정 호출과 중복되는 낭비다.
    for _ in range(max_iters if tools else 0):
        result = await provider.chat_completion(
            messages=messages,
            model=model,
            temperature=0.0,
            max_tokens=max_tokens,
            reasoning_effort=reasoning_effort,
            tools=tools,
            tool_choice="auto",
        )
        _accumulate_usage(usage_total, result.get("usage", {}))
        tool_calls = result.get("tool_calls") or []
        if not tool_calls:
            break

        messages.append(
            {"role": "assistant", "content": result.get("content") or "", "tool_calls": tool_calls}
        )
        for tc in tool_calls:
            fn = tc.get("function", {}) if isinstance(tc, dict) else {}
            name = fn.get("name", "")
            try:
                args = json.loads(fn.get("arguments") or "{}")
            except json.JSONDecodeError:
                args = {}
            assert pmi_path is not None  # tools 는 pmi_path 있을 때만 노출된다
            tool_result = dispatch_tool(name, args, pmi_path=pmi_path)
            tool_call_log.append({"tool": name, "args": args, "result": tool_result})
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tc.get("id", ""),
                    "content": json.dumps(tool_result, ensure_ascii=False),
                }
            )

    messages.append(
        {
            "role": "user",
            "content": (
                "이제 도구 호출을 멈추고 모든 correction_index 에 대한 최종 keep/revert 판정을 "
                "JSON 으로만 출력하라. decisions 의 길이는 적용된 교정 목록의 길이와 같아야 한다."
            ),
        }
    )
    final_result = await provider.chat_completion(
        messages=messages,
        model=model,
        temperature=0.0,
        max_tokens=max_tokens,
        reasoning_effort=reasoning_effort,
        response_format=_DECISIONS_SCHEMA,
    )
    _accumulate_usage(usage_total, final_result.get("usage", {}))

    parsed = _parse_json_robust(final_result.get("content", "")) or {}
    raw_decisions = parsed.get("decisions", []) if isinstance(parsed, dict) else []
    decisions = [
        d
        for d in raw_decisions
        if isinstance(d, dict)
        and isinstance(d.get("correction_index"), int)
        and d.get("verdict") in ("keep", "revert")
    ]
    if len(decisions) != len(raw_decisions):
        logger.warning(
            "[%s] judge 가 %d개 중 %d개만 유효한 decision 을 반환했다",
            step_name,
            len(raw_decisions),
            len(decisions),
        )

    return {"decisions": decisions, "tool_call_log": tool_call_log, "usage": usage_total}