import express from 'express';
import FormData from 'form-data';
import { v4 as uuidv4 } from 'uuid';
import fetch from 'node-fetch';
import cors from 'cors';
import Logger from './logger.js';
import dotenv from 'dotenv';

// 初始化环境变量
dotenv.config();

// 配置管理
const CONFIG = {
    SERVER: {
        PORT: process.env.PORT || 25526,
        BODY_LIMIT: '5mb',
        CORS_OPTIONS: {
            origin: '*',
            methods: ['GET', 'POST', 'OPTIONS'],
            allowedHeaders: ['Content-Type', 'Authorization'],
            credentials: true
        }
    },
    API: {
        API_KEY: process.env.API_KEY || "sk-123456",
        AUTH_TOKEN: process.env.AUTH_TOKEN,
        CT0: process.env.CT0,
        ENDPOINTS: {
            CHAT: 'https://grok.x.com/2/grok/add_response.json',
            CREATE_CONVERSATION: 'https://x.com/i/api/graphql/vvC5uy7pWWHXS2aDi1FZeA/CreateGrokConversation',
            DELETE_CONVERSATION: 'https://x.com/i/api/graphql/TlKHSWVMVeaa-i7dqQqFQA/ConversationItem_DeleteConversationMutation',
            UPLOAD_IMAGE: 'https://x.com/i/api/2/grok/attachment.json'
        }
    },
    MODELS: {
        "grok-3": "grok-3",
        "grok-3-deepsearch": "grok-3",
        "grok-3-reasoning": "grok-3",
        "grok-3-imageGen": "grok-3",
    },
    IS_IMG_GEN: false,
    IS_THINKING: false
};

// HTTP 请求头配置
const DEFAULT_HEADERS = {
    'authorization': 'Bearer AAAAAAAAAAAAAAAAAAAAANRILgAAAAAAnNwIzUejRCOuH5E6I8xnZz4puTs%3D1Zv7ttfk8LF81IUq16cHjhLTvJu4FA33AGWWjCpTnA',
    'sec-ch-ua': '"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
    'sec-ch-ua-mobile': '?0',
    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36',
    'accept': '*/*',
    'content-type': 'text/plain;charset=UTF-8',
    'origin': 'https://x.com',
    'sec-fetch-site': 'same-site',
    'sec-fetch-mode': 'cors',
    'accept-encoding': 'gzip, deflate, br, zstd',
    'accept-language': 'zh-CN,zh;q=0.9',
    'priority': 'u=1, i'
};

// 工具类
class Utils {
    static generateRandomString(length, charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ') {
        return Array(length).fill(null)
            .map(() => charset[Math.floor(Math.random() * charset.length)])
            .join('');
    }

    static createAuthHeaders() {
        return {
            ...DEFAULT_HEADERS,
            'x-csrf-token': CONFIG.API.CT0,
            'cookie': `auth_token=${CONFIG.API.AUTH_TOKEN};ct0=${CONFIG.API.CT0}`
        };
    }

    static getImageMimeType(base64String) {
        const matches = base64String.match(/data:([a-zA-Z0-9]+\/[a-zA-Z0-9-.+]+);base64,/);
        return matches ? matches[1] : 'image/jpeg';
    }

    static async handleApiResponse(response, errorMessage) {
        if (!response.ok) {
            throw new Error(`${errorMessage} Status: ${response.status}`);
        }
        return await response.json();
    }
}

// 会话管理类
class ConversationManager {
    static async generateNewId() {
        const response = await fetch(CONFIG.API.ENDPOINTS.CREATE_CONVERSATION, {
            method: 'POST',
            headers: Utils.createAuthHeaders(),
            body: JSON.stringify({
                variables: {},
                queryId: "vvC5uy7pWWHXS2aDi1FZeA"
            })
        });

        const data = await Utils.handleApiResponse(response, '创建会话失败!');
        return data.data.create_grok_conversation.conversation_id;
    }

    static async deleteConversation(conversationId) {
        if (!conversationId) return;

        await fetch(CONFIG.API.ENDPOINTS.DELETE_CONVERSATION, {
            method: 'POST',
            headers: Utils.createAuthHeaders(),
            body: JSON.stringify({
                variables: { conversationId },
                queryId: "TlKHSWVMVeaa-i7dqQqFQA"
            })
        });
    }
}

// 消息处理类
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
        };
    }

    static processMessageContent(content) {
        if (typeof content === 'string') return content;
        if (Array.isArray(content)) {
            if (content.some(item => item.type === 'image_url')) return null;
            return content
                .filter(item => item.type === 'text')
                .map(item => item.text)
                .join('\n');
        }
        if (typeof content === 'object') return content.text || null;
        return null;
    }
}

// Grok API 客户端类
class TwitterGrokApiClient {
    constructor(modelId) {
        if (!CONFIG.MODELS[modelId]) {
            throw new Error(`不支持的模型: ${modelId}`);
        }
        this.modelId = CONFIG.MODELS[modelId];
        this.modelType = {
            isDeepSearch: modelId === 'grok-3-deepsearch',
            isReasoning: modelId === 'grok-3-reasoning'
        };
    }

    async uploadImage(imageData) {
        const formData = new FormData();
        const imageBuffer = Buffer.from(imageData.split(',')[1], 'base64');
        const mimeType = Utils.getImageMimeType(imageData);

        formData.append('photo', imageBuffer, {
            filename: 'image.png',
            contentType: mimeType
        });

        const response = await fetch(CONFIG.API.ENDPOINTS.UPLOAD_IMAGE, {
            method: 'POST',
            headers: {
                ...Utils.createAuthHeaders(),
                ...formData.getHeaders()
            },
            body: formData
        });

        return await Utils.handleApiResponse(response, '图片上传失败');
    }

    async transformMessages(messages) {
        const processedMessages = [];
        for (let i = 0; i < messages.length; i++) {
            const isLastTwoMessages = i >= messages.length - 2;
            const content = await this.processMessageContent(messages[i], isLastTwoMessages);
            if (content) {
                processedMessages.push(content);
            }
        }
        return processedMessages;
    }

    async processMessageContent(msg, isLastTwoMessages) {
        const { role, content } = msg;
        let message = '';
        let fileAttachments = [];

        if (typeof content === 'string') {
            message = content;
        } else if (Array.isArray(content) || typeof content === 'object') {
            const { text, imageAttachments } = await this.processComplexContent(content, isLastTwoMessages);
            message = text;
            fileAttachments = imageAttachments;
        }

        return {
            message,
            sender: role === 'user' ? 1 : 2,
            ...(role === 'user' && { fileAttachments })
        };
    }

    async processComplexContent(content, isLastTwoMessages) {
        let text = '';
        let imageAttachments = [];

        const processItem = async (item) => {
            if (item.type === 'text') {
                text += item.text;
            } else if (item.type === 'image_url' && item.image_url.url.includes('data:image')) {
                if (isLastTwoMessages) {
                    const uploadResult = await this.uploadImage(item.image_url.url);
                    if (Array.isArray(uploadResult)) {
                        imageAttachments.push(...uploadResult);
                    }
                } else {
                    text += '[图片]';
                }
            }
        };

        if (Array.isArray(content)) {
            await Promise.all(content.map(processItem));
        } else {
            await processItem(content);
        }

        return { text, imageAttachments };
    }

    async prepareChatRequest(request) {
        const responses = await this.transformMessages(request.messages);
        const conversationId = await ConversationManager.generateNewId();

        return {
            responses,
            systemPromptName: "",
            grokModelOptionId: this.modelId,
            conversationId,
            returnSearchResults: this.modelType.isReasoning,
            returnCitations: this.modelType.isReasoning,
            promptMetadata: {
                promptSource: "NATURAL",
                action: "INPUT"
            },
            imageGenerationCount: 1,
            requestFeatures: {
                eagerTweets: false,
                serverHistory: false
            },
            enableCustomization: true,
            enableSideBySide: false,
            toolOverrides: {
                imageGen: request.model === 'grok-3-imageGen',
            },
            isDeepsearch: this.modelType.isDeepSearch,
            isReasoning: this.modelType.isReasoning
        };
    }
}

// 响应处理类
class ResponseHandler {
    static async handleStreamResponse(response, model, res) {
        res.setHeader('Content-Type', 'text/event-stream');
        res.setHeader('Cache-Control', 'no-cache');
        res.setHeader('Connection', 'keep-alive');

        const reader = response.body;
        let buffer = '';
        CONFIG.IS_IMG_GEN = false;
        CONFIG.IS_THINKING = false;

        try {
            for await (const chunk of reader) {
                const lines = (buffer + chunk.toString()).split('\n');
                buffer = lines.pop() || '';

                for (const line of lines) {
                    if (!line.trim()) continue;
                    await this.processStreamLine(JSON.parse(line), model, res);
                }
            }

            res.write('data: [DONE]\n\n');
            res.end();
        } catch (error) {
            Logger.error('Stream response error:', error, 'ChatAPI');
            throw error;
        }
    }

    static async processStreamLine(jsonData, model, res) {
        if (jsonData.result?.doImgGen) {
            CONFIG.IS_IMG_GEN = true;
            return;
        }

        if (CONFIG.IS_IMG_GEN && jsonData.result?.event?.imageAttachmentUpdate?.progress === 100) {
            await this.handleImageGeneration(jsonData, model, res);
            return;
        }

        if (!CONFIG.IS_IMG_GEN && jsonData.result?.message) {
            await this.handleTextMessage(jsonData, model, res);
        }
    }

    static async handleImageGeneration(jsonData, model, res) {
        const imageUrl = jsonData.result.event.imageAttachmentUpdate.imageUrl;
        const imageResponse = await fetch(imageUrl, {
            method: 'GET',
            headers: Utils.createAuthHeaders()
        });

        if (!imageResponse.ok) {
            throw new Error(`Image request failed: ${imageResponse.status}`);
        }

        const imageBuffer = await imageResponse.arrayBuffer();
        const base64Image = Buffer.from(imageBuffer).toString('base64');
        const imageContentType = imageResponse.headers.get('content-type');
        const message = `![image](data:${imageContentType};base64,${base64Image})`;

        const responseData = MessageProcessor.createChatResponse(message, model, true);
        res.write(`data: ${JSON.stringify(responseData)}\n\n`);
    }

    static async handleTextMessage(jsonData, model, res) {
        let message = jsonData.result.message;

        switch (model) {
            case "grok-3-reasoning":
                if (!CONFIG.IS_THINKING && jsonData.result?.isThinking) {
                    message = "<think>" + message;
                    CONFIG.IS_THINKING = true;
                } else if (CONFIG.IS_THINKING && !jsonData.result?.isThinking) {
                    message = "</think>" + message;
                    CONFIG.IS_THINKING = false;
                }
                break;
            case "grok-3-deepsearch":
                if (jsonData.result?.messageTag !== "final") return;
                break;
        }

        const responseData = MessageProcessor.createChatResponse(message, model, true);
        res.write(`data: ${JSON.stringify(responseData)}\n\n`);
    }

    static async handleNormalResponse(response, model, res) {
        const reader = response.body;
        let buffer = '';
        let fullResponse = '';
        let imageUrl = null;

        try {
            for await (const chunk of reader) {
                const lines = (buffer + chunk.toString()).split('\n');
                buffer = lines.pop() || '';

                for (const line of lines) {
                    if (!line.trim()) continue;
                    const result = await this.processNormalLine(JSON.parse(line), model, CONFIG.IS_THINKING);
                    fullResponse += result.text || '';
                    imageUrl = result.imageUrl || imageUrl;
                    CONFIG.IS_THINKING = result.isThinking;
                }
            }

            if (imageUrl) {
                await this.sendImageResponse(imageUrl, model, res);
            } else {
                const responseData = MessageProcessor.createChatResponse(fullResponse, model);
                res.json(responseData);
            }
        } catch (error) {
            Logger.error('Normal response error:', error, 'ChatAPI');
            throw error;
        }
    }

    static async processNormalLine(jsonData, model, isThinking) {
        let result = { text: '', imageUrl: null, isThinking };

        if (jsonData.result?.message) {
            switch (model) {
                case "grok-3-reasoning":
                    result = this.processReasoningMessage(jsonData, isThinking);
                    break;
                case "grok-3-deepsearch":
                    if (jsonData.result?.messageTag === "final") {
                        result.text = jsonData.result.message;
                    }
                    break;
                default:
                    result.text = jsonData.result.message;
            }
        }

        if (jsonData.result?.event?.imageAttachmentUpdate?.progress === 100) {
            result.imageUrl = jsonData.result.event.imageAttachmentUpdate.imageUrl;
        }

        return result;
    }

    static processReasoningMessage(jsonData, isThinking) {
        let result = { text: '', isThinking };

        if (jsonData.result?.isThinking && !isThinking) {
            result.text = "<think>" + jsonData.result.message;
            result.isThinking = true;
        } else if (isThinking && !jsonData.result?.isThinking) {
            result.text = "</think>" + jsonData.result.message;
            result.isThinking = false;
        } else {
            result.text = jsonData.result.message;
        }

        return result;
    }

    static async sendImageResponse(imageUrl, model, res) {
        const response = await fetch(imageUrl, {
            method: 'GET',
            headers: Utils.createAuthHeaders()
        });

        if (!response.ok) {
            throw new Error(`Image request failed: ${response.status}`);
        }

        const imageBuffer = await response.arrayBuffer();
        const base64Image = Buffer.from(imageBuffer).toString('base64');
        const imageContentType = response.headers.get('content-type');

        const responseData = MessageProcessor.createChatResponse(
            `![image](data:${imageContentType};base64,${base64Image})`,
            model
        );
        res.json(responseData);
    }
}

// Express 应用配置
const app = express();
app.use(Logger.requestLogger);
app.use(cors(CONFIG.SERVER.CORS_OPTIONS));
app.use(express.json({ limit: CONFIG.SERVER.BODY_LIMIT }));
app.use(express.urlencoded({ extended: true, limit: CONFIG.SERVER.BODY_LIMIT }));

// API 路由
app.get('/v1/models', (req, res) => {
    res.json({
        object: "list",
        data: Object.keys(CONFIG.MODELS).map(model => ({
            id: model,
            object: "model",
            created: Math.floor(Date.now() / 1000),
            owned_by: "xgrok",
        }))
    });
});

app.post('/v1/chat/completions', async (req, res) => {
    try {
        const authToken = req.headers.authorization?.replace('Bearer ', '');
        if (authToken !== CONFIG.API.API_KEY) {
            return res.status(401).json({ error: 'Unauthorized' });
        }

        const grokClient = new TwitterGrokApiClient(req.body.model);
        const requestPayload = await grokClient.prepareChatRequest(req.body);

        const response = await fetch(CONFIG.API.ENDPOINTS.CHAT, {
            method: 'POST',
            headers: Utils.createAuthHeaders(),
            body: JSON.stringify(requestPayload)
        });

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

        await (req.body.stream
            ? ResponseHandler.handleStreamResponse(response, req.body.model, res)
            : ResponseHandler.handleNormalResponse(response, req.body.model, res));

    } catch (error) {
        Logger.error('Chat Completions Request Error', error, 'ChatAPI');
        res.status(500).json({
            error: {
                message: error.message,
                type: 'server_error',
                param: null,
                code: error.code || null
            }
        });
    } finally {
        if (req.body.conversationId) {
            await ConversationManager.deleteConversation(req.body.conversationId);
        }
    }
});

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

// 启动服务器
app.listen(CONFIG.SERVER.PORT, () => {
    Logger.info(`服务器运行在端口 ${CONFIG.SERVER.PORT}`, 'Server');
});