const express = require('express');
const fetch = require('node-fetch');

const app = express();
const PORT = process.env.PORT || 4777;
const A2BASE_API_URL = 'https://a2abase.ai/api';

app.use(express.json());

// 模型列表
const MODELS_DATA = {
  "models": [
    {
      "id": "gemini/gemini-ultra",
      "display_name": "Gemini Ultra",
      "short_name": "Gemini Ultra",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 5.0,
      "output_cost_per_million_tokens": 20.0,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output", "thinking"],
      "recommended": false,
      "priority": 106
    },
    {
      "id": "gemini/gemini-3-deep-think-exp",
      "display_name": "Gemini 3 Deep Think",
      "short_name": "Gemini 3 Deep Think",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 2.5,
      "output_cost_per_million_tokens": 10.0,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling", "vision", "thinking"],
      "recommended": false,
      "priority": 105
    },
    {
      "id": "gemini/gemini-3-flash-preview",
      "display_name": "Gemini 3 Flash",
      "short_name": "Gemini 3 Flash",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.3,
      "output_cost_per_million_tokens": 1.5,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 104
    },
    {
      "id": "gemini/gemini-3-pro-preview",
      "display_name": "Gemini 3",
      "short_name": "Gemini 3",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.36,
      "output_cost_per_million_tokens": 3.0,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 103
    },
    {
      "id": "anthropic/claude-sonnet-4-5-20250929",
      "display_name": "Claude Sonnet 4.5",
      "short_name": "Claude Sonnet 4.5",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 3.6,
      "output_cost_per_million_tokens": 18.0,
      "context_window": 200000,
      "capabilities": ["chat", "function_calling", "vision", "thinking", "structured_output"],
      "recommended": false,
      "priority": 102
    },
    {
      "id": "gemini/gemini-2.5-flash",
      "display_name": "Gemini 2.5 Flash",
      "short_name": "Gemini 2.5 Flash",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.36,
      "output_cost_per_million_tokens": 3.0,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 102
    },
    {
      "id": "anthropic/claude-haiku-4-5-20251001",
      "display_name": "Claude Haiku 4.5",
      "short_name": "Claude Haiku 4.5",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 1.0,
      "output_cost_per_million_tokens": 5.0,
      "context_window": 200000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 101
    },
    {
      "id": "gemini/gemini-2.5-pro",
      "display_name": "Gemini 2.5 Pro",
      "short_name": "Gemini 2.5 Pro",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 1.25,
      "output_cost_per_million_tokens": 5.0,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output", "thinking"],
      "recommended": false,
      "priority": 101
    },
    {
      "id": "openrouter/z-ai/glm-4.7",
      "display_name": "GLM 4.7",
      "short_name": "GLM 4.7",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.5,
      "output_cost_per_million_tokens": 2.2,
      "context_window": 205000,
      "capabilities": ["chat", "function_calling", "code_interpreter", "thinking"],
      "recommended": false,
      "priority": 100
    },
    {
      "id": "openrouter/x-ai/grok-code-fast-1",
      "display_name": "Grok Code Fast 1",
      "short_name": "Grok Code Fast 1",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.24,
      "output_cost_per_million_tokens": 1.8,
      "context_window": 256000,
      "capabilities": ["chat", "function_calling", "code_interpreter"],
      "recommended": false,
      "priority": 100
    },
    {
      "id": "openrouter/z-ai/glm-4.6",
      "display_name": "GLM 4.6",
      "short_name": "GLM 4.6",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.48,
      "output_cost_per_million_tokens": 2.1,
      "context_window": 202752,
      "capabilities": ["chat", "function_calling"],
      "recommended": false,
      "priority": 99
    },
    {
      "id": "openrouter/minimax/minimax-m2",
      "display_name": "Minimax M2",
      "short_name": "Minimax M2",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.306,
      "output_cost_per_million_tokens": 1.224,
      "context_window": 204800,
      "capabilities": ["chat", "function_calling", "vision"],
      "recommended": false,
      "priority": 98
    },
    {
      "id": "openai/gpt-5.2",
      "display_name": "GPT-5.2",
      "short_name": "GPT-5.2",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 5.0,
      "output_cost_per_million_tokens": 15.0,
      "context_window": 400000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output", "thinking"],
      "recommended": true,
      "priority": 97
    },
    {
      "id": "openai/gpt-5-mini",
      "display_name": "GPT-5 Mini",
      "short_name": "GPT-5 Mini",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.3,
      "output_cost_per_million_tokens": 2.4,
      "context_window": 400000,
      "capabilities": ["chat", "function_calling", "structured_output"],
      "recommended": false,
      "priority": 96
    },
    {
      "id": "openrouter/deepseek/deepseek-chat",
      "display_name": "DeepSeek Chat",
      "short_name": "DeepSeek Chat",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.46,
      "output_cost_per_million_tokens": 1.07,
      "context_window": 128000,
      "capabilities": ["chat", "function_calling"],
      "recommended": false,
      "priority": 95
    },
    {
      "id": "openai/o3-mini",
      "display_name": "o3-mini",
      "short_name": "o3-mini",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 3.5,
      "output_cost_per_million_tokens": 14.0,
      "context_window": 128000,
      "capabilities": ["chat", "thinking"],
      "recommended": false,
      "priority": 95
    },
    {
      "id": "openai/gpt-4.1-mini",
      "display_name": "GPT-4.1 Mini",
      "short_name": "GPT-4.1 Mini",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.48,
      "output_cost_per_million_tokens": 1.92,
      "context_window": 1047576,
      "capabilities": ["chat", "function_calling", "structured_output"],
      "recommended": false,
      "priority": 94
    },
    {
      "id": "openrouter/x-ai/grok-4-fast",
      "display_name": "Grok 4 Fast",
      "short_name": "Grok 4 Fast",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.24,
      "output_cost_per_million_tokens": 0.6,
      "context_window": 2000000,
      "capabilities": ["chat", "function_calling"],
      "recommended": false,
      "priority": 93
    },
    {
      "id": "openai/gpt-4o-2024-11-20",
      "display_name": "GPT-4o (Nov 2024)",
      "short_name": "GPT-4o (Nov 2024)",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 2.5,
      "output_cost_per_million_tokens": 10.0,
      "context_window": 128000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 91
    },
    {
      "id": "openai/gpt-4o",
      "display_name": "GPT-4o",
      "short_name": "GPT-4o",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 2.5,
      "output_cost_per_million_tokens": 10.0,
      "context_window": 128000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 90
    },
    {
      "id": "gemini/gemini-2.0-flash-exp",
      "display_name": "Gemini 2.0 Flash",
      "short_name": "Gemini 2.0 Flash",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.075,
      "output_cost_per_million_tokens": 0.3,
      "context_window": 1000000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 89
    },
    {
      "id": "openai/gpt-4o-mini",
      "display_name": "GPT-4o Mini",
      "short_name": "GPT-4o Mini",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.18,
      "output_cost_per_million_tokens": 0.72,
      "context_window": 128000,
      "capabilities": ["chat", "function_calling", "vision", "structured_output"],
      "recommended": false,
      "priority": 88
    },
    {
      "id": "gemini/gemini-2.0-flash-lite",
      "display_name": "Gemini 2.0 Flash-Lite",
      "short_name": "Gemini 2.0 Flash-Lite",
      "requires_subscription": false,
      "input_cost_per_million_tokens": 0.0375,
      "output_cost_per_million_tokens": 0.15,
      "context_window": 1000000,
      "capabilities": ["chat", "function_calling", "vision"],
      "recommended": false,
      "priority": 86
    }
  ],
  "subscription_tier": "free",
  "total_models": 23,
  "allowed_models_count": 23
};

// GET /v1/models - 返回模型列表
app.get('/v1/models', (req, res) => {
  res.json({
    object: 'list',
    data: MODELS_DATA.models.map(model => ({
      id: model.id,
      object: 'model',
      created: Date.now(),
      owned_by: 'a2base',
      permission: [],
      root: model.id,
      parent: null
    }))
  });
});

// POST /v1/chat/completions - 处理聊天请求
app.post('/v1/chat/completions', async (req, res) => {
  try {
    // 从Authorization header中提取apiKey和agentId
    const authHeader = req.headers.authorization;
    if (!authHeader || !authHeader.startsWith('Bearer ')) {
      return res.status(401).json({ error: { message: 'Missing or invalid Authorization header', type: 'invalid_request_error' } });
    }

    const token = authHeader.substring(7); // 去除 'Bearer '
    const [apiKey, agentId] = token.split('|');

    if (!apiKey || !agentId) {
      return res.status(401).json({ error: { message: 'Invalid token format. Expected: apiKey|agentId', type: 'invalid_request_error' } });
    }

    const { messages, model, stream = false, temperature, max_tokens } = req.body;

    if (!messages || !Array.isArray(messages) || messages.length === 0) {
      return res.status(400).json({ error: { message: 'Messages array is required', type: 'invalid_request_error' } });
    }

    // 步骤1: 创建Thread
    const threadResponse = await fetch(`${A2BASE_API_URL}/threads`, {
      method: 'POST',
      headers: {
        'X-API-Key': apiKey
      }
    });

    if (!threadResponse.ok) {
      const errorText = await threadResponse.text();
      return res.status(threadResponse.status).json({ error: { message: `Failed to create thread: ${errorText}`, type: 'api_error' } });
    }

    const threadData = await threadResponse.json();
    const threadId = threadData.thread_id;

    // 步骤2: 添加消息到Thread
    await fetch(`${A2BASE_API_URL}/threads/${threadId}/messages`, {
      method: 'POST',
      headers: {
        'X-API-Key': apiKey,
        'Content-Type': 'application/json'
      },
      body: JSON.stringify({
        content: messages.map(msg => `${msg.role === 'user' ? 'User' : 'Assistant'}: ${msg.content}`).join('\n'),
        type: 'user',
        is_llm_message: true
      })
    });

    // 步骤3: 启动Agent执行
    const startResponse = await fetch(`${A2BASE_API_URL}/thread/${threadId}/agent/start`, {
      method: 'POST',
      headers: {
        'X-API-Key': apiKey,
        'Content-Type': 'application/json'
      },
      body: JSON.stringify({
        agent_id: agentId,
        model_name: model || 'gemini/gemini-3-pro-preview',
        stream: true
      })
    });

    if (!startResponse.ok) {
      const errorText = await startResponse.text();
      return res.status(startResponse.status).json({ error: { message: `Failed to start agent: ${errorText}`, type: 'api_error' } });
    }

    const startData = await startResponse.json();
    const agentRunId = startData.agent_run_id;

    // 步骤4: 获取流式响应
    const streamResponse = await fetch(`${A2BASE_API_URL}/agent-run/${agentRunId}/stream`, {
      method: 'GET',
      headers: {
        'X-API-Key': apiKey,
        'Accept': 'text/event-stream'
      }
    });

    if (!streamResponse.ok) {
      const errorText = await streamResponse.text();
      return res.status(streamResponse.status).json({ error: { message: `Failed to get stream: ${errorText}`, type: 'api_error' } });
    }

    const chatId = `chatcmpl-${Date.now()}`;
    const created = Math.floor(Date.now() / 1000);

    if (stream) {
      // 流式返回
      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');

      let buffer = '';

      streamResponse.body.on('data', (chunk) => {
        // console.log('Received chunk: ', chunk.toString());
        buffer += chunk.toString();
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const jsonStr = line.substring(6);
            try {
              const data = JSON.parse(jsonStr);

              // 处理assistant类型的消息
              if (data.type === 'assistant' && data.content) {
                const contentObj = JSON.parse(data.content);
                const metaObj = JSON.parse(data.metadata || '{}');
                if (contentObj.content && metaObj.stream_status === 'chunk') {
                  const openaiChunk = {
                    id: chatId,
                    object: 'chat.completion.chunk',
                    created: created,
                    model: model || 'gemini/gemini-2.5-pro',
                    choices: [{
                      index: 0,
                      delta: {
                        content: contentObj.content
                      },
                      finish_reason: null
                    }]
                  };
                  res.write(`data: ${JSON.stringify(openaiChunk)}\n\n`);
                }
              }

              // 处理完成状态
              if (data.status === 'completed' || data.type === 'assistant_response_end') {
                // const finalChunk = {
                //   id: chatId,
                //   object: 'chat.completion.chunk',
                //   created: created,
                //   model: model || 'gemini/gemini-2.5-pro',
                //   choices: [{
                //     index: 0,
                //     delta: {},
                //     finish_reason: 'stop'
                //   }]
                // };
                // res.write(`data: ${JSON.stringify(finalChunk)}\n\n`);
                res.write('data: [DONE]\n\n');
                res.end();
              }
            } catch (e) {
              // 忽略解析错误
            }
          }
        }
      });

      streamResponse.body.on('end', () => {
        if (!res.writableEnded) {
          res.end();
        }
      });

      streamResponse.body.on('error', (err) => {
        console.error('Stream error:', err);
        if (!res.writableEnded) {
          res.end();
        }
      });

    } else {
      // 非流式返回 - 收集所有内容后一次性返回
      let fullContent = '';
      let buffer = '';

      await new Promise((resolve, reject) => {
        streamResponse.body.on('data', (chunk) => {
          buffer += chunk.toString();
          const lines = buffer.split('\n');
          buffer = lines.pop() || '';

          for (const line of lines) {
            if (line.startsWith('data: ')) {
              const jsonStr = line.substring(6);
              try {
                const data = JSON.parse(jsonStr);

                if (data.type === 'assistant' && data.content) {
                  const contentObj = JSON.parse(data.content);
                  const metaObj = JSON.parse(data.metadata || '{}');
                  if (contentObj.content && metaObj.stream_status === 'chunk') {
                    fullContent += contentObj.content;
                  }
                }

                if (data.status === 'completed' || data.type === 'assistant_response_end') {
                  resolve();
                }
              } catch (e) {
                // 忽略解析错误
              }
            }
          }
        });

        streamResponse.body.on('end', resolve);
        streamResponse.body.on('error', reject);
      });

      // 返回OpenAI格式的响应
      const response = {
        id: chatId,
        object: 'chat.completion',
        created: created,
        model: model || 'gemini/gemini-2.5-pro',
        choices: [{
          index: 0,
          message: {
            role: 'assistant',
            content: fullContent
          },
          finish_reason: 'stop'
        }],
        usage: {
          prompt_tokens: 0,
          completion_tokens: 0,
          total_tokens: 0
        }
      };

      res.json(response);
    }

  } catch (error) {
    console.error('Error:', error);
    res.status(500).json({ error: { message: error.message, type: 'internal_error' } });
  }
});

// 健康检查端点
app.get('/health', (req, res) => {
  res.json({ status: 'ok' });
});

app.listen(PORT, () => {
  console.log(`A2Base to OpenAI API server running on port ${PORT}`);
});