#!/usr/bin/env python3
"""Use local LM Studio OpenAI-compatible API to polish markdown chapters."""

from __future__ import annotations

import argparse
import json
import re
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path


DEFAULT_BASE_URL = "http://localhost:1234/v1"
DEFAULT_MODEL = "gemma4-12b-qat-uncensored-hauhaucs-balanced"


def read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8-sig")


def write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text.rstrip() + "\n", encoding="utf-8")


def text_stats(text: str) -> dict[str, int]:
    return {
        "chars": len(text),
        "nonspace": len(re.findall(r"\S", text)),
        "zh": len(re.findall(r"[\u4e00-\u9fff]", text)),
    }


def request_json(url: str, payload: dict | None = None, timeout: int = 600) -> dict:
    data = None if payload is None else json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(url, data=data)
    if payload is not None:
        req.add_header("Content-Type", "application/json")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8")
    return json.loads(raw)


def list_models(base_url: str) -> None:
    payload = request_json(f"{base_url.rstrip('/')}/models", timeout=10)
    for item in payload.get("data", []):
        print(item.get("id", ""))


def polish_text(
    text: str,
    system_prompt: str,
    *,
    base_url: str,
    model: str,
    temperature: float,
    top_p: float | None,
    presence_penalty: float | None,
    frequency_penalty: float | None,
    repeat_penalty: float | None,
    max_tokens: int | None,
    timeout: int,
) -> str:
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": text},
        ],
        "temperature": temperature,
    }
    if top_p is not None:
        payload["top_p"] = top_p
    if presence_penalty is not None:
        payload["presence_penalty"] = presence_penalty
    if frequency_penalty is not None:
        payload["frequency_penalty"] = frequency_penalty
    if repeat_penalty is not None:
        payload["repeat_penalty"] = repeat_penalty
    if max_tokens is not None:
        payload["max_tokens"] = max_tokens

    result = request_json(f"{base_url.rstrip('/')}/chat/completions", payload, timeout=timeout)
    try:
        return result["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError, TypeError) as exc:
        raise RuntimeError(f"Unexpected LM Studio response: {result}") from exc


def collect_inputs(input_path: Path) -> list[Path]:
    if input_path.is_file():
        return [input_path]
    if input_path.is_dir():
        return sorted(input_path.glob("*.md"))
    raise FileNotFoundError(f"Input not found: {input_path}")


def output_path_for(src: Path, input_root: Path, out: Path) -> Path:
    if input_root.is_file() and out.suffix.lower() == ".md":
        return out
    return out / src.name


def main() -> int:
    parser = argparse.ArgumentParser(description="Polish chapters with local LM Studio.")
    parser.add_argument("--input", "-i", type=Path, help="Input markdown file or directory.")
    parser.add_argument("--out", "-o", type=Path, default=Path("中间文件/润色输出"), help="Output file or directory.")
    parser.add_argument("--system", type=Path, default=Path("脚本/lmstudio-润色-system.md"), help="System prompt file.")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", type=float, default=None)
    parser.add_argument("--presence-penalty", type=float, default=None)
    parser.add_argument("--frequency-penalty", type=float, default=None)
    parser.add_argument("--repeat-penalty", type=float, default=None, help="LM Studio/llama.cpp repeat penalty when supported.")
    parser.add_argument("--max-tokens", type=int, default=None)
    parser.add_argument("--timeout", type=int, default=1800)
    parser.add_argument("--list-models", action="store_true", help="List LM Studio models and exit.")
    parser.add_argument("--dry-run", action="store_true", help="Show planned tasks without calling LM Studio.")
    args = parser.parse_args()

    if args.list_models:
        list_models(args.base_url)
        return 0
    if args.input is None:
        parser.error("--input is required unless --list-models is used")

    inputs = collect_inputs(args.input)
    if not inputs:
        raise FileNotFoundError(f"No markdown files under: {args.input}")
    system_prompt = read_text(args.system)
    for src in inputs:
        dst = output_path_for(src, args.input, args.out)
        print(f"{src} -> {dst}", file=sys.stderr)
        source_text = read_text(src)
        source_stats = text_stats(source_text)
        if args.dry_run:
            print(
                f"source chars={source_stats['chars']} nonspace={source_stats['nonspace']} zh={source_stats['zh']}",
                file=sys.stderr,
            )
            continue
        started = time.time()
        try:
            polished = polish_text(
                source_text,
                system_prompt,
                base_url=args.base_url,
                model=args.model,
                temperature=args.temperature,
                top_p=args.top_p,
                presence_penalty=args.presence_penalty,
                frequency_penalty=args.frequency_penalty,
                repeat_penalty=args.repeat_penalty,
                max_tokens=args.max_tokens,
                timeout=args.timeout,
            )
        except urllib.error.URLError as exc:
            print(f"LM Studio request failed: {exc}", file=sys.stderr)
            return 1
        write_text(dst, polished)
        polished_stats = text_stats(polished)
        print(
            f"wrote {dst} ({time.time() - started:.1f}s) "
            f"source_zh={source_stats['zh']} output_zh={polished_stats['zh']} "
            f"source_nonspace={source_stats['nonspace']} output_nonspace={polished_stats['nonspace']}",
            file=sys.stderr,
        )
        print(dst)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
