"""
测试脚本：上传 PDF / DOCX / TXT 文件到阿里云 OSS → 获取 URL → 发送给 GLM-4.6V 识别

支持的文件类型:
  - .pdf   → 上传 OSS 后以 file_url 类型发送 URL 给 GLM
  - .docx  → 上传 OSS 后以 file_url 类型发送 URL 给 GLM
  - .txt   → 上传 OSS 后以 file_url 类型发送 URL 给 GLM

用法:
    cd server
    source ~/.bashrc && source .venv/bin/activate
    python -m utils.test_oss_doc_glm --file <本地文件路径> [--prompt "请总结这份文件"]
"""

import argparse
import sys
from pathlib import Path

# 确保 server 目录在 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from utils.oss_uploader import upload_file
from utils.glm_adapter import glm_chat_sync


# 文件类型分组
DOC_EXTS = {".pdf", ".doc", ".docx", ".xlsx", ".xls", ".pptx", ".ppt"}
TXT_EXTS = {".txt", ".md", ".csv", ".json", ".log"}
IMG_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"}
# 所有可通过 file_url 发送的类型
FILE_URL_EXTS = DOC_EXTS | TXT_EXTS


def detect_file_type(suffix: str) -> str:
    """根据后缀判断文件类别: 'file_url' / 'image' / 'unknown'"""
    suffix = suffix.lower()
    if suffix in FILE_URL_EXTS:
        return "file_url"
    elif suffix in IMG_EXTS:
        return "image"
    return "unknown"


def build_messages_for_file_url(file_url: str, prompt: str) -> list:
    """
    为文档/文本文件构建消息。
    使用 file_url 类型，直接传递 OSS 的 URL 给 GLM。
    """
    return [
        {
            "role": "user",
            "content": [
                {
                    "type": "file_url",
                    "file_url": {"url": file_url},
                },
                {
                    "type": "text",
                    "text": prompt,
                },
            ],
        }
    ]


def build_messages_for_image(file_url: str, prompt: str) -> list:
    """为图片文件构建消息，使用 image_url 类型。"""
    return [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": file_url}},
                {"type": "text", "text": prompt},
            ],
        }
    ]


def main():
    parser = argparse.ArgumentParser(
        description="上传 PDF/DOCX/TXT 文件到 OSS 并让 GLM-4.6V 识别"
    )
    parser.add_argument("--file", required=True, help="要上传的本地文件路径")
    parser.add_argument(
        "--prompt", default="请总结这份文件的主要内容", help="发给 GLM 的提示词"
    )
    parser.add_argument(
        "--model", default="glm-4.6v", help="GLM 模型名称（默认: glm-4.6v）"
    )
    args = parser.parse_args()

    file_path = Path(args.file).resolve()
    if not file_path.exists():
        print(f"❌ 文件不存在: {file_path}")
        sys.exit(1)

    suffix = file_path.suffix.lower()
    file_type = detect_file_type(suffix)
    print(f"📂 文件信息:")
    print(f"   路径: {file_path}")
    print(f"   后缀: {suffix}")
    print(f"   类型: {file_type}")
    print(f"   大小: {file_path.stat().st_size / 1024:.1f} KB")
    print()

    # ── 第一步：上传文件到 OSS ────────────────────────────────
    print(f"📤 正在上传文件到阿里云 OSS...")
    oss_result = upload_file(str(file_path))
    file_url = oss_result["url"]
    print(f"✅ OSS 上传成功!")
    print(f"   URL: {file_url}")
    print(f"   ETag: {oss_result['etag']}")
    print()

    # ── 第二步：根据文件类型构建消息 ──────────────────────────
    print(f"🔧 正在构建 GLM 消息...")

    if file_type == "file_url":
        # PDF / DOCX / TXT 等：使用 file_url 类型发送 OSS URL
        print(f"   策略: 使用 file_url 发送 OSS 链接")
        messages = build_messages_for_file_url(file_url, args.prompt)
    elif file_type == "image":
        # 图片：使用 image_url
        print(f"   策略: 使用 image_url 发送 OSS 链接")
        messages = build_messages_for_image(file_url, args.prompt)
    else:
        print(f"❌ 不支持的文件类型: {suffix}")
        print(f"   支持: {', '.join(sorted(FILE_URL_EXTS | IMG_EXTS))}")
        sys.exit(1)

    print()

    # ── 第三步：发送给 GLM 识别 ──────────────────────────────
    print(f"🤖 正在请求 GLM ({args.model}) 识别文件...")
    print(f"   提示词: {args.prompt}")
    print()

    try:
        result = glm_chat_sync(
            messages=messages,
            model=args.model,
            temperature=0.7,
            max_tokens=4096,
        )

        print("─" * 60)
        print("📝 GLM 回复:")
        print("─" * 60)
        print(result["content"])
        print("─" * 60)

        if result.get("usage"):
            usage = result["usage"]
            print(
                f"\n📊 Token 用量: 输入 {usage['promptTokens']} | "
                f"输出 {usage['completionTokens']} | "
                f"总计 {usage['totalTokens']}"
            )

        print(f"\n✅ 测试完成! 使用模型: {result.get('model', args.model)}")

    except Exception as e:
        print(f"\n❌ GLM 请求失败: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()