"""Command line interface for kbase."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, Optional

from .kb_manager import KBManager
from .pdf_parser import parse_pdfs, parse_single_pdf
from .pipeline import build_kb_from_pdfs


def _resolve_arg_value(
    cli_value: Optional[str],
    config: Dict[str, Any],
    key: str,
    *,
    required: bool = True,
) -> str:
    value = cli_value if cli_value else config.get(key)
    if required and not value:
        raise ValueError(f"缺少必填参数: {key}")
    return str(value)


def _resolve_optional_arg_value(
    cli_value: Optional[str],
    config: Dict[str, Any],
    key: str,
) -> Optional[str]:
    value = cli_value if cli_value is not None else config.get(key)
    if value is None or value == "":
        return None
    return str(value)


def _load_config_file(path: Optional[str]) -> Dict[str, Any]:
    if not path:
        return {}
    config_path = Path(path)
    if not config_path.exists():
        raise FileNotFoundError(f"配置文件不存在: {config_path}")

    suffix = config_path.suffix.lower()
    if suffix in {".json"}:
        return json.loads(config_path.read_text(encoding="utf-8"))
    if suffix in {".yaml", ".yml"}:
        try:
            import yaml  # type: ignore
        except Exception as exc:  # pragma: no cover
            raise RuntimeError("读取 YAML 需要安装 pyyaml。") from exc
        payload = yaml.safe_load(config_path.read_text(encoding="utf-8"))
        return payload if isinstance(payload, dict) else {}
    raise ValueError(f"不支持的配置文件格式: {config_path}")


def _print_result(payload: Any, as_json: bool) -> None:
    if as_json:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    else:
        print(payload)


def _cmd_parse(args: argparse.Namespace) -> int:
    config = _load_config_file(args.config).get("parse", {})
    input_dir = _resolve_arg_value(args.input, config, "input")
    output_dir = _resolve_arg_value(args.output, config, "output")
    result = parse_pdfs(
        input_dir=input_dir,
        output_dir=output_dir,
        parser_backend=args.parser_backend or config.get("parser_backend", "auto"),
        image_backend=args.image_backend or config.get("image_backend", "auto"),
        extract_images=not args.no_images,
        return_base64=args.return_base64,
        docling_artifacts_path=_resolve_optional_arg_value(
            args.docling_artifacts_path,
            config,
            "docling_artifacts_path",
        ),
    )
    _print_result(result, args.json)
    return 0


def _cmd_parse_one(args: argparse.Namespace) -> int:
    config = _load_config_file(args.config).get("parse_one", {})
    file_path = _resolve_arg_value(args.file, config, "file")
    output_dir = _resolve_optional_arg_value(args.output, config, "output")
    if output_dir is None:
        # 默认写回到 PDF 所在目录
        output_dir = str(Path(file_path).expanduser().resolve().parent)

    result = parse_single_pdf(
        file_path=file_path,
        output_dir=output_dir,
        parser_backend=args.parser_backend or config.get("parser_backend", "auto"),
        image_backend=args.image_backend or config.get("image_backend", "auto"),
        extract_images=not args.no_images,
        return_base64=args.return_base64,
        docling_artifacts_path=_resolve_optional_arg_value(
            args.docling_artifacts_path,
            config,
            "docling_artifacts_path",
        ),
    )

    stem = Path(file_path).stem
    payload = {
        "source": str(Path(file_path)),
        "output_dir": str(Path(output_dir)),
        "text_file": str(Path(output_dir) / f"{stem}_text.md"),
        "json_file": str(Path(output_dir) / f"{stem}_parsed.json"),
        "metadata": result.get("metadata", {}),
    }
    _print_result(payload, args.json)
    return 0


def _cmd_build(args: argparse.Namespace) -> int:
    config = _load_config_file(args.config).get("build", {})
    name = _resolve_arg_value(args.name, config, "name")
    input_dir = _resolve_arg_value(args.input, config, "input")
    output_dir = _resolve_arg_value(args.output, config, "output")
    result = build_kb_from_pdfs(
        name=name,
        input_dir=input_dir,
        output_dir=output_dir,
        embedding_model=args.embedding_model or config.get("embedding_model", "all-MiniLM-L6-v2"),
        parser_backend=args.parser_backend or config.get("parser_backend", "auto"),
        image_backend=args.image_backend or config.get("image_backend", "auto"),
        docling_artifacts_path=_resolve_optional_arg_value(
            args.docling_artifacts_path,
            config,
            "docling_artifacts_path",
        ),
        mode=args.mode or config.get("mode", "rebuild"),
        overwrite=args.overwrite or bool(config.get("overwrite", False)),
        chunk_size=args.chunk_size or int(config.get("chunk_size", 800)),
        chunk_overlap=args.chunk_overlap or int(config.get("chunk_overlap", 80)),
        register_in_manager=False,
    )
    payload = {
        "kb_name": result.kb_name,
        "kb_path": result.kb_path,
        "parse_stats": result.parse_stats,
        "build_result": result.build_result.to_dict(),
    }
    _print_result(payload, args.json)
    return 0


def _cmd_query(args: argparse.Namespace) -> int:
    config = _load_config_file(args.config).get("query", {})
    db_path = _resolve_arg_value(args.db, config, "db")
    query_text = _resolve_arg_value(args.q, config, "q")

    kb_name = args.kb_name or config.get("kb_name", "default")
    manager = KBManager(default_embedding_model=args.embedding_model or config.get("embedding_model"))
    manager.add_kb(kb_name, db_path, embedding_model=args.embedding_model or config.get("embedding_model"))

    if args.json:
        payload = manager.query_json(
            query=query_text,
            top_k=args.top_k or int(config.get("top_k", 5)),
            kb_name=kb_name,
            collection_name=args.collection or config.get("collection"),
        )
        _print_result(payload, as_json=True)
    else:
        xml = manager.query(
            query=query_text,
            top_k=args.top_k or int(config.get("top_k", 5)),
            kb_name=kb_name,
            collection_name=args.collection or config.get("collection"),
        )
        print(xml)
    manager.close()
    return 0


def _cmd_eval(args: argparse.Namespace) -> int:
    """
    Simple retrieval eval for jsonl data:
    {"query": "...", "must_contain": ["keyword1", "keyword2"]}
    """
    manager = KBManager(default_embedding_model=args.embedding_model)
    manager.add_kb(args.kb_name, args.db, embedding_model=args.embedding_model)

    total = 0
    hit = 0
    for line in Path(args.input).read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        sample = json.loads(line)
        query = str(sample.get("query", ""))
        required = [str(item).lower() for item in sample.get("must_contain", [])]
        payload = manager.query_json(query=query, top_k=args.top_k, kb_name=args.kb_name)
        docs = " ".join(item.get("document", "").lower() for item in payload.get("results", []))
        total += 1
        if not required or all(term in docs for term in required):
            hit += 1

    manager.close()
    report = {"total": total, "hit": hit, "recall_at_k": (hit / total if total else 0.0)}
    _print_result(report, args.json)
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="kbase", description="kbase CLI")
    parser.add_argument("--config", help="配置文件（json/yaml）", default=None)
    sub = parser.add_subparsers(dest="command", required=True)

    parse_cmd = sub.add_parser("parse", help="解析目录中的全部 PDF 到结构化文本")
    parse_cmd.add_argument("--input", required=False, help="PDF 输入目录")
    parse_cmd.add_argument("--output", required=False, help="解析输出目录")
    parse_cmd.add_argument("--parser-backend", default=None, choices=["auto", "docling", "pymupdf"])
    parse_cmd.add_argument("--image-backend", default=None, choices=["auto", "pymupdf", "pypdfium2", "none"])
    parse_cmd.add_argument("--docling-artifacts-path", default=None, help="Docling 模型目录（artifacts_path）")
    parse_cmd.add_argument("--no-images", action="store_true", help="不提取图片")
    parse_cmd.add_argument("--return-base64", action="store_true", help="在 JSON 中返回 base64 图片")
    parse_cmd.add_argument("--json", action="store_true")
    parse_cmd.set_defaults(func=_cmd_parse)

    parse_one_cmd = sub.add_parser("parse-one", help="解析单个 PDF（默认输出到 PDF 同目录）")
    parse_one_cmd.add_argument("--file", required=False, help="PDF 文件路径")
    parse_one_cmd.add_argument("--output", required=False, help="输出目录；不传时默认 PDF 同目录")
    parse_one_cmd.add_argument("--parser-backend", default=None, choices=["auto", "docling", "pymupdf"])
    parse_one_cmd.add_argument("--image-backend", default=None, choices=["auto", "pymupdf", "pypdfium2", "none"])
    parse_one_cmd.add_argument("--docling-artifacts-path", default=None, help="Docling 模型目录（artifacts_path）")
    parse_one_cmd.add_argument("--no-images", action="store_true", help="不提取图片")
    parse_one_cmd.add_argument("--return-base64", action="store_true", help="在 JSON 中返回 base64 图片")
    parse_one_cmd.add_argument("--json", action="store_true")
    parse_one_cmd.set_defaults(func=_cmd_parse_one)

    build_cmd = sub.add_parser("build", help="一站式构建 KB（parse + vector build）")
    build_cmd.add_argument("--name", required=False, help="知识库名称")
    build_cmd.add_argument("--input", required=False, help="PDF 输入目录")
    build_cmd.add_argument("--output", required=False, help="输出根目录")
    build_cmd.add_argument("--embedding-model", default=None)
    build_cmd.add_argument("--parser-backend", default=None, choices=["auto", "docling", "pymupdf"])
    build_cmd.add_argument("--image-backend", default=None, choices=["auto", "pymupdf", "pypdfium2", "none"])
    build_cmd.add_argument("--docling-artifacts-path", default=None, help="Docling 模型目录（artifacts_path）")
    build_cmd.add_argument("--mode", default=None, choices=["rebuild", "append", "sync"])
    build_cmd.add_argument("--overwrite", action="store_true")
    build_cmd.add_argument("--chunk-size", type=int, default=None)
    build_cmd.add_argument("--chunk-overlap", type=int, default=None)
    build_cmd.add_argument("--json", action="store_true")
    build_cmd.set_defaults(func=_cmd_build)

    query_cmd = sub.add_parser("query", help="查询向量库")
    query_cmd.add_argument("--db", required=False, help="chroma_store 路径")
    query_cmd.add_argument("--kb-name", default="default", help="KB 注册名")
    query_cmd.add_argument("--q", required=False, help="查询文本")
    query_cmd.add_argument("--top-k", type=int, default=None)
    query_cmd.add_argument("--collection", default=None)
    query_cmd.add_argument("--embedding-model", default=None)
    query_cmd.add_argument("--json", action="store_true")
    query_cmd.set_defaults(func=_cmd_query)

    eval_cmd = sub.add_parser("eval", help="离线评估 recall@k")
    eval_cmd.add_argument("--db", required=True, help="chroma_store 路径")
    eval_cmd.add_argument("--kb-name", default="default")
    eval_cmd.add_argument("--input", required=True, help="jsonl 评估集路径")
    eval_cmd.add_argument("--top-k", type=int, default=5)
    eval_cmd.add_argument("--embedding-model", default=None)
    eval_cmd.add_argument("--json", action="store_true")
    eval_cmd.set_defaults(func=_cmd_eval)

    return parser


def main(argv: Optional[list[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return int(args.func(args))


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())
