#!/usr/bin/env python3
"""
临时健康检查脚本：验证本机 Vertex AI 基础配置是否完好。
使用说明：
  python3 tmp_rovodev_vertex_check.py
运行完成后请将输出粘贴给我，我会据此给出修复建议。
完成诊断后可删除该文件。
"""

import os
import json
import sys
from pathlib import Path

STATUS = []

def ok(msg):
    STATUS.append((True, msg))
    print(f"✅ {msg}")

def warn(msg):
    STATUS.append((None, msg))
    print(f"⚠️ {msg}")

def fail(msg):
    STATUS.append((False, msg))
    print(f"❌ {msg}")

def has(var):
    return bool(os.getenv(var))

print("🧪 Vertex AI 本地配置健康检查\n")

# 1) 关键环境变量与文件
adc_path = Path.home() / ".config/gcloud/application_default_credentials.json"
if adc_path.exists():
    ok(f"找到 ADC (Application Default Credentials): {adc_path}")
else:
    warn("未找到 ADC 文件 (~/.config/gcloud/application_default_credentials.json)。如果你改用服务账号密钥也可忽略此项")

vars_to_check = [
    "GOOGLE_GENAI_USE_VERTEXAI",
    "GOOGLE_CLOUD_PROJECT_ID",
    "GOOGLE_CLOUD_LOCATION",
    "VERTEX_MEMORY_BANK_ENABLED",
    "GOOGLE_API_KEY",
    "GOOGLE_SERVICE_ACCOUNT_KEY",
]

print("\n🔍 环境变量：")
for v in vars_to_check:
    val = os.getenv(v)
    if not val:
        if v in ("GOOGLE_SERVICE_ACCOUNT_KEY", "GOOGLE_API_KEY"):
            # 可选
            warn(f"{v}: 未设置 (可选)")
        else:
            fail(f"{v}: 未设置")
    else:
        if v.endswith("KEY"):
            ok(f"{v}: 已设置 (不显示具体值)")
        else:
            ok(f"{v}: {val}")

# 2) 依赖与版本
print("\n🔍 依赖：google-cloud-aiplatform")
try:
    import google
    from google.cloud import aiplatform
    ok(f"google-cloud-aiplatform 可导入，版本: {aiplatform.__version__}")
except Exception as e:
    fail(f"无法导入 google-cloud-aiplatform: {e}")

# 3) 验证默认凭据可用性（若存在）
print("\n🔍 默认凭据 (ADC) 可用性：")
try:
    import google.auth
    from google.auth.transport.requests import Request
    creds, proj = google.auth.default(scopes=[
        "https://www.googleapis.com/auth/cloud-platform",
    ])
    # 刷新一次，验证有效性
    try:
        creds.refresh(Request())
        ok("默认凭据可用并成功刷新访问令牌")
    except Exception as e:
        warn(f"找到默认凭据，但刷新失败（可能未登录或网络受限）：{e}")
    if proj:
        ok(f"默认凭据解析到的 Project: {proj}")
    else:
        warn("默认凭据未解析出 Project ID")
except Exception as e:
    warn(f"未能通过 google.auth.default() 获取默认凭据：{e}")

# 4) 初始化 Vertex AI（使用环境中的 Project/Location）
print("\n🔍 Vertex AI 初始化：")
project = os.getenv("GOOGLE_CLOUD_PROJECT_ID")
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
if project:
    try:
        from google.cloud import aiplatform
        aiplatform.init(project=project, location=location)
        ok(f"Vertex AI 初始化成功: project={project}, location={location}")
    except Exception as e:
        fail(f"Vertex AI 初始化失败: {e}")
else:
    fail("缺少 GOOGLE_CLOUD_PROJECT_ID，无法初始化 Vertex AI")

# 5) 汇总
print("\n📋 检查摘要：")
all_good = True
for s, msg in STATUS:
    if s is False:
        all_good = False

if all_good:
    print("\n🎉 结论：看起来你的 Vertex AI 本地配置是可用的。")
else:
    print("\n🚧 结论：存在未通过项。请根据上面的 ❌/⚠️ 提示逐项修复后重试。")

print("\n提示：\n- 启用 Vertex 模式：export GOOGLE_GENAI_USE_VERTEXAI=TRUE\n- 必填：export GOOGLE_CLOUD_PROJECT_ID=你的项目ID\n- 可选：export GOOGLE_CLOUD_LOCATION=us-central1\n- 建议使用 ADC：gcloud auth application-default login\n- 或设置服务账号密钥：export GOOGLE_SERVICE_ACCOUNT_KEY='(JSON 字符串或路径)'\n")