#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import json
import argparse
from typing import Any, Dict

import requests


def build_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="调用本地 llama-server (OpenAI兼容) 进行推理，支持自定义系统/用户提示词"
    )
    parser.add_argument(
        "--base-url",
        default=os.getenv("SIMPLE_BASE_URL", "http://127.0.0.1:8081/v1"),
        help="llama-server 的基础URL（默认: http://127.0.0.1:8081/v1，或环境变量 SIMPLE_BASE_URL）",
    )
    parser.add_argument(
        "--model",
        default=os.getenv("SIMPLE_MODEL", "local-model"),
        help="模型名称（默认: local-model，或环境变量 SIMPLE_MODEL）",
    )
    parser.add_argument(
        "--system",
        default="You are a helpful assistant.",
        help="系统提示词（system role）",
    )
    parser.add_argument(
        "--system-file",
        default=None,
        help="系统提示词文件路径（txt）；若提供，则覆盖 --system 的字符串",
    )
    parser.add_argument(
        "--user",
        default=None,
        help="用户提示词（user role）；若不传则从交互式输入读取",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.2,
        help="采样温度（默认: 0.2）",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=4096,
        help="最大生成Token数（默认: 4096）",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=120.0,
        help="HTTP超时时间秒（默认: 120）",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="打印完整返回JSON",
    )
    return parser.parse_args()


def call_llama_server(
    base_url: str,
    model: str,
    system_prompt: str,
    user_prompt: str,
    temperature: float,
    max_tokens: int,
    timeout: float,
) -> Dict[str, Any]:
    endpoint = base_url.rstrip("/") + "/chat/completions"
    headers: Dict[str, str] = {"Content-Type": "application/json"}

    # 兼容需要API Key的代理/服务（llama-server通常不强制）
    api_key = os.getenv("OPENAI_API_KEY")
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    payload: Dict[str, Any] = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    resp = requests.post(endpoint, headers=headers, data=json.dumps(payload), timeout=timeout)
    resp.raise_for_status()
    return resp.json()


def main() -> None:
    args = build_args()

    user_prompt = args.user
    if not user_prompt:
        try:
            user_prompt = input("请输入用户提示词: ")
        except KeyboardInterrupt:
            print("\n已取消。")
            sys.exit(1)

    # 解析系统提示词：优先使用 --system-file
    system_prompt = args.system
    if args.system_file:
        try:
            with open(args.system_file, "r", encoding="utf-8") as f:
                system_prompt = f.read()
        except Exception as e:
            print("\n❌ 读取系统提示词文件失败：")
            print(str(e))
            sys.exit(1)

    try:
        print("--- llama-server 推理 ---")
        print(f"Base URL: {args.base_url}")
        print(f"Model: {args.model}")
        if args.system_file:
            print(f"System(from file): {args.system_file}")
        else:
            print(f"System: {system_prompt}")
        print(f"User: {user_prompt}")

        data = call_llama_server(
            base_url=args.base_url,
            model=args.model,
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
            timeout=args.timeout,
        )

        if args.verbose:
            print("\n完整返回JSON:")
            print(json.dumps(data, ensure_ascii=False, indent=2))

        # 尝试按OpenAI兼容格式提取assistant内容
        content = None
        try:
            content = data["choices"][0]["message"]["content"]
        except Exception:
            pass

        if content is not None:
            print("\n模型输出:")
            print(content)
        else:
            # 兜底打印
            print("\n无法按OpenAI兼容字段解析内容，原始返回如下：")
            print(json.dumps(data, ensure_ascii=False))

    except requests.exceptions.RequestException as e:
        print("\n❌ 请求失败：请确认 llama-server 已在 8081 端口启动并可访问。")
        print(f"详情: {e}")
        sys.exit(2)
    except Exception as e:
        print("\n❌ 发生未预期的错误：")
        print(str(e))
        sys.exit(3)


if __name__ == "__main__":
    main()