import express from 'express';
import fetch from 'node-fetch';
import dotenv from 'dotenv';
import cors from 'cors';
import { launch } from 'puppeteer';
import { v4 as uuidv4 } from 'uuid';
import path from 'path';
import fs from 'fs';

dotenv.config();
// 配置常量
const CONFIG = {
    MODELS: {
        'grok-latest': 'grok-latest',
        'grok-latest-image': 'grok-latest'
    },
    API: {
        BASE_URL: "https://grok.com",
        API_KEY: process.env.API_KEY || "sk-123456",
        SSO_TOKEN: null,//登录时才有的认证cookie,这里暂时用不到，之后可能需要
        SIGNATURE_COOKIE: null
    },
    SERVER: {
        PORT: process.env.PORT || 3000,
        BODY_LIMIT: '5mb'
    },
    RETRY: {
        MAX_ATTEMPTS: 3,//重试次数
        DELAY_BASE: 1000 // 基础延迟时间（毫秒）
    },
    CHROME_PATH: process.env.CHROME_PATH || 'C:/Program Files/Google/Chrome/Application/chrome.exe'// 替换为你的 Chrome 实际路径
};


// 请求头配置
const DEFAULT_HEADERS = {
    'accept': '*/*',
    'accept-language': 'zh-CN,zh;q=0.9',
    'accept-encoding': 'gzip, deflate, br, zstd',
    'content-type': 'text/plain;charset=UTF-8',
    'Connection': 'keep-alive',
    'origin': 'https://grok.com',
    'priority': 'u=1, i',
    'sec-ch-ua': '"Chromium";v="130", "Google Chrome";v="130", "Not?A_Brand";v="99"',
    'sec-ch-ua-mobile': '?0',
    'sec-ch-ua-platform': '"Windows"',
    'sec-fetch-dest': 'empty',
    'sec-fetch-mode': 'cors',
    'sec-fetch-site': 'same-origin',
    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36',
    'baggage': 'sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c'
};

class Utils {
    static async extractGrokHeaders() {
        try {
            // 启动浏览器
            const browser = await launch({
                executablePath: CONFIG.CHROME_PATH,
                headless: true
            });

            const page = await browser.newPage();
            await page.goto(`${CONFIG.API.BASE_URL}`, { waitUntil: 'networkidle0' });

            // 获取所有 Cookies
            const cookies = await page.cookies();
            const targetHeaders = ['x-anonuserid', 'x-challenge', 'x-signature'];
            const extractedHeaders = {};
            // 遍历 Cookies
            for (const cookie of cookies) {
                // 检查是否为目标头信息
                if (targetHeaders.includes(cookie.name.toLowerCase())) {
                    extractedHeaders[cookie.name.toLowerCase()] = cookie.value;
                }
            }
            // 关闭浏览器
            await browser.close();
            // 打印并返回提取的头信息
            console.log('提取的头信息:', JSON.stringify(extractedHeaders, null, 2));
            return extractedHeaders;

        } catch (error) {
            console.error('获取头信息出错:', error);
            return null;
        }
    }
    static async get_signature() {
        console.log("刷新认证信息");
        let retryCount = 0;
        while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
            let headers = await Utils.extractGrokHeaders();
            if (headers) {
                console.log("获取认证信息成功");
                CONFIG.API.SIGNATURE_COOKIE = { cookie: `x-anonuserid=${headers["x-anonuserid"]}; x-challenge=${headers["x-challenge"]}; x-signature=${headers["x-signature"]}` };
                fs.writeFileSync(path.resolve(process.cwd(), 'signature.json'), JSON.stringify(CONFIG.API.SIGNATURE_COOKIE));//保存认证信息
                return CONFIG.API.SIGNATURE_COOKIE;
            }
            retryCount++;
            if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
                throw new Error(`获取认证信息失败!`);
            }
            await new Promise(resolve => setTimeout(resolve, CONFIG.RETRY.DELAY_BASE * retryCount));
        }
    }
    static async retryRequestWithNewSignature(originalRequest) {
        console.log("认证信息已过期，尝试刷新认证信息并重新请求～");
        try {
            // 刷新认证信息
            await Utils.get_signature();

            // 重新执行原始请求
            return await originalRequest();
        } catch (error) {
            // 第二次失败直接抛出错误
            throw new Error(`重试请求失败: ${error.message}`);
        }
    }
    static async handleError(error, res, originalRequest = null) {
        console.error('Error:', error);

        // 如果是500错误且提供了原始请求函数，尝试重新获取签名并重试
        if (error.status === 500 && originalRequest) {
            try {
                const result = await Utils.retryRequestWithNewSignature(originalRequest);
                return result;
            } catch (retryError) {
                console.error('重试失败:', retryError);
                return res.status(500).json({
                    error: {
                        message: retryError.message,
                        type: 'server_error',
                        param: null,
                        code: retryError.code || null
                    }
                });
            }
        }

        // 其他错误直接返回
        res.status(500).json({
            error: {
                message: error.message,
                type: 'server_error',
                param: null,
                code: error.code || null
            }
        });
    }
}


class GrokApiClient {
    constructor(modelId) {
        if (!CONFIG.MODELS[modelId]) {
            throw new Error(`不支持的模型: ${modelId}`);
        }
        this.modelId = CONFIG.MODELS[modelId];
    }

    processMessageContent(content) {
        if (typeof content === 'string') return content;
        return null;
    }
    // 获取图片类型
    getImageType(base64String) {
        let mimeType = 'image/jpeg';
        if (base64String.includes('data:image')) {
            const matches = base64String.match(/data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,/);
            if (matches) {
                mimeType = matches[1];
            }
        }
        const extension = mimeType.split('/')[1];
        const fileName = `image.${extension}`;

        return {
            mimeType: mimeType,
            fileName: fileName
        };
    }

    async uploadBase64Image(base64Data, url) {
        try {
            // 处理 base64 数据
            let imageBuffer;
            if (base64Data.includes('data:image')) {
                imageBuffer = base64Data.split(',')[1];
            } else {
                imageBuffer = base64Data
            }
            const { mimeType, fileName } = this.getImageType(base64Data);
            let uploadData = {
                rpc: "uploadFile",
                req: {
                    fileName: fileName,
                    fileMimeType: mimeType,
                    content: imageBuffer
                }
            };
            console.log("发送图片请求");
            // 发送请求
            const response = await fetch(url, {
                method: 'POST',
                headers: {
                    ...CONFIG.DEFAULT_HEADERS,
                    ...CONFIG.API.SIGNATURE_COOKIE
                },
                body: JSON.stringify(uploadData)
            });

            if (!response.ok) {
                console.error(`上传图片失败,状态码:${response.status},原因:${response.error}`);
                return '';
            }

            const result = await response.json();
            console.log('上传图片成功:', result);
            return result.fileMetadataId;

        } catch (error) {
            console.error('上传图片失败:', error);
            return '';
        }
    }

    async prepareChatRequest(request) {
        const processImageUrl = async (content) => {
            if (content.type === 'image_url' && content.image_url.url.includes('data:image')) {
                const imageResponse = await this.uploadBase64Image(
                    content.image_url.url,
                    `${CONFIG.API.BASE_URL}/api/rpc`
                );
                return imageResponse;
            }
            return null;
        };
        let fileAttachments = [];
        let messages = '';
        let lastRole = null;
        let lastContent = '';

        for (const current of request.messages) {
            const role = current.role === 'assistant' ? 'assistant' : 'user';
            let textContent = '';
            // 处理消息内容
            if (Array.isArray(current.content)) {
                // 处理数组内的所有内容
                for (const item of current.content) {
                    if (item.type === 'image_url') {
                        // 如果是图片且是最后一条消息，则处理图片
                        if (current === request.messages[request.messages.length - 1]) {
                            const processedImage = await processImageUrl(item);
                            if (processedImage) fileAttachments.push(processedImage);
                        }
                        textContent += (textContent ? '\n' : '') + "[图片]";//图片占位符
                    } else if (item.type === 'text') {
                        textContent += (textContent ? '\n' : '') + item.text;
                    }
                }
            } else if (typeof current.content === 'object' && current.content !== null) {
                // 处理单个对象内容
                if (current.content.type === 'image_url') {
                    // 如果是图片且是最后一条消息，则处理图片
                    if (current === request.messages[request.messages.length - 1]) {
                        const processedImage = await processImageUrl(current.content);
                        if (processedImage) fileAttachments.push(processedImage);
                    }
                    textContent += (textContent ? '\n' : '') + "[图片]";//图片占位符
                } else if (current.content.type === 'text') {
                    textContent = current.content.text;
                }
            } else {
                // 处理普通文本内容
                textContent = this.processMessageContent(current.content);
            }
            // 添加文本内容到消息字符串
            if (textContent) {
                if (role === lastRole) {
                    // 如果角色相同，合并消息内容
                    lastContent += '\n' + textContent;
                    messages = messages.substring(0, messages.lastIndexOf(`${role.toUpperCase()}: `)) +
                        `${role.toUpperCase()}: ${lastContent}\n`;
                } else {
                    // 如果角色不同，添加新的消息
                    messages += `${role.toUpperCase()}: ${textContent}\n`;
                    lastContent = textContent;
                    lastRole = role;
                }
            } else if (current === request.messages[request.messages.length - 1] && fileAttachments.length > 0) {
                // 如果是最后一条消息且有图片附件，添加空消息占位
                messages += `${role.toUpperCase()}: [图片]\n`;
            }
        }

        if (fileAttachments.length > 4) {
            fileAttachments = fileAttachments.slice(0, 4); // 最多上传4张
        }

        messages = messages.trim();

        return {
            message: messages,
            modelName: this.modelId,
            disableSearch: false,
            imageAttachments: [],
            returnImageBytes: false,
            returnRawGrokInXaiRequest: false,
            fileAttachments: fileAttachments,
            enableImageStreaming: false,
            imageGenerationCount: 1,
            toolOverrides: {
                imageGen: request.model === 'grok-latest-image',
                webSearch: false,
                xSearch: false,
                xMediaSearch: false,
                trendsSearch: false,
                xPostAnalyze: false
            }
        };
    }
}

class MessageProcessor {
    static createChatResponse(message, model, isStream = false) {
        const baseResponse = {
            id: `chatcmpl-${uuidv4()}`,
            created: Math.floor(Date.now() / 1000),
            model: model
        };

        if (isStream) {
            return {
                ...baseResponse,
                object: 'chat.completion.chunk',
                choices: [{
                    index: 0,
                    delta: {
                        content: message
                    }
                }]
            };
        }

        return {
            ...baseResponse,
            object: 'chat.completion',
            choices: [{
                index: 0,
                message: {
                    role: 'assistant',
                    content: message
                },
                finish_reason: 'stop'
            }],
            usage: null
        };
    }
}

// 中间件配置
const app = express();
app.use(express.json({ limit: '5mb' }));
app.use(express.urlencoded({ extended: true, limit: '5mb' }));
app.use(cors({
    origin: '*',
    methods: ['GET', 'POST', 'OPTIONS'],
    allowedHeaders: ['Content-Type', 'Authorization']
}));
// API路由
app.get('/hf/v1/models', (req, res) => {
    res.json({
        object: "list",
        data: Object.keys(CONFIG.MODELS).map((model, index) => ({
            id: model,
            object: "model",
            created: Math.floor(Date.now() / 1000),
            owned_by: "xai",
        }))
    });
});

app.post('/hf/v1/chat/completions', async (req, res) => {
    const authToken = req.headers.authorization?.replace('Bearer ', '');
    if (authToken !== CONFIG.API.API_KEY) {
        return res.status(401).json({ error: 'Unauthorized' });
    }
    if (req.body.model === 'grok-latest-image' && req.body.stream) {
        return res.status(400).json({ error: '该模型不支持流式' });
    }

    const makeRequest = async () => {
        if (!CONFIG.API.SIGNATURE_COOKIE) {
            await Utils.get_signature();
            CONFIG.API.SIGNATURE_COOKIE = JSON.parse(fs.readFileSync(path.resolve(process.cwd(), 'signature.json')));
        }
        const grokClient = new GrokApiClient(req.body.model);
        const requestPayload = await grokClient.prepareChatRequest(req.body);
        //创建新对话        
        const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
            method: 'POST',
            headers: {
                ...DEFAULT_HEADERS,
                ...CONFIG.API.SIGNATURE_COOKIE
            },
            body: JSON.stringify({
                rpc: "createConversation",
                req: {
                    temporary: false
                }
            })
        });
        if (!newMessageReq.ok) {
            throw new Error(`上游服务请求失败! status: ${newMessageReq.status}`);
        }

        // 获取响应文本
        const responseText = await newMessageReq.json();
        const conversationId = responseText.conversationId;
        console.log("会话ID:conversationId", conversationId);
        if (!conversationId) {
            throw new Error(`创建会话失败! status: ${newMessageReq.status}`);
        }
        //发送对话
        const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
            method: 'POST',
            headers: {
                "accept": "text/event-stream",
                "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
                "content-type": "text/plain;charset=UTF-8",
                "Connection": "keep-alive",
                ...CONFIG.API.SIGNATURE_COOKIE
            },
            body: JSON.stringify(requestPayload)
        });

        if (!response.ok) {
            throw new Error(`上游服务请求失败! status: ${response.status}`);
        }

        if (req.body.stream) {
            await handleStreamResponse(response, req.body.model, res);
        } else {
            await handleNormalResponse(response, req.body.model, res);
        }
    }
    try {
        await makeRequest();
    } catch (error) {
        await Utils.handleError(error, res, makeRequest);
    }
});

async function handleStreamResponse(response, model, res) {
    res.setHeader('Content-Type', 'text/event-stream');
    res.setHeader('Cache-Control', 'no-cache');
    res.setHeader('Connection', 'keep-alive');

    try {
        // 获取响应文本
        const responseText = await response.text();
        const lines = responseText.split('\n');

        for (const line of lines) {
            if (!line.startsWith('data: ')) continue;

            const data = line.slice(6);
            if (data === '[DONE]') {
                res.write('data: [DONE]\n\n');
                continue;
            }
            console.log("data", data);
            let linejosn = JSON.parse(data);
            const token = linejosn.token;
            if (token && token.length > 0) {
                const responseData = MessageProcessor.createChatResponse(token, model, true);
                res.write(`data: ${JSON.stringify(responseData)}\n\n`);
            }
        }

        res.end();
    } catch (error) {
        Utils.handleError(error, res);
    }
}

async function handleNormalResponse(response, model, res) {
    let fullResponse = '';
    let imageUrl = '';

    try {
        const responseText = await response.text();
        const lines = responseText.split('\n');

        for (const line of lines) {
            if (!line.startsWith('data: ')) continue;

            const data = line.slice(6);
            if (data === '[DONE]') continue;
            let linejosn = JSON.parse(data);
            if (linejosn.response === "token") {
                const token = linejosn.token;
                if (token && token.length > 0) {
                    fullResponse += token;
                }
            } else if (linejosn.response === "modelResponse" && model === 'grok-latest-image') {
                if (linejosn?.modelResponse?.generatedImageUrls) {
                    imageUrl = linejosn.modelResponse.generatedImageUrls;
                }
            }
        }

        if (imageUrl) {
            const dataImage = await handleImageResponse(imageUrl);
            fullResponse = [{
                type: "image_url",
                image_url: {
                    url: dataImage
                }
            }]
            const responseData = MessageProcessor.createChatResponse(fullResponse, model);
            res.json(responseData);
        } else {
            const responseData = MessageProcessor.createChatResponse(fullResponse, model);
            res.json(responseData);
        }
    } catch (error) {
        Utils.handleError(error, res);
    }
}
async function handleImageResponse(imageUrl) {
    //对服务器发送图片请求
    const MAX_RETRIES = 3;
    let retryCount = 0;
    let imageBase64Response;

    while (retryCount < MAX_RETRIES) {
        try {
            //发送图片请求获取图片
            imageBase64Response = await fetch(`https://assets.grok.com/${imageUrl}`, {
                method: 'GET',
                headers: {
                    ...DEFAULT_HEADERS,
                    ...CONFIG.API.SIGNATURE_COOKIE
                }
            });

            if (imageBase64Response.ok) {
                break; // 如果请求成功，跳出重试循环
            }

            retryCount++;
            if (retryCount === MAX_RETRIES) {
                throw new Error(`上游服务请求失败! status: ${imageBase64Response.status}`);
            }

            // 等待一段时间后重试（可以使用指数退避）
            await new Promise(resolve => setTimeout(resolve, 500 * retryCount));

        } catch (error) {
            retryCount++;
            if (retryCount === MAX_RETRIES) {
                throw error;
            }
            // 等待一段时间后重试
            await new Promise(resolve => setTimeout(resolve, 500 * retryCount));
        }
    }

    const arrayBuffer = await imageBase64Response.arrayBuffer();
    const imagebuffer = Buffer.from(arrayBuffer);

    const base64String = imagebuffer.toString('base64');

    // 获取图片类型（例如：image/jpeg, image/png）
    const imageContentType = imageBase64Response.headers.get('content-type');

    // 构建 Data URL返回
    return `data:${imageContentType};base64,${base64String}`;
}

// 404处理
app.use((req, res) => {
    res.status(404).send('请求路径不存在');
});

// 启动服务器
app.listen(CONFIG.SERVER.PORT, () => {
    console.log(`服务器已启动，监听端口: ${CONFIG.SERVER.PORT}`);
});