from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from pathlib import Path

from .config import infer_kind, load_config


PROJECT_DIR = Path(__file__).resolve().parent.parent
DEFAULT_CONFIG = PROJECT_DIR / "config" / "config.json"


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="localpilot", description="LocalPilot 本地模型后端")
    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="配置 JSON 路径")
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("models", help="列出配置模型")
    sub.add_parser("doctor", help="检查本机运行环境")
    serve = sub.add_parser("serve", help="启动 OpenAI 兼容 API")
    serve.add_argument("--host")
    serve.add_argument("--port", type=int)
    sub.add_parser("tui", help="启动 Textual TUI")
    load = sub.add_parser("load", help="预加载模型")
    load.add_argument("model_id")
    inspect = sub.add_parser("inspect", help="推断模型格式")
    inspect.add_argument("path")
    return parser


def _doctor() -> dict[str, object]:
    result: dict[str, object] = {"python": sys.version, "executable": sys.executable}
    try:
        import torch

        result["torch"] = torch.__version__
        result["cuda_available"] = bool(torch.cuda.is_available())
        result["cuda_device"] = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
        result["cuda_version"] = torch.version.cuda
    except Exception as exc:
        result["torch_error"] = str(exc)
    try:
        import onnxruntime as ort

        result["onnxruntime"] = ort.__version__
        result["onnx_providers"] = ort.get_available_providers()
    except Exception as exc:
        result["onnx_error"] = str(exc)
    result["conda_prefix"] = os.getenv("CONDA_PREFIX")
    return result


def main() -> None:
    args = _parser().parse_args()
    config_path = Path(args.config)
    config = load_config(config_path)

    if args.command == "models":
        print(json.dumps([
            {"id": model.id, "kind": infer_kind(model), "path": model.path, "enabled": model.enabled}
            for model in config.models
            if model.enabled
        ], ensure_ascii=False, indent=2))
        return
    if args.command == "doctor":
        print(json.dumps(_doctor(), ensure_ascii=False, indent=2))
        return
    if args.command == "inspect":
        from .config import ModelConfig

        spec = ModelConfig(id=Path(args.path).stem, path=args.path)
        print(json.dumps({"path": args.path, "kind": infer_kind(spec)}, ensure_ascii=False, indent=2))
        return
    if args.command == "serve":
        import uvicorn

        from .server import create_app

        app = create_app(config_path)
        uvicorn.run(
            app,
            host=args.host or config.runtime.host,
            port=args.port or config.runtime.port,
            log_level="info",
        )
        return
    if args.command == "tui":
        from .tui import LocalPilotTUI

        LocalPilotTUI(config).run()
        return
    if args.command == "load":
        from .manager import ModelManager

        async def run() -> None:
            manager = ModelManager(config)
            try:
                backend = await manager.load(args.model_id)
                print(json.dumps(backend.info(), ensure_ascii=False, indent=2))
            finally:
                await manager.shutdown()

        asyncio.run(run())


if __name__ == "__main__":
    main()