"""pdf2zh-next-mcp: MCP server for PDF translation with full-context LLM translation.

2-pass 방식으로 동작:
  Pass 1 (extract_segments): pdf2zh-next로 레이아웃 분석 + 세그먼트 추출
  LLM이 전체 세그먼트를 한 번에 번역
  Pass 2 (assemble_translated): 번역 결과를 주입하여 최종 PDF 생성
"""

import json
import os
import platform
import shlex
import shutil
import subprocess
import tempfile
import uuid
from pathlib import Path

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("pdf2zh-next-mcp")

SCRIPTS_DIR = Path(__file__).parent / "scripts"
COLLECT_SCRIPT = SCRIPTS_DIR / "collect.py"
LOOKUP_SCRIPT = SCRIPTS_DIR / "lookup.py"
SESSIONS_DIR = Path(tempfile.gettempdir()) / "pdf2zh-mcp-sessions"


def _find_pdf2zh_binary() -> str | None:
    """크로스 플랫폼 pdf2zh_next 바이너리 탐색. 못 찾으면 None 반환."""
    system = platform.system()
    home = Path.home()

    # 1. uv tool 설치 경로 (플랫폼별)
    if system == "Windows":
        appdata = os.environ.get("APPDATA", str(home / "AppData" / "Roaming"))
        candidates = [
            Path(appdata) / "uv" / "data" / "tools" / "pdf2zh-next" / "Scripts" / "pdf2zh_next.exe",
        ]
    else:  # macOS, Linux
        candidates = [
            home / ".local" / "share" / "uv" / "tools" / "pdf2zh-next" / "bin" / "pdf2zh_next",
        ]

    for path in candidates:
        if path.exists():
            return str(path)

    # 2. PATH에서 탐색 (shutil.which는 크로스 플랫폼)
    found = shutil.which("pdf2zh_next")
    if found:
        return found

    return None


def _find_python_binary() -> str | None:
    """크로스 플랫폼 Python 바이너리 탐색. 못 찾으면 None 반환.

    pdf2zh-next의 가상환경 내 Python을 우선 사용한다.
    CLITranslator가 이 Python으로 collect.py/lookup.py를 실행하므로
    시스템 Python과 독립적이어야 한다.
    """
    system = platform.system()
    home = Path.home()

    if system == "Windows":
        appdata = os.environ.get("APPDATA", str(home / "AppData" / "Roaming"))
        candidates = [
            Path(appdata) / "uv" / "data" / "tools" / "pdf2zh-next" / "Scripts" / "python.exe",
        ]
    else:
        candidates = [
            home / ".local" / "share" / "uv" / "tools" / "pdf2zh-next" / "bin" / "python3",
        ]

    for path in candidates:
        if path.exists():
            return str(path)

    # fallback: 시스템 Python
    for name in ("python3", "python"):
        found = shutil.which(name)
        if found:
            return found

    return None


# Lazy 초기화 + 캐싱
_pdf2zh_bin: str | None = None
_python_bin: str | None = None


def _get_pdf2zh_bin() -> str | None:
    global _pdf2zh_bin
    if _pdf2zh_bin is None:
        _pdf2zh_bin = _find_pdf2zh_binary()
    return _pdf2zh_bin


def _get_python_bin() -> str | None:
    global _python_bin
    if _python_bin is None:
        _python_bin = _find_python_binary()
    return _python_bin


def _check_binaries() -> str | None:
    """바이너리 확인. 문제 있으면 에러 JSON 문자열 반환, 정상이면 None."""
    if _get_pdf2zh_bin() is None:
        return json.dumps(
            {
                "status": "error",
                "message": (
                    "pdf2zh-next가 설치되어 있지 않습니다.\n"
                    "다음 명령어로 설치해주세요: uv tool install pdf2zh-next\n"
                    "설치 확인: pdf2zh_next --version"
                ),
            },
            ensure_ascii=False,
        )
    if _get_python_bin() is None:
        return json.dumps(
            {
                "status": "error",
                "message": "Python을 찾을 수 없습니다. Python 3.12 이상이 설치되어 있는지 확인해주세요.",
            },
            ensure_ascii=False,
        )
    return None


def _get_session_dir(session_id: str) -> Path:
    return SESSIONS_DIR / session_id


def _build_cli_command(script: Path, *args: Path) -> str:
    """shlex.quote로 경로를 감싸서 공백/백슬래시에 안전한 CLI 커맨드 문자열을 생성한다."""
    parts = [shlex.quote(_get_python_bin()), shlex.quote(str(script))]
    parts.extend(shlex.quote(str(a)) for a in args)
    return " ".join(parts)


def _run_pdf2zh(
    file: str,
    lang_in: str,
    lang_out: str,
    cli_command: str,
    output_dir: Path,
    pool_max_workers: int = 0,
    debug: bool = False,
) -> subprocess.CompletedProcess:
    """pdf2zh_next CLI를 실행한다."""
    cmd = [
        _get_pdf2zh_bin(),
        "--clitranslator",
        "--clitranslator-command", cli_command,
        "--lang-in", lang_in,
        "--lang-out", lang_out,
        "--output", str(output_dir),
        "--disable-config-auto-save",
        "--ignore-cache",
    ]
    if pool_max_workers > 0:
        cmd.extend(["--pool-max-workers", str(pool_max_workers)])
    if debug:
        cmd.append("--debug")
    cmd.append(file)

    return subprocess.run(
        cmd,
        capture_output=True,
        encoding="utf-8",
        errors="replace",
        timeout=600,  # 10분 타임아웃
    )


def _unique_path(path: Path) -> Path:
    """파일이 이미 존재하면 (1), (2) 등 번호를 붙여 고유한 경로를 반환한다.

    예: paper-mono.pdf → paper-mono (1).pdf → paper-mono (2).pdf → ...
    """
    if not path.exists():
        return path
    stem = path.stem
    suffix = path.suffix
    parent = path.parent
    counter = 1
    while True:
        new_path = parent / f"{stem} ({counter}){suffix}"
        if not new_path.exists():
            return new_path
        counter += 1


# BabelDOC 디버그 작업 디렉토리
_BABELDOC_WORKING_DIR = Path.home() / ".cache" / "babeldoc" / "working"


def _extract_formula_hints(file_stem: str) -> dict[str, dict[str, str]]:
    """translate_tracking.json에서 수식 플레이스홀더 → 원본 유니코드 매핑을 추출한다.

    Returns:
        {segment_text: {placeholder: formula_chars}} 형태의 딕셔너리.
        예: {"some text {v1} here": {"{v1}": "[N+1, d]"}}
    """
    tracking_file = _BABELDOC_WORKING_DIR / file_stem / "translate_tracking.json"
    if not tracking_file.exists():
        return {}

    try:
        data = json.loads(tracking_file.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return {}

    # 세그먼트(input) → {placeholder: formula_chars} 매핑 구축
    formula_map: dict[str, dict[str, str]] = {}
    for page in data.get("page", []):
        if page is None:
            continue
        for para in page.get("paragraph", []):
            if para is None:
                continue
            placeholders = para.get("placeholders", [])
            if not placeholders:
                continue
            input_text = para.get("input", "")
            if not input_text:
                continue
            hints = {}
            for ph in placeholders:
                placeholder = ph.get("placeholder", "")
                formula_chars = ph.get("formula_chars", "")
                if placeholder and formula_chars:
                    hints[placeholder] = formula_chars
            if hints:
                formula_map[input_text] = hints

    return formula_map


def _build_formula_hint_text(
    segments: list[str], formula_map: dict[str, dict[str, str]]
) -> str:
    """세그먼트 목록과 수식 매핑으로부터 번역 힌트 텍스트를 생성한다."""
    if not formula_map:
        return ""

    hints = []
    for seg in segments:
        seg_hints = formula_map.get(seg, {})
        if seg_hints:
            for placeholder, formula in seg_hints.items():
                hints.append(f"  {placeholder} = {formula}")

    if not hints:
        return ""

    # 중복 제거 (같은 매핑이 여러 세그먼트에 있을 수 있음)
    unique_hints = list(dict.fromkeys(hints))
    return (
        "\n\n수식 참고 (Formula Reference):\n"
        "아래는 각 플레이스홀더가 나타내는 실제 수식/기호입니다. "
        "번역 시 문맥을 이해하는 데 참고하세요. 플레이스홀더 자체는 그대로 유지해야 합니다.\n"
        + "\n".join(unique_hints)
    )


def _cleanup_debug_files(file_stem: str) -> None:
    """BabelDOC 디버그 작업 디렉토리를 정리한다."""
    debug_dir = _BABELDOC_WORKING_DIR / file_stem
    if debug_dir.exists():
        shutil.rmtree(debug_dir, ignore_errors=True)


def _collect_output_files(
    temp_output_dir: Path, final_dir: Path, file_stem: str, glossary_json: str = ""
) -> dict[str, str]:
    """임시 출력 디렉토리에서 생성된 PDF를 찾아 최종 디렉토리로 이동한다.

    파일명 중복 시 번호를 붙여서 기존 파일을 보존한다.
    """
    output_files: dict[str, str] = {}

    # 기대 패턴으로 찾기
    mono_pdf = temp_output_dir / f"{file_stem}-mono.pdf"
    dual_pdf = temp_output_dir / f"{file_stem}-dual.pdf"

    # fallback: BabelDOC 출력 패턴 탐색
    if not mono_pdf.exists():
        candidates = (
            list(temp_output_dir.glob(f"{file_stem}*mono*.pdf"))
            + list(temp_output_dir.glob(f"{file_stem}*translated*.pdf"))
        )
        if candidates:
            mono_pdf = candidates[0]

    if not dual_pdf.exists():
        candidates = list(temp_output_dir.glob(f"{file_stem}*dual*.pdf"))
        if candidates:
            dual_pdf = candidates[0]

    # 최종 위치로 이동 (번호 붙이기)
    if mono_pdf.exists():
        dest = _unique_path(final_dir / f"{file_stem}-mono.pdf")
        shutil.move(str(mono_pdf), str(dest))
        output_files["mono"] = str(dest)

    if dual_pdf.exists():
        dest = _unique_path(final_dir / f"{file_stem}-dual.pdf")
        shutil.move(str(dual_pdf), str(dest))
        output_files["dual"] = str(dest)

    # 용어집 저장
    if glossary_json:
        try:
            glossary = json.loads(glossary_json)
            dest = _unique_path(final_dir / f"{file_stem}-glossary.json")
            dest.write_text(
                json.dumps(glossary, ensure_ascii=False, indent=2), encoding="utf-8"
            )
            output_files["glossary"] = str(dest)
        except json.JSONDecodeError:
            pass

    return output_files


@mcp.tool()
def extract_segments(
    file: str,
    lang_in: str = "en",
    lang_out: str = "ko",
    mode: str = "basic",
    output_dir: str = "",
) -> str:
    """PDF에서 번역 대상 텍스트 세그먼트를 추출합니다.

    pdf2zh-next로 레이아웃을 분석하고 번역이 필요한 텍스트 세그먼트를 추출합니다.
    반환된 세그먼트를 번역한 후 assemble_translated 툴에 전달해주세요.
    중요: {v0}, {v1} 같은 수식 플레이스홀더는 그대로 유지해주세요.

    Args:
        file: PDF 파일의 절대 경로
        lang_in: 원본 언어 코드 (기본값: en)
        lang_out: 번역 대상 언어 코드 (기본값: ko)
        mode: 번역 모드. "basic"(텍스트만) 또는 "visual"(PDF 첨부 시각 맥락 활용 + 용어집 저장)
        output_dir: 출력 디렉토리 경로. 비어 있으면 입력 PDF와 같은 디렉토리에 저장.
    """
    # 바이너리 확인
    err = _check_binaries()
    if err:
        return err

    file_path = Path(file)
    if not file_path.exists():
        return json.dumps(
            {"status": "error", "message": f"파일을 찾을 수 없습니다: {file}"},
            ensure_ascii=False,
        )

    # 세션 생성
    session_id = uuid.uuid4().hex[:12]
    session_dir = _get_session_dir(session_id)
    session_dir.mkdir(parents=True, exist_ok=True)

    # 출력 디렉토리 결정
    resolved_output_dir = Path(output_dir) if output_dir else file_path.parent
    if output_dir and not resolved_output_dir.exists():
        resolved_output_dir.mkdir(parents=True, exist_ok=True)

    # 세션 메타데이터 저장
    meta = {
        "file": str(file_path.absolute()),
        "lang_in": lang_in,
        "lang_out": lang_out,
        "output_dir": str(resolved_output_dir.absolute()),
    }
    (session_dir / "meta.json").write_text(
        json.dumps(meta, ensure_ascii=False), encoding="utf-8"
    )

    segments_file = session_dir / "segments.json"
    pass1_output = session_dir / "pass1_output"
    pass1_output.mkdir(exist_ok=True)

    # Pass 1: 세그먼트 수집 (--debug로 수식 매핑 정보도 생성)
    cli_command = _build_cli_command(COLLECT_SCRIPT, segments_file)
    result = _run_pdf2zh(
        file=str(file_path.absolute()),
        lang_in=lang_in,
        lang_out=lang_out,
        cli_command=cli_command,
        output_dir=pass1_output,
        pool_max_workers=1,  # 순서 보장
        debug=True,
    )

    if result.returncode != 0:
        return json.dumps(
            {
                "status": "error",
                "message": "pdf2zh-next 실행 실패",
                "stderr": result.stderr[-2000:] if result.stderr else "",
            },
            ensure_ascii=False,
        )

    if not segments_file.exists():
        return json.dumps(
            {"status": "error", "message": "세그먼트가 추출되지 않았습니다."},
            ensure_ascii=False,
        )

    segments = json.loads(segments_file.read_text(encoding="utf-8"))

    # 수식 힌트 추출 (--debug로 생성된 translate_tracking.json에서)
    file_stem = file_path.stem
    formula_map = _extract_formula_hints(file_stem)
    formula_hint_text = _build_formula_hint_text(segments, formula_map)

    # 디버그 파일 및 Pass 1 출력물 정리
    _cleanup_debug_files(file_stem)
    shutil.rmtree(pass1_output, ignore_errors=True)

    # 공통 instruction (두 모드 동일)
    common_instruction = (
        f"위 {len(segments)}개 세그먼트를 [{lang_in}]에서 [{lang_out}]로 번역해주세요.\n\n"
        "번역 전에 반드시:\n"
        "1. 모든 세그먼트를 먼저 읽고 핵심 전문 용어를 파악하세요\n"
        "2. 일관된 용어를 사용하여 번역하세요\n"
        "3. 일부 세그먼트는 페이지 경계에서 문장이 잘려 있을 수 있습니다. "
        "앞뒤 세그먼트를 참고하여 자연스럽게 이어지도록 번역하세요\n\n"
        "규칙:\n"
        "- {v0}, {v1} 등의 수식 플레이스홀더는 절대 변경하지 마세요\n"
        "- 학술 논문에 적합한 격식체를 사용하세요\n"
        "- 널리 알려진 고유명사나 약어는 원문 그대로 유지하세요\n"
    )

    if mode == "visual":
        instruction = (
            common_instruction
            + "\n첨부된 PDF를 참고하여 그림, 표, 수식의 시각적 맥락을 반영해서 번역하세요.\n\n"
            "출력 형식:\n"
            "1. 먼저 용어집을 JSON 객체로 출력하세요: "
            '{"term": "번역", ...}\n'
            f"2. 그 다음 번역된 {len(segments)}개 세그먼트를 JSON 배열로 출력하세요\n"
            "3. assemble_translated 툴을 호출하여 session_id, "
            "translated_segments_json, glossary_json을 전달해주세요."
        )
    else:
        instruction = (
            common_instruction
            + f"\n번역된 {len(segments)}개 세그먼트를 JSON 배열로 출력 후 "
            "assemble_translated 툴을 호출하여 "
            "session_id와 translated_segments_json을 전달해주세요."
        )

    # 수식 힌트를 모든 모드에 추가
    if formula_hint_text:
        instruction += formula_hint_text

    return json.dumps(
        {
            "status": "success",
            "session_id": session_id,
            "segment_count": len(segments),
            "segments": segments,
            "mode": mode,
            "instruction": instruction,
        },
        ensure_ascii=False,
    )


@mcp.tool()
def assemble_translated(
    session_id: str,
    translated_segments_json: str,
    glossary_json: str = "",
) -> str:
    """번역된 세그먼트로 최종 PDF를 생성합니다.

    extract_segments에서 받은 session_id와 번역된 세그먼트 배열을 전달하면,
    pdf2zh-next가 번역을 주입하여 최종 PDF를 조립합니다.

    Args:
        session_id: extract_segments에서 반환된 세션 ID
        translated_segments_json: 번역된 세그먼트의 JSON 배열 문자열
        glossary_json: (선택) 용어집 JSON 문자열. visual 모드에서 생성된 용어집.
    """
    # 바이너리 확인
    err = _check_binaries()
    if err:
        return err

    session_dir = _get_session_dir(session_id)
    meta_file = session_dir / "meta.json"

    if not meta_file.exists():
        return json.dumps(
            {"status": "error", "message": f"세션을 찾을 수 없습니다: {session_id}"},
            ensure_ascii=False,
        )

    # 메타데이터 로드
    meta = json.loads(meta_file.read_text(encoding="utf-8"))
    file = meta["file"]
    lang_in = meta["lang_in"]
    lang_out = meta["lang_out"]
    final_dir = Path(meta.get("output_dir", str(Path(file).parent)))

    # 번역 파싱
    try:
        translated = json.loads(translated_segments_json)
    except json.JSONDecodeError as e:
        return json.dumps(
            {"status": "error", "message": f"JSON 파싱 실패: {e}"},
            ensure_ascii=False,
        )

    if not isinstance(translated, list):
        return json.dumps(
            {"status": "error", "message": "번역 결과는 JSON 배열이어야 합니다."},
            ensure_ascii=False,
        )

    # 원본 세그먼트 로드
    segments_file = session_dir / "segments.json"
    segments = json.loads(segments_file.read_text(encoding="utf-8"))

    if len(translated) != len(segments):
        return json.dumps(
            {
                "status": "error",
                "message": f"세그먼트 수 불일치: 원본 {len(segments)}개, 번역 {len(translated)}개",
            },
            ensure_ascii=False,
        )

    # 번역 딕셔너리 생성 + 저장
    translations = dict(zip(segments, translated))
    translations_file = session_dir / "translations.json"
    translations_file.write_text(
        json.dumps(translations, ensure_ascii=False), encoding="utf-8"
    )

    # 임시 출력 디렉토리 (pdf2zh-next가 여기에 생성 → 최종 위치로 이동)
    temp_output_dir = session_dir / "pass2_output"
    temp_output_dir.mkdir(exist_ok=True)

    # Pass 2: PDF 조립
    cli_command = _build_cli_command(LOOKUP_SCRIPT, translations_file)
    result = _run_pdf2zh(
        file=file,
        lang_in=lang_in,
        lang_out=lang_out,
        cli_command=cli_command,
        output_dir=temp_output_dir,
    )

    if result.returncode != 0:
        return json.dumps(
            {
                "status": "error",
                "message": "pdf2zh-next PDF 조립 실패",
                "stderr": result.stderr[-2000:] if result.stderr else "",
            },
            ensure_ascii=False,
        )

    # 생성된 PDF를 최종 디렉토리로 이동 (번호 붙이기)
    file_stem = Path(file).stem
    output_files = _collect_output_files(
        temp_output_dir, final_dir, file_stem, glossary_json
    )

    if not output_files or ("mono" not in output_files and "dual" not in output_files):
        return json.dumps(
            {
                "status": "error",
                "message": "PDF가 생성되지 않았습니다. pdf2zh-next 출력을 확인해주세요.",
                "stderr": result.stderr[-2000:] if result.stderr else "",
            },
            ensure_ascii=False,
        )

    # 세션 정리
    shutil.rmtree(session_dir, ignore_errors=True)

    return json.dumps(
        {
            "status": "success",
            "message": f"번역 완료! {len(segments)}개 세그먼트 처리됨.",
            "output": output_files,
        },
        ensure_ascii=False,
    )


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
