app.py

"""
API 反代服务 - 主应用程序
Gradio 管理界面 + FastAPI API 端点
"""
import os
import json
import gradio as gr
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager

from models import OpenAIChatRequest, OpenAIMessage, SUPPORTED_MODELS, ServiceConfig
from account_manager import account_manager, config_manager
from load_balancer import load_balancer
from proxy_service import stream_chat_completion, chat_completion

# ============ 服务配置 ============
config = ServiceConfig()

# 从环境变量获取管理员凭据
ADMIN_USERNAME = os.environ.get("ADMIN_USERNAME", "admin")
ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "antigravity")


# ============ FastAPI 应用 ============
@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    print("🚀 API 反代服务启动中...")
    yield
    print("👋 API 反代服务关闭")


app = FastAPI(
    title="Antigravity API Proxy",
    description="OpenAI 兼容的 API 反代服务",
    version="1.0.0",
    lifespan=lifespan
)

# CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============ API 端点 ============

@app.get("/v1/models")
async def list_models():
    """列出支持的模型"""
    return {
        "object": "list",
        "data": [
            {
                "id": model["id"],
                "object": "model",
                "owned_by": "antigravity",
                "permission": []
            }
            for model in SUPPORTED_MODELS
        ]
    }


@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    """
    OpenAI 兼容的 Chat Completion API
    """
    # 验证 API Key
    auth_header = request.headers.get("Authorization", "")
    if not auth_header.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing API key")
    
    # 验证 API Key 是否匹配
    provided_key = auth_header[7:]  # 去掉 "Bearer " 前缀
    expected_key = config_manager.get_api_key()
    if provided_key != expected_key:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    # 获取请求体
    body = await request.json()
    
    try:
        openai_request = OpenAIChatRequest(
            model=body.get("model", "gemini-2.5-flash"),
            messages=[OpenAIMessage(**msg) for msg in body.get("messages", [])],
            temperature=body.get("temperature", 1.0),
            top_p=body.get("top_p", 0.95),
            max_tokens=body.get("max_tokens", 8192),
            stream=body.get("stream", False)
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid request: {e}")
    
    # 获取可用账号
    account = await load_balancer.get_next_account()
    if not account:
        raise HTTPException(status_code=503, detail="No available accounts")
    
    # 处理请求
    if openai_request.stream:
        return StreamingResponse(
            stream_chat_completion(openai_request, account),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            }
        )
    else:
        result = await chat_completion(openai_request, account)
        if "error" in result:
            raise HTTPException(status_code=500, detail=result["error"])
        return JSONResponse(content=result)


@app.get("/health")
async def health_check():
    """健康检查"""
    stats = account_manager.get_stats()
    return {
        "status": "healthy",
        "accounts": stats.total_accounts,
        "available": stats.available_accounts
    }


@app.post("/admin/import-accounts")
async def import_accounts(request: Request):
    """
    批量导入账号 (需要管理员密码)
    请求体格式:
    {
        "password": "admin_password",
        "accounts": [
            {"email": "xxx@gmail.com", "refresh_token": "1//xxx..."},
            ...
        ]
    }
    """
    body = await request.json()
    
    # 验证管理员密码
    if body.get("password") != ADMIN_PASSWORD:
        raise HTTPException(status_code=401, detail="Invalid admin password")
    
    accounts_data = body.get("accounts", [])
    if not accounts_data:
        raise HTTPException(status_code=400, detail="No accounts provided")
    
    results = []
    for acc in accounts_data:
        email = acc.get("email")
        refresh_token = acc.get("refresh_token")
        
        if not email or not refresh_token:
            results.append({"email": email, "status": "error", "message": "Missing email or refresh_token"})
            continue
        
        try:
            account = account_manager.add_account(
                email=email,
                access_token="pending",  # 将在首次使用时自动刷新
                refresh_token=refresh_token,
                expires_in=0  # 立即过期，强制刷新
            )
            results.append({"email": email, "status": "success", "id": account.id})
        except Exception as e:
            results.append({"email": email, "status": "error", "message": str(e)})
    
    return {
        "imported": len([r for r in results if r["status"] == "success"]),
        "failed": len([r for r in results if r["status"] == "error"]),
        "results": results
    }


# ============ Gradio 管理界面 ============

def get_accounts_table():
    """获取账号列表表格数据"""
    accounts = account_manager.get_all_accounts()
    if not accounts:
        return [["暂无账号", "-", "-", "-", "-", "-"]]
    
    return [
        [
            acc.id,
            acc.email,
            "✅ 正常" if acc.is_available() else "❌ 冷却中",
            str(acc.total_requests),
            f"{acc.successful_requests / acc.total_requests * 100:.1f}%" if acc.total_requests > 0 else "-",
            acc.last_used.strftime("%Y-%m-%d %H:%M") if acc.last_used else "-"
        ]
        for acc in accounts
    ]


def add_account(email: str, access_token: str, refresh_token: str, project_id: str):
    """添加账号"""
    if not email or not access_token or not refresh_token:
        return "❌ 请填写完整信息", get_accounts_table()
    
    try:
        account = account_manager.add_account(
            email=email,
            access_token=access_token,
            refresh_token=refresh_token,
            project_id=project_id if project_id else None
        )
        return f"✅ 账号 {email} 添加成功！", get_accounts_table()
    except Exception as e:
        return f"❌ 添加失败: {e}", get_accounts_table()


def delete_account(email_to_delete: str):
    """删除账号"""
    if not email_to_delete or email_to_delete == "暂无账号":
        return "❌ 请选择要删除的账号", get_accounts_table()
    
    # 查找账号 ID
    accounts = account_manager.get_all_accounts()
    target_account = None
    for acc in accounts:
        if acc.email == email_to_delete:
            target_account = acc
            break
    
    if not target_account:
        return f"❌ 账号 {email_to_delete} 不存在", get_accounts_table()
    
    try:
        if account_manager.remove_account(target_account.id):
            return f"✅ 账号 {email_to_delete} 已删除", get_accounts_table()
        else:
            return f"❌ 删除账号 {email_to_delete} 失败", get_accounts_table()
    except Exception as e:
        return f"❌ 删除失败: {e}", get_accounts_table()


def get_account_emails():
    """获取所有账号邮箱列表 (用于更新 Dropdown)"""
    accounts = account_manager.get_all_accounts()
    if not accounts:
        return gr.update(choices=["暂无账号"], value=None)
    emails = [acc.email for acc in accounts]
    return gr.update(choices=emails, value=None)


def _get_account_emails_list():
    """获取账号邮箱列表 (用于初始化)"""
    accounts = account_manager.get_all_accounts()
    if not accounts:
        return ["暂无账号"]
    return [acc.email for acc in accounts]


def update_api_key(new_api_key: str):
    """更新 API Key"""
    if not new_api_key or len(new_api_key.strip()) < 3:
        return "❌ API Key 不能为空且至少需要 3 个字符", config_manager.get_api_key()
    
    if config_manager.set_api_key(new_api_key.strip()):
        return f"✅ API Key 已更新为: {new_api_key.strip()}", new_api_key.strip()
    else:
        return "❌ 更新 API Key 失败", config_manager.get_api_key()


def get_service_info():
    """获取服务信息"""
    stats = account_manager.get_stats()
    current_api_key = config_manager.get_api_key()
    
    # 获取当前 Space URL
    space_url = os.environ.get("SPACE_HOST", "localhost:7860")
    if not space_url.startswith("http"):
        space_url = f"https://{space_url}"
    
    return f"""
## 🔗 API 配置

**Base URL:** `{space_url}/v1`

**API Key:** `{current_api_key}`

**支持的模型:**
- `gemini-2.5-pro` / `gemini-2.5-flash` / `gemini-2.5-flash-lite`
- `gemini-3-flash` / `gemini-3-pro` / `gemini-3-pro-high`
- `claude-sonnet-4-5-thinking` / `claude-sonnet-4-5` / `claude-opus-4-5`

---

## 📊 服务状态

| 指标 | 值 |
|------|-----|
| 总账号数 | {stats.total_accounts} |
| 可用账号 | {stats.available_accounts} |
| 总请求数 | {stats.total_requests} |
| 成功率 | {stats.success_rate * 100:.1f}% |

---

## 🛠️ 使用方法

### cURL 示例
```bash
curl {space_url}/v1/chat/completions \\
  -H "Content-Type: application/json" \\
  -H "Authorization: Bearer {current_api_key}" \\
  -d '{{
    "model": "gemini-2.5-flash",
    "messages": [{{"role": "user", "content": "Hello"}}],
    "stream": true
  }}'
```

### NextChat / Cherry Studio 配置
1. API 类型: OpenAI
2. Base URL: `{space_url}/v1`
3. API Key: `{current_api_key}`
"""


# 创建 Gradio 界面
with gr.Blocks(
    title="Antigravity API Proxy",
    theme=gr.themes.Soft(
        primary_hue="purple",
        secondary_hue="blue",
    ),
    css="""
    .container { max-width: 1200px; margin: auto; }
    .logo { font-size: 2em; font-weight: bold; }
    .danger-btn { background-color: #dc3545 !important; }
    """
) as demo:
    gr.Markdown(
        """
        # 🚀 Antigravity API Proxy
        ### 将 Gemini/Claude 转换为 OpenAI 兼容 API
        """,
        elem_classes="logo"
    )
    
    with gr.Tabs():
        # 仪表盘
        with gr.Tab("📊 仪表盘"):
            service_info = gr.Markdown(get_service_info)
            refresh_btn = gr.Button("🔄 刷新状态")
            refresh_btn.click(fn=get_service_info, outputs=service_info)
        
        # 账号管理
        with gr.Tab("👤 账号管理"):
            gr.Markdown("### 添加新账号")
            gr.Markdown("> 需要从 Antigravity 桌面应用导出账号信息，或手动获取 OAuth Token")
            
            with gr.Row():
                email_input = gr.Textbox(label="邮箱", placeholder="example@gmail.com")
                project_id_input = gr.Textbox(label="Project ID (选填)", placeholder="可留空")
            
            access_token_input = gr.Textbox(
                label="Access Token", 
                placeholder="ya29.xxx... 或填 pending",
                lines=2
            )
            refresh_token_input = gr.Textbox(
                label="Refresh Token",
                placeholder="1//xxx...",
                lines=2
            )
            
            add_btn = gr.Button("➕ 添加账号", variant="primary")
            result_text = gr.Textbox(label="结果", interactive=False)
            
            gr.Markdown("---")
            gr.Markdown("### 账号列表")
            accounts_table = gr.Dataframe(
                headers=["ID", "邮箱", "状态", "请求数", "成功率", "最后使用"],
                value=get_accounts_table,
                interactive=False
            )
            
            gr.Markdown("### 删除账号")
            with gr.Row():
                delete_email_dropdown = gr.Dropdown(
                    label="选择要删除的账号",
                    choices=_get_account_emails_list(),
                    interactive=True
                )
                delete_btn = gr.Button("🗑️ 删除账号", variant="stop")
            delete_result = gr.Textbox(label="删除结果", interactive=False)
            
            # 绑定事件
            add_btn.click(
                fn=add_account,
                inputs=[email_input, access_token_input, refresh_token_input, project_id_input],
                outputs=[result_text, accounts_table]
            ).then(
                fn=get_account_emails,
                outputs=delete_email_dropdown
            )
            
            delete_btn.click(
                fn=delete_account,
                inputs=[delete_email_dropdown],
                outputs=[delete_result, accounts_table]
            ).then(
                fn=get_account_emails,
                outputs=delete_email_dropdown
            )
        
        # 设置
        with gr.Tab("⚙️ 设置"):
            gr.Markdown("### API Key 配置")
            gr.Markdown("> 修改后，所有 API 请求需要使用新的 API Key")
            
            with gr.Row():
                api_key_input = gr.Textbox(
                    label="当前 API Key",
                    value=config_manager.get_api_key,
                    placeholder="输入新的 API Key",
                    interactive=True
                )
                save_api_key_btn = gr.Button("💾 保存 API Key", variant="primary")
            
            api_key_result = gr.Textbox(label="结果", interactive=False)
            
            save_api_key_btn.click(
                fn=update_api_key,
                inputs=[api_key_input],
                outputs=[api_key_result, api_key_input]
            )
            
            gr.Markdown("---")
            gr.Markdown("### 环境变量配置")
            gr.Markdown("""
            > 以下配置在 HF Spaces 环境下通过 Secrets 设置
            
            | 配置项 | 环境变量 | 默认值 |
            |--------|----------|--------|
            | 数据目录 | `DATA_DIR` | `./data` |
            | 管理员用户名 | `ADMIN_USERNAME` | `admin` |
            | 管理员密码 | `ADMIN_PASSWORD` | `antigravity` |
            """)


# 挂载 Gradio 到 FastAPI (带密码保护)
app = gr.mount_gradio_app(
    app, 
    demo, 
    path="/",
    auth=(ADMIN_USERNAME, ADMIN_PASSWORD),
    auth_message="🔐 请输入管理员凭据登录 Antigravity API Proxy"
)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)