main.py

import asyncio
import json
import sys
import uuid
import base64
import re
import os
import argparse
import time
from datetime import datetime, timezone
from typing import List, Optional

import httpx
import uvicorn
from fastapi import (
    BackgroundTasks,
    FastAPI,
    HTTPException,
    Request,
    Response,
    status,
)
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

from bearer_token import BearerTokenGenerator

from fastapi import Depends, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer

# 模型列表
MODELS = ["gpt-4o", "gpt-4o-mini", "claude-3-5-sonnet", "claude"]

# 默认端口
INITIAL_PORT = 3000

# 外部API的URL
EXTERNAL_API_URL = "https://api.chaton.ai/chats/stream"

# 初始化FastAPI应用
app = FastAPI()

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],  # 允许GET, POST, OPTIONS方法
    allow_headers=["Content-Type", "Authorization"],  # 允许的头部
)

# 挂载静态文件路由以提供 images 目录的内容
app.mount("/images", StaticFiles(directory="images"), name="images")

# 辅助函数
def send_error_response(message: str, status_code: int = 400):
    """构建错误响应，并确保包含CORS头"""
    error_json = {"error": message}
    headers = {
        "Access-Control-Allow-Origin": "*",
        "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
        "Access-Control-Allow-Headers": "Content-Type, Authorization",
    }
    return JSONResponse(status_code=status_code, content=error_json, headers=headers)

def extract_path_from_markdown(markdown: str) -> Optional[str]:
    """
    提取 Markdown 图片链接中的路径，匹配以 https://spc.unk/ 开头的 URL
    """
    pattern = re.compile(r'!\[.*?\]\(https://spc\.unk/(.*?)\)')
    match = pattern.search(markdown)
    if match:
        return match.group(1)
    return None

async def fetch_get_url_from_storage(storage_url: str) -> Optional[str]:
    """
    从 storage URL 获取 JSON 并提取 getUrl
    """
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(storage_url)
            if response.status_code != 200:
                print(f"获取 storage URL 失败，状态码: {response.status_code}")
                return None
            json_response = response.json()
            return json_response.get("getUrl")
        except Exception as e:
            print(f"Error fetching getUrl from storage: {e}")
            return None

async def download_image(image_url: str) -> Optional[bytes]:
    """
    下载图像
    """
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(image_url)
            if response.status_code == 200:
                return response.content
            else:
                print(f"下载图像失败，状态码: {response.status_code}")
                return None
        except Exception as e:
            print(f"Error downloading image: {e}")
            return None

def save_base64_image(base64_str: str, images_dir: str = "images") -> str:
    """
    将Base64编码的图片保存到images目录，返回文件名
    """
    if not os.path.exists(images_dir):
        os.makedirs(images_dir)
    image_data = base64.b64decode(base64_str)
    filename = f"{uuid.uuid4()}.png"  # 默认保存为png格式
    file_path = os.path.join(images_dir, filename)
    with open(file_path, "wb") as f:
        f.write(image_data)
    return filename

def is_base64_image(url: str) -> bool:
    """
    判断URL是否为Base64编码的图片
    """
    return url.startswith("data:image/")

# 添加 HTTPBearer 实例
security = HTTPBearer()

# 添加 API_KEY 验证函数
def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security)):
    api_key = os.environ.get("API_KEY")
    if api_key is None:
        raise HTTPException(status_code=500, detail="API_KEY not set in environment variables")
    if credentials.credentials != api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return credentials.credentials

# 根路径GET请求处理
@app.get("/")
async def root():
    return JSONResponse(content={
        "service": "AI Chat Completion Proxy",
        "usage": {
            "endpoint": "/ai/v1/chat/completions",
            "method": "POST",
            "headers": {
                "Content-Type": "application/json",
                "Authorization": "Bearer YOUR_API_KEY"
            },
            "body": {
                "model": "One of: " + ", ".join(MODELS),
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Hello, who are you?"}
                ],
                "stream": False,
                "temperature": 0.7,
                "max_tokens": 8000
            }
        },
        "availableModels": MODELS,
        "endpoints": {
            "/ai/v1/chat/completions": "Chat completion endpoint",
            "/ai/v1/images/generations": "Image generation endpoint",
            "/ai/v1/models": "List available models"
        },
        "note": "Replace YOUR_API_KEY with your actual API key."
    })

# 返回模型列表
@app.get("/ai/v1/models")
async def list_models():
    """返回可用模型列表。"""
    models = [
        {
            "id": model,
            "object": "model",
            "created": int(time.time()),
            "owned_by": "chaton",
            "permission": [],
            "root": model,
            "parent": None,
        } for model in MODELS
    ]
    return JSONResponse(content={
        "object": "list",
        "data": models
    })

# 聊天完成处理
@app.post("/ai/v1/chat/completions")
async def chat_completions(request: Request, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
    """
    处理聊天完成请求
    """
    try:
        request_body = await request.json()
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON")

    # 打印接收到的请求
    print("Received Completion JSON:", json.dumps(request_body, ensure_ascii=False))

    # 处理消息内容
    messages = request_body.get("messages", [])
    temperature = request_body.get("temperature", 1.0)
    top_p = request_body.get("top_p", 1.0)
    max_tokens = request_body.get("max_tokens", 8000)
    model = request_body.get("model", "gpt-4o")
    is_stream = request_body.get("stream", False)  # 获取 stream 字段

    has_image = False
    has_text = False

    # 清理和提取消息内容
    cleaned_messages = []
    for message in messages:
        content = message.get("content", "")
        if isinstance(content, list):
            text_parts = []
            images = []
            for item in content:
                if "text" in item:
                    text_parts.append(item.get("text", ""))
                elif "image_url" in item:
                    has_image = True
                    image_info = item.get("image_url", {})
                    url = image_info.get("url", "")
                    if is_base64_image(url):
                        # 解码并保存图片
                        base64_str = url.split(",")[1]
                        filename = save_base64_image(base64_str)
                        base_url = app.state.base_url
                        image_url = f"{base_url}/images/{filename}"
                        images.append({"data": image_url})
                    else:
                        images.append({"data": url})
            extracted_content = " ".join(text_parts).strip()
            if extracted_content:
                has_text = True
                message["content"] = extracted_content
                if images:
                    message["images"] = images
                cleaned_messages.append(message)
                print("Extracted:", extracted_content)
            else:
                if images:
                    has_image = True
                    message["content"] = ""
                    message["images"] = images
                    cleaned_messages.append(message)
                    print("Extracted image only.")
                else:
                    print("Deleted message with empty content.")
        elif isinstance(content, str):
            content_str = content.strip()
            if content_str:
                has_text = True
                message["content"] = content_str
                cleaned_messages.append(message)
                print("Retained content:", content_str)
            else:
                print("Deleted message with empty content.")
        else:
            print("Deleted non-expected type of content message.")

    if not cleaned_messages:
        raise HTTPException(status_code=400, detail="所有消息的内容均为空。")

    # 验证模型
    if model not in MODELS:
        model = "gpt-4o"

    # 构建新的请求JSON
    new_request_json = {
        "function_image_gen": False,
        "function_web_search": True,
        "max_tokens": max_tokens,
        "model": model,
        "source": "chat/free",
        "temperature": temperature,
        "top_p": top_p,
        "messages": cleaned_messages,
    }

    modified_request_body = json.dumps(new_request_json, ensure_ascii=False)
    print("Modified Request JSON:", modified_request_body)

    # 获取Bearer Token
    tmp_token = BearerTokenGenerator.get_bearer(modified_request_body)
    if not tmp_token:
        raise HTTPException(status_code=500, detail="无法生成 Bearer Token")

    bearer_token, formatted_date = tmp_token

    headers = {
        "Date": formatted_date,
        "Client-time-zone": "-05:00",
        "Authorization": bearer_token,
        "User-Agent": "ChatOn_Android/1.53.502",
        "Accept-Language": "en-US",
        "X-Cl-Options": "hb",
        "Content-Type": "application/json; charset=UTF-8",
    }

    if is_stream:
        # 流式响应处理
        async def event_generator():
            async with httpx.AsyncClient(timeout=None) as client_stream:
                try:
                    async with client_stream.stream("POST", EXTERNAL_API_URL, headers=headers, content=modified_request_body) as streamed_response:
                        async for line in streamed_response.aiter_lines():
                            if line.startswith("data: "):
                                data = line[6:].strip()
                                if data == "[DONE]":
                                    # 通知客户端流结束
                                    yield "data: [DONE]\n\n"
                                    break
                                try:
                                    sse_json = json.loads(data)
                                    if "choices" in sse_json:
                                        for choice in sse_json["choices"]:
                                            delta = choice.get("delta", {})
                                            content = delta.get("content")
                                            if content:
                                                new_sse_json = {
                                                    "choices": [
                                                        {
                                                            "index": choice.get("index", 0),
                                                            "delta": {"content": content},
                                                        }
                                                    ],
                                                    "created": sse_json.get(
                                                        "created", int(datetime.now(timezone.utc).timestamp())
                                                    ),
                                                    "id": sse_json.get(
                                                        "id", str(uuid.uuid4())
                                                    ),
                                                    "model": sse_json.get("model", "gpt-4o"),
                                                    "system_fingerprint": f"fp_{uuid.uuid4().hex[:12]}",
                                                }
                                                new_sse_line = f"data: {json.dumps(new_sse_json, ensure_ascii=False)}\n\n"
                                                yield new_sse_line
                                except json.JSONDecodeError:
                                    print("JSON解析错误")
                                    continue
                except httpx.RequestError as exc:
                    print(f"外部API请求失败: {exc}")
                    yield f"data: {{\"error\": \"外部API请求失败: {str(exc)}\"}}\n\n"

        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                # CORS头已通过中间件处理，无需在这里重复添加
            },
        )
    else:
        # 非流式响应处理
        async with httpx.AsyncClient(timeout=None) as client:
            try:
                response = await client.post(
                    EXTERNAL_API_URL,
                    headers=headers,
                    content=modified_request_body,
                    timeout=None
                )

                if response.status_code != 200:
                    raise HTTPException(
                        status_code=response.status_code,
                        detail=f"API 错误: {response.status_code}",
                    )

                sse_lines = response.text.splitlines()
                content_builder = ""
                images_urls = []

                for line in sse_lines:
                    if line.startswith("data: "):
                        data = line[6:].strip()
                        if data == "[DONE]":
                            break
                        try:
                            sse_json = json.loads(data)
                            if "choices" in sse_json:
                                for choice in sse_json["choices"]:
                                    if "delta" in choice:
                                        delta = choice["delta"]
                                        if "content" in delta:
                                            content_builder += delta["content"]
                        except json.JSONDecodeError:
                            print("JSON解析错误")
                            continue

                openai_response = {
                    "id": f"chatcmpl-{uuid.uuid4()}",
                    "object": "chat.completion",
                    "created": int(datetime.now(timezone.utc).timestamp()),
                    "model": model,
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": content_builder,
                            },
                            "finish_reason": "stop",
                        }
                    ],
                }

                # 处理图片（如果有）
                if has_image:
                    images = []
                    for message in cleaned_messages:
                        if "images" in message:
                            for img in message["images"]:
                                images.append({"data": img["data"]})
                    openai_response["choices"][0]["message"]["images"] = images

                return JSONResponse(content=openai_response, status_code=200)
            except httpx.RequestError as exc:
                raise HTTPException(status_code=500, detail=f"请求失败: {str(exc)}")
            except Exception as exc:
                raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(exc)}")

# 图像生成处理
@app.post("/ai/v1/images/generations")
async def images_generations(request: Request, api_key: str = Depends(verify_api_key)):
    """
    处理图像生成请求
    """
    try:
        request_body = await request.json()
    except json.JSONDecodeError:
        return send_error_response("Invalid JSON", status_code=400)

    print("Received Image Generations JSON:", json.dumps(request_body, ensure_ascii=False))

    # 验证必需的字段
    if "prompt" not in request_body:
        return send_error_response("缺少必需的字段: prompt", status_code=400)

    user_prompt = request_body.get("prompt", "").strip()
    response_format = request_body.get("response_format", "b64_json").strip()

    if not user_prompt:
        return send_error_response("Prompt 不能为空。", status_code=400)

    print(f"Prompt: {user_prompt}")

    # 构建新的 TextToImage JSON 请求体
    text_to_image_json = {
        "function_image_gen": True,
        "function_web_search": True,
        "image_aspect_ratio": "1:1",
        "image_style": "photographic",  # 暂时固定 image_style
        "max_tokens": 8000,
        "messages": [
            {
                "content": "You are a helpful artist, please based on imagination draw a picture.",
                "role": "system"
            },
            {
                "content": "Draw: " + user_prompt,
                "role": "user"
            }
        ],
        "model": "gpt-4o",  # 固定 model，只能gpt-4o或gpt-4o-mini
        "source": "chat/pro_image"  # 固定 source
    }

    modified_request_body = json.dumps(text_to_image_json, ensure_ascii=False)
    print("Modified Request JSON:", modified_request_body)

    # 获取Bearer Token
    tmp_token = BearerTokenGenerator.get_bearer(modified_request_body, path="/chats/stream")
    if not tmp_token:
        return send_error_response("无法生成 Bearer Token", status_code=500)

    bearer_token, formatted_date = tmp_token

    headers = {
        "Date": formatted_date,
        "Client-time-zone": "-05:00",
        "Authorization": bearer_token,
        "User-Agent": "ChatOn_Android/1.53.502",
        "Accept-Language": "en-US",
        "X-Cl-Options": "hb",
        "Content-Type": "application/json; charset=UTF-8",
    }

    async with httpx.AsyncClient(timeout=None) as client:
        try:
            response = await client.post(
                EXTERNAL_API_URL, headers=headers, content=modified_request_body, timeout=None
            )
            if response.status_code != 200:
                return send_error_response(f"API 错误: {response.status_code}", status_code=500)

            # 初始化用于拼接 URL 的字符串
            url_builder = ""

            # 读取 SSE 流并拼接 URL
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:].strip()
                    if data == "[DONE]":
                        break
                    try:
                        sse_json = json.loads(data)
                        if "choices" in sse_json:
                            for choice in sse_json["choices"]:
                                delta = choice.get("delta", {})
                                content = delta.get("content")
                                if content:
                                    url_builder += content
                    except json.JSONDecodeError:
                        print("JSON解析错误")
                        continue

            image_markdown = url_builder
            # Step 1: 检查Markdown文本是否为空
            if not image_markdown:
                print("无法从 SSE 流中构建图像 Markdown。")
                return send_error_response("无法从 SSE 流中构建图像 Markdown。", status_code=500)

            # Step 2, 3, 4, 5: 处理图像
            extracted_path = extract_path_from_markdown(image_markdown)
            if not extracted_path:
                print("无法从 Markdown 中提取路径。")
                return send_error_response("无法从 Markdown 中提取路径。", status_code=500)

            print(f"提取的路径: {extracted_path}")

            # Step 5: 拼接最终的存储URL
            storage_url = f"https://api.chaton.ai/storage/{extracted_path}"
            print(f"存储URL: {storage_url}")

            # 获取最终下载URL
            final_download_url = await fetch_get_url_from_storage(storage_url)
            if not final_download_url:
                return send_error_response("无法从 storage URL 获取最终下载链接。", status_code=500)

            print(f"Final Download URL: {final_download_url}")

            # 下载图像
            image_bytes = await download_image(final_download_url)
            if not image_bytes:
                return send_error_response("无法从 URL 下载图像。", status_code=500)

            # 转换为 Base64
            image_base64 = base64.b64encode(image_bytes).decode('utf-8')

            # 将图片保存到images目录并构建可访问的URL
            filename = save_base64_image(image_base64)
            base_url = app.state.base_url
            accessible_url = f"{base_url}/images/{filename}"

            # 根据 response_format 返回相应的响应
            if response_format.lower() == "b64_json":
                response_json = {
                    "data": [
                        {
                            "b64_json": image_base64
                        }
                    ]
                }
                return JSONResponse(content=response_json, status_code=200)
            else:
                # 构建包含可访问URL的响应
                response_json = {
                    "data": [
                        {
                            "url": accessible_url
                        }
                    ]
                }
                return JSONResponse(content=response_json, status_code=200)
        except httpx.RequestError as exc:
            print(f"请求失败: {exc}")
            return send_error_response(f"请求失败: {str(exc)}", status_code=500)
        except Exception as exc:
            print(f"内部服务器错误: {exc}")
            return send_error_response(f"内部服务器错误: {str(exc)}", status_code=500)

# 运行服务器
def main():
    parser = argparse.ArgumentParser(description="启动ChatOn API服务器")
    parser.add_argument('--base_url', type=str, default='http://localhost', help='Base URL for accessing images')
    parser.add_argument('--port', type=int, default=INITIAL_PORT, help='服务器监听端口')
    args = parser.parse_args()
    base_url = args.base_url
    port = args.port

    # 检查 API_KEY 是否设置
    if not os.environ.get("API_KEY"):
        print("警告: API_KEY 环境变量未设置。客户端验证将无法正常工作。")

    # 确保 images 目录存在
    if not os.path.exists("images"):
        os.makedirs("images")

    # 设置 FastAPI 应用的 state
    app.state.base_url = base_url

    print(f"Server started on port {port} with base_url: {base_url}")

    # 运行FastAPI应用
    uvicorn.run(app, host="0.0.0.0", port=port)

async def get_available_port(start_port: int = INITIAL_PORT, end_port: int = 65535) -> int:
    """查找可用的端口号"""
    for port in range(start_port, end_port + 1):
        try:
            server = await asyncio.start_server(lambda r, w: None, host="0.0.0.0", port=port)
            server.close()
            await server.wait_closed()
            return port
        except OSError:
            continue
    raise RuntimeError(f"No available ports between {start_port} and {end_port}")

if __name__ == "__main__":
    main()