src/deberta_ime/cli.py

from __future__ import annotations

import argparse
import io
import json
import sqlite3
import sys
from collections.abc import Sequence
from contextlib import ExitStack
from dataclasses import asdict
from pathlib import Path
from typing import TextIO

from .deberta import MODEL_ID, MODEL_REVISION, DebertaCandidateScorer
from .domain import Candidate, CandidateScorer, RerankConfig, RerankRequest, RerankResult
from .mozc import MOZC_REVISION, MozcDictionaryIndex, build_mozc_index
from .profiles import MOZC_PROFILE_SELECTION_REVISION, MOZC_PROFILES
from .reranker import Reranker
from .sidecar import serve


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="deberta-ime")
    subparsers = parser.add_subparsers(dest="command", required=True)
    rerank = subparsers.add_parser("rerank", help="rerank JSON Lines from stdin")
    rerank.add_argument("--prior-weight", type=float, default=2.0)
    rerank.add_argument("--min-margin", type=float, default=0.1)
    rerank.add_argument("--device", default="cpu")
    rerank.add_argument("--cache-dir", type=Path)
    rerank.add_argument("--offline", action="store_true")
    rerank.add_argument("--input", type=Path, help="UTF-8 JSONL file (default: stdin)")
    rerank.add_argument("--output", type=Path, help="UTF-8 JSONL file (default: stdout)")
    serve_parser = subparsers.add_parser(
        "serve", help="run a persistent UTF-8 JSON Lines reranking sidecar"
    )
    serve_parser.add_argument("--device", default="cpu")
    serve_parser.add_argument("--cache-dir", type=Path)
    serve_parser.add_argument("--offline", action="store_true")
    mozc_index = subparsers.add_parser(
        "mozc-index", help="build a local SQLite index from Mozc OSS dictionary files"
    )
    mozc_index.add_argument("--dictionary-dir", type=Path, required=True)
    mozc_index.add_argument("--output", type=Path, required=True)
    mozc_index.add_argument("--source-revision", default=MOZC_REVISION)
    mozc_rerank = subparsers.add_parser(
        "mozc-rerank", help="generate finite candidates from a Mozc index and rerank JSON Lines"
    )
    mozc_rerank.add_argument("--index", type=Path, required=True)
    mozc_rerank.add_argument("--profile", choices=tuple(MOZC_PROFILES), default="incremental")
    mozc_rerank.add_argument("--candidate-limit", type=int, default=8)
    mozc_rerank.add_argument("--device", default="cpu")
    mozc_rerank.add_argument("--cache-dir", type=Path)
    mozc_rerank.add_argument("--offline", action="store_true")
    mozc_rerank.add_argument("--input", type=Path, help="UTF-8 JSONL file (default: stdin)")
    mozc_rerank.add_argument("--output", type=Path, help="UTF-8 JSONL file (default: stdout)")
    return parser


def _request_from_payload(payload: object) -> RerankRequest:
    if not isinstance(payload, dict):
        raise ValueError("input must be a JSON object")
    reading = payload.get("reading")
    raw_candidates = payload.get("candidates")
    if not isinstance(reading, str):
        raise ValueError("reading must be a string")
    if not isinstance(raw_candidates, list):
        raise ValueError("candidates must be a list")

    candidates: list[Candidate] = []
    for raw_candidate in raw_candidates:
        if isinstance(raw_candidate, str):
            candidates.append(Candidate(raw_candidate))
            continue
        if not isinstance(raw_candidate, dict) or not isinstance(
            raw_candidate.get("surface"), str
        ):
            raise ValueError("each candidate must be a string or an object with surface")
        prior_score = raw_candidate.get("prior_score", 0.0)
        if not isinstance(prior_score, int | float):
            raise ValueError("candidate prior_score must be numeric")
        candidates.append(Candidate(raw_candidate["surface"], float(prior_score)))

    left = _string_list(payload.get("left_context", []), "left_context")
    right = _string_list(payload.get("right_context", []), "right_context")
    return RerankRequest(
        reading=reading,
        candidates=tuple(candidates),
        left_context=left,
        right_context=right,
    )


def _string_list(value: object, name: str) -> tuple[str, ...]:
    if not isinstance(value, list) or not all(isinstance(item, str) for item in value):
        raise ValueError(f"{name} must be a list of strings")
    return tuple(value)


def _mozc_request_from_payload(
    payload: object,
    index: MozcDictionaryIndex,
    *,
    candidate_limit: int,
    profile: str,
) -> RerankRequest:
    if not isinstance(payload, dict):
        raise ValueError("input must be a JSON object")
    reading = payload.get("reading")
    if not isinstance(reading, str):
        raise ValueError("reading must be a string")
    right_context = _string_list(payload.get("right_context", []), "right_context")
    if profile == "incremental" and right_context:
        raise ValueError("incremental profile does not accept right_context")
    return RerankRequest(
        reading=reading,
        candidates=index.lookup(reading, limit=candidate_limit),
        left_context=_string_list(payload.get("left_context", []), "left_context"),
        right_context=right_context,
    )


def _result_payload(result: RerankResult) -> dict[str, object]:
    return {
        "ok": True,
        "schema_version": 1,
        "model": {"id": MODEL_ID, "revision": MODEL_REVISION},
        "reading": result.reading,
        "decision": {
            "changed": result.changed,
            "reason": result.reason,
            "margin": result.margin,
        },
        "ranked": [
            {
                "surface": item.surface,
                "original_rank": item.original_rank,
                "prior_score": item.prior_score,
                "model_score": item.model_score,
                "combined_score": item.combined_score,
            }
            for item in result.ranked
        ],
    }


def run(
    argv: Sequence[str] | None = None,
    *,
    stdin: TextIO | None = None,
    stdout: TextIO | None = None,
    stderr: TextIO | None = None,
    scorer: CandidateScorer | None = None,
) -> int:
    args = _parser().parse_args(argv)
    error_stream = stderr or sys.stderr
    if args.command == "serve":
        active_scorer = scorer or DebertaCandidateScorer(
            device=args.device,
            cache_dir=args.cache_dir,
            local_files_only=args.offline,
        )
        input_stream = stdin
        if input_stream is None:
            if isinstance(sys.stdin, io.TextIOWrapper):
                sys.stdin.reconfigure(encoding="utf-8", errors="strict")
            input_stream = sys.stdin
        output_stream = stdout
        if output_stream is None:
            if isinstance(sys.stdout, io.TextIOWrapper):
                sys.stdout.reconfigure(
                    encoding="utf-8",
                    errors="strict",
                    newline="\n",
                    write_through=True,
                )
            output_stream = sys.stdout
        return serve(input_stream, output_stream, active_scorer)
    if args.command == "mozc-index":
        try:
            manifest = build_mozc_index(
                args.dictionary_dir,
                args.output,
                source_revision=args.source_revision,
            )
        except (OSError, ValueError, sqlite3.Error) as error:
            error_stream.write(f"error: {error}\n")
            return 2
        output_stream = stdout or sys.stdout
        output_stream.write(
            json.dumps(
                {"ok": True, "schema_version": 1, "manifest": asdict(manifest)},
                ensure_ascii=False,
                separators=(",", ":"),
            )
        )
        output_stream.write("\n")
        return 0
    if args.command == "mozc-rerank":
        if (
            args.input is not None
            and args.output is not None
            and args.input.resolve() == args.output.resolve()
        ):
            error_stream.write("error: --input and --output must be different paths\n")
            return 2
        active_scorer = scorer or DebertaCandidateScorer(
            device=args.device,
            cache_dir=args.cache_dir,
            local_files_only=args.offline,
        )
        selected_profile = MOZC_PROFILES[args.profile]
        reranker = Reranker(active_scorer, selected_profile)
        try:
            index = MozcDictionaryIndex(args.index)
        except (OSError, ValueError, sqlite3.Error) as error:
            error_stream.write(f"error: cannot open Mozc index: {error}\n")
            return 2
        with index, ExitStack() as stack:
            input_stream = (
                stack.enter_context(args.input.open(encoding="utf-8"))
                if args.input is not None
                else (stdin or sys.stdin)
            )
            output_stream = (
                stack.enter_context(args.output.open("w", encoding="utf-8", newline="\n"))
                if args.output is not None
                else (stdout or sys.stdout)
            )
            exit_code = 0
            for line_number, raw_line in enumerate(input_stream, start=1):
                if not raw_line.strip():
                    continue
                try:
                    request = _mozc_request_from_payload(
                        json.loads(raw_line),
                        index,
                        candidate_limit=args.candidate_limit,
                        profile=args.profile,
                    )
                    payload = _result_payload(reranker.rerank(request))
                    payload["candidate_source"] = {
                        "kind": "mozc_oss_dictionary_index",
                        "source_revision": index.manifest.source_revision,
                        "limit": args.candidate_limit,
                    }
                    payload["profile"] = args.profile
                    payload["profile_config"] = {
                        "prior_weight": selected_profile.prior_weight,
                        "min_margin": selected_profile.min_margin,
                        "selected_on_revision": MOZC_PROFILE_SELECTION_REVISION,
                    }
                except (json.JSONDecodeError, ValueError, TypeError) as error:
                    payload = {
                        "ok": False,
                        "schema_version": 1,
                        "line": line_number,
                        "error": str(error),
                    }
                    exit_code = 2
                output_stream.write(
                    json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
                )
                output_stream.write("\n")
            return exit_code
    if (
        args.input is not None
        and args.output is not None
        and args.input.resolve() == args.output.resolve()
    ):
        error_stream.write("error: --input and --output must be different paths\n")
        return 2
    active_scorer = scorer or DebertaCandidateScorer(
        device=args.device,
        cache_dir=args.cache_dir,
        local_files_only=args.offline,
    )
    reranker = Reranker(
        active_scorer,
        RerankConfig(prior_weight=args.prior_weight, min_margin=args.min_margin),
    )
    with ExitStack() as stack:
        input_stream = (
            stack.enter_context(args.input.open(encoding="utf-8"))
            if args.input is not None
            else (stdin or sys.stdin)
        )
        output_stream = (
            stack.enter_context(args.output.open("w", encoding="utf-8", newline="\n"))
            if args.output is not None
            else (stdout or sys.stdout)
        )
        exit_code = 0
        for line_number, raw_line in enumerate(input_stream, start=1):
            if not raw_line.strip():
                continue
            try:
                request = _request_from_payload(json.loads(raw_line))
                payload = _result_payload(reranker.rerank(request))
            except (json.JSONDecodeError, ValueError, TypeError) as error:
                payload = {
                    "ok": False,
                    "schema_version": 1,
                    "line": line_number,
                    "error": str(error),
                }
                exit_code = 2
            output_stream.write(json.dumps(payload, ensure_ascii=False, separators=(",", ":")))
            output_stream.write("\n")
        return exit_code


def main() -> None:
    raise SystemExit(run())


if __name__ == "__main__":
    main()