const Koa = require('koa')
const Router = require('koa-router')
const bodyParser = require('koa-bodyparser')
const models = require("./models")
const crypto = require('crypto')

const app = new Koa()
const router = new Router()

// 使用 bodyParser 中间件
app.use(bodyParser())

// 配置 bodyParser
app.use(bodyParser({
  enableTypes: ['json', 'form', 'text'],
  jsonLimit: '30mb',  
  formLimit: '30mb',  
  textLimit: '30mb',  
}))

// 添加 CORS 中间件
app.use(async (ctx, next) => {
  ctx.set('Access-Control-Allow-Origin', '*');
  ctx.set('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization');
  ctx.set('Access-Control-Allow-Methods', 'POST, GET, PUT, DELETE, OPTIONS');
  ctx.set('Access-Control-Max-Age', '3600');
  
  if (ctx.method === 'OPTIONS') {
    ctx.status = 204;
    return;
  }
  
  await next();
});

// 添加请求日志中间件
app.use(async (ctx, next) => {
  const start = Date.now()
  console.log(`[${new Date().toISOString()}] ${ctx.method} ${ctx.url} - Request started`)
  try {
    await next()
  } catch (err) {
    throw err
  } finally {
    const ms = Date.now() - start
    console.log(`[${new Date().toISOString()}] ${ctx.method} ${ctx.url} - Response ${ctx.status} - ${ms}ms`)
  }
})

// 错误处理中间件
app.use(async (ctx, next) => {
  try {
    await next()
    if (ctx.status === 404 && !ctx.body) {
      ctx.status = 404
      ctx.body = {
        error: 'Not Found',
        message: `The requested path ${ctx.path} was not found`,
        availableEndpoints: [
          '/v1/chat/completions',
          '/hf/v1/chat/completions',
          '/v1/models',
          '/hf/v1/models'
        ]
      }
    }
  } catch (err) {
    console.error('Error details:', {
      path: ctx.path,
      method: ctx.method,
      error: err.message,
      stack: err.stack
    })
    ctx.status = err.status || 500
    ctx.body = {
      error: err.name || 'Internal Server Error',
      message: err.message,
      path: ctx.path
    }
    ctx.app.emit('error', err, ctx)
  }
})

const makeRequest = async (session_id, requestModel, messages) => {
  console.log('开始请求:', { session_id: '***', model: requestModel })
  try {
    const myHeaders = new Headers()
    myHeaders.append("Cookie", `session_id=${session_id}`)
    myHeaders.append("User-Agent", "Apifox/1.0.0 (https://apifox.com)");
    myHeaders.append("Content-Type", "application/json")
    myHeaders.append("Accept", "*/*")
    myHeaders.append("Host", "www.genspark.ai")
    myHeaders.append("Connection", "keep-alive")

    const body = JSON.stringify({
      "type": "COPILOT_MOA_CHAT",
      "current_query_string": "type=COPILOT_MOA_CHAT",
      "messages": messages,
      "action_params": {},
      "extra_data": {
        "models": [
          models[`${requestModel}`] || models["claude-3-5-sonnet-20241022"]
        ],
        "run_with_another_model": false,
        "writingContent": null
      }
    })

    const requestConfig = {
      method: 'POST',
      headers: myHeaders,
      body: body,
      redirect: 'follow'
    };

    const response = await fetch("https://www.genspark.ai/api/copilot/ask", requestConfig)
    
    if (!response.ok) {
      throw new Error(`HTTP error! status: ${response.status}`)
    }
    
    // 验证响应是否有效
    if (!response.body) {
      throw new Error('Invalid response: body is missing')
    }

    return response
  } catch (error) {
    console.error('请求出错:', error)
    // 添加重试逻辑
    if (error.message.includes('HTTP error! status: 500')) {
      console.log('尝试重新请求...')
      await new Promise(resolve => setTimeout(resolve, 1000)) // 等待1秒后重试
      return makeRequest(session_id, requestModel, messages)
    }
    throw error
  }
}

const parseStreamData = (data) => {
  try {
    if (!data) return null
    const cleanData = data.replace("data: ", '').trim()
    if (!cleanData || cleanData === '[DONE]') return null
    return JSON.parse(cleanData)
  } catch (error) {
    console.error('解析数据出错:', error, '原始数据:', data)
    return null
  }
}

const handleChatCompletions = async (ctx) => {
  const { messages, stream = false, model = 'claude-3-5-sonnet' } = ctx.request.body
  const session_id = ctx.get('Authorization')?.replace('Bearer ', '')

  console.log('收到请求:', {
    path: ctx.path,
    model,
    stream,
    messages_count: messages?.length
  })

  if (!session_id) {
    ctx.status = 401
    ctx.body = { error: '未提供有效的 session_id' }
    return
  }

  try {
    const response = await makeRequest(session_id, model, messages)

    const messageId = crypto.randomUUID()
    const reader = response.body.getReader()
    let resBody = {
      id: `chatcmpl-${messageId}`,
      object: 'chat.completion',
      created: Math.floor(Date.now() / 1000),
      model,
      choices: [],
      usage: {
        prompt_tokens: 0,
        completion_tokens: 0,
        total_tokens: 0,
      },
    }

    if (stream == "true" || stream == true) {
      ctx.status = 200
      ctx.set({
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive',
      })
      
      // 发送初始消息
      ctx.res.write(`data: ${JSON.stringify({
        "id": `chatcmpl-${messageId}`,
        "choices": [
          {
            "index": 0,
            "delta": {
              "content": "",
              "role": "assistant"
            }
          }
        ],
        "created": Math.floor(Date.now() / 1000),
        "model": model,
        "object": "chat.completion.chunk"
      })}\n\n`)
    } else {
      ctx.status = 200
      ctx.set({
        'Content-Type': 'application/json',
      })
    }

    try {
      while (true) {
        const { done, value } = await reader.read()
        if (done) {
          if (stream == "true" || stream == true) {
            ctx.res.write('data: [DONE]\n\n')
          }
          break
        }

        const text = new TextDecoder().decode(value)
        const textContent = [...text.matchAll(/data:.*?(?=(\n\ndata:|$))/gs)]

        for (const item of textContent) {
          if (!item[0]) continue
          
          const content = parseStreamData(item[0])
          if (!content) continue

          if (stream == "true" || stream == true) {
            if (content.delta) {
              try {
                ctx.res.write(`data: ${JSON.stringify({
                  "id": `chatcmpl-${messageId}`,
                  "choices": [
                    {
                      "index": 0,
                      "delta": {
                        "content": content.delta
                      }
                    }
                  ],
                  "created": Math.floor(Date.now() / 1000),
                  "model": model,
                  "object": "chat.completion.chunk"
                })}\n\n`)
              } catch (e) {
                console.error('写入流数据失败:', e)
              }
            }
          } else {
            if (content.field_value && 
                content.field_name !== 'session_state.answer_is_finished' && 
                content.field_name !== 'content' && 
                content.field_name !== 'session_state' && 
                !content.delta && 
                content.type !== 'project_field') {
              resBody.choices.push({
                index: 0,
                message: {
                  role: 'assistant',
                  content: content.field_value,
                },
                finish_reason: 'stop',
              })
              resBody.usage.total_tokens = content.field_value.length
            }
          }
        }
      }

      if (stream == "false" || stream == false) {
        if (resBody.choices.length === 0) {
          resBody.choices.push({
            index: 0,
            message: {
              role: 'assistant',
              content: '抱歉，我无法处理您的请求。',
            },
            finish_reason: 'stop',
          })
        }
        ctx.body = resBody
      } else {
        ctx.res.end()
      }
    } catch (error) {
      console.error('响应处理出错:', error)
      if (stream == "true" || stream == true) {
        try {
          ctx.res.write(`data: ${JSON.stringify({
            "id": `chatcmpl-${messageId}`,
            "choices": [
              {
                "index": 0,
                "delta": {
                  "content": "\n\n抱歉，处理您的请求时出现错误。"
                }
              }
            ],
            "created": Math.floor(Date.now() / 1000),
            "model": model,
            "object": "chat.completion.chunk"
          })}\n\n`)
          ctx.res.write('data: [DONE]\n\n')
        } catch (e) {
          console.error('发送错误信息失败:', e)
        } finally {
          ctx.res.end()
        }
      } else {
        ctx.status = 500
        ctx.body = { error: '响应处理失败', details: error.toString() }
      }
    }
  } catch (error) {
    console.error('请求处理出错:', error)
    ctx.status = error.status || 500
    ctx.body = { 
      error: error.message || '请求处理失败',
      details: error.toString()
    }
  }
}

// 创建获取模型列表处理函数
const handleModels = async (ctx) => {
  ctx.body = {
    object: "list",
    data: Object.keys(models).map(model => ({
      id: model,
      object: "model",
      created: 1706745938,
      owned_by: "genspark"
    }))
  }
}

// 注册路由
router.post('/v1/chat/completions', handleChatCompletions)
router.post('/hf/v1/chat/completions', handleChatCompletions)
router.get('/v1/models', handleModels)
router.get('/hf/v1/models', handleModels)

router.get('/', async (ctx) => {
  ctx.body = {
    status: 'ok',
    message: 'API server is running',
    version: '1.0.0',
    endpoints: [
      '/v1/chat/completions',
      '/hf/v1/chat/completions', 
      '/v1/models',
      '/hf/v1/models'
    ]
  }
})

app.use(router.routes()).use(router.allowedMethods())

const PORT = process.env.PORT || 3000
app.listen(PORT, '0.0.0.0', () => {
  console.log('='.repeat(50))
  console.log(`服务器启动于 http://0.0.0.0:${PORT}`)
  console.log('环境变量:', {
    PORT: process.env.PORT,
    NODE_ENV: process.env.NODE_ENV
  })
  console.log('='.repeat(50))
});