import * as vscode from 'vscode';
import fetch from 'node-fetch';
import { ApiKeyManager } from './apiKeyManager';

interface ZAIToolCall {
    id: string;
    type: 'function';
    function: {
        name: string;
        arguments: string;
    };
}

interface ZAIMessage {
    role: 'user' | 'assistant' | 'system' | 'tool';
    content?: string;
    tool_calls?: ZAIToolCall[];
    tool_call_id?: string;
    name?: string;
}

interface ZAITool {
    type: 'function';
    function: {
        name: string;
        description?: string;
        parameters?: any;
    };
}

export class ZAIAPIClient {
    private apiEndpoint: string;
    private timeout: number;
    private enableThinkingMode: boolean;

    constructor(private apiKeyManager: ApiKeyManager) {
        const config = vscode.workspace.getConfiguration('zai');
        this.apiEndpoint = config.get<string>('apiEndpoint', 'https://api.z.ai/api/coding/paas/v4');
        this.timeout = config.get<number>('timeout', 60000);
        this.enableThinkingMode = config.get<boolean>('enableThinkingMode', true);
    }

    /**
     * Convert VS Code messages to Z.AI format
     */
    private convertMessages(messages: readonly vscode.LanguageModelChatRequestMessage[]): ZAIMessage[] {
        const converted: ZAIMessage[] = [];

        for (const msg of messages) {
            const role = this.mapRole(msg.role);
            const textParts: string[] = [];
            const toolCalls: ZAIToolCall[] = [];
            const toolResults: Array<{ id: string; content: any }> = [];

            for (const part of msg.content) {
                if (part instanceof vscode.LanguageModelTextPart) {
                    textParts.push(part.value);
                } else if (part instanceof vscode.LanguageModelToolCallPart) {
                    toolCalls.push({
                        id: part.callId ?? `${part.name}-${toolCalls.length}`,
                        type: 'function',
                        function: {
                            name: part.name,
                            arguments: JSON.stringify(part.input ?? {})
                        }
                    });
                } else if (part instanceof vscode.LanguageModelToolResultPart) {
                    toolResults.push({
                        id: part.callId,
                        content: part.content
                    });
                }
            }

            for (const result of toolResults) {
                converted.push({
                    role: 'tool',
                    tool_call_id: result.id,
                    content: typeof result.content === 'string' ? result.content : JSON.stringify(result.content)
                });
            }

            const content = textParts.join('');
            if (content || toolCalls.length > 0 || role === 'system') {
                const message: ZAIMessage = { role };
                if (content) {
                    message.content = content;
                }
                if (toolCalls.length > 0) {
                    message.tool_calls = toolCalls;
                }
                converted.push(message);
            }
        }

        return converted;
    }

    private mapRole(role: vscode.LanguageModelChatMessageRole): ZAIMessage['role'] {
        const toolRole = (vscode.LanguageModelChatMessageRole as any).Tool;
        if (toolRole && role === toolRole) {
            return 'tool';
        }

        switch (role) {
            case vscode.LanguageModelChatMessageRole.Assistant:
                return 'assistant';
            default:
                return 'user';
        }
    }

    /**
     * Convert VS Code tools to Z.AI format
     */
    private convertTools(tools: readonly vscode.LanguageModelChatTool[] | undefined): ZAITool[] | undefined {
        if (!tools || tools.length === 0) {
            return undefined;
        }

        return tools.map(tool => ({
            type: 'function' as const,
            function: {
                name: tool.name,
                description: tool.description,
                parameters: tool.inputSchema
            }
        }));
    }

    /**
     * Stream chat completion from Z.AI API
     */
    async streamChatCompletion(
        modelId: string,
        messages: readonly vscode.LanguageModelChatRequestMessage[],
        options: vscode.ProvideLanguageModelChatResponseOptions,
        progress: vscode.Progress<vscode.LanguageModelResponsePart>,
        token: vscode.CancellationToken
    ): Promise<void> {

        const apiKey = await this.apiKeyManager.getApiKey();
        if (!apiKey) {
            throw new Error('API key not configured');
        }

        const convertedMessages = this.convertMessages(messages);
        const convertedTools = this.convertTools(options.tools);

        const requestBody: any = {
            model: modelId,
            messages: convertedMessages,
            stream: true,
            max_tokens: 4096,
        };

        // Add tools if present
        if (convertedTools && convertedTools.length > 0) {
            requestBody.tools = convertedTools;
            requestBody.tool_choice = 'auto';
        }

        // Add thinking mode if enabled
        if (this.enableThinkingMode) {
            requestBody.enable_thinking = true;
        }

        try {
            const controller = new AbortController();
            const timeoutId = setTimeout(() => controller.abort(), this.timeout);

            // Set up cancellation
            const cancellationListener = token.onCancellationRequested(() => {
                controller.abort();
            });

            const response = await fetch(`${this.apiEndpoint}/chat/completions`, {
                method: 'POST',
                headers: {
                    'Authorization': `Bearer ${apiKey}`,
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(requestBody),
                signal: controller.signal as any
            });

            clearTimeout(timeoutId);
            cancellationListener.dispose();

            if (!response.ok) {
                const errorText = await response.text();
                throw new Error(`Z.AI API error (${response.status}): ${errorText}`);
            }

            if (!response.body) {
                throw new Error('Response body is null');
            }

            // Process the streaming response
            await this.processStream(response.body, progress, token);

        } catch (error) {
            if (error instanceof Error) {
                if (error.name === 'AbortError') {
                    throw new vscode.CancellationError();
                }
                throw error;
            }
            throw new Error(`Unexpected error: ${String(error)}`);
        }
    }

    /**
     * Process the streaming response from Z.AI
     */
    private async processStream(
        stream: NodeJS.ReadableStream,
        progress: vscode.Progress<vscode.LanguageModelResponsePart>,
        token: vscode.CancellationToken
    ): Promise<void> {

        let buffer = '';
        let currentToolCall: any = null;

        for await (const chunk of stream) {
            if (token.isCancellationRequested) {
                throw new vscode.CancellationError();
            }

            buffer += chunk.toString();
            const lines = buffer.split('\n');
            buffer = lines.pop() || '';

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

                const data = line.slice(6).trim();
                if (data === '[DONE]') {
                    // Finish any pending tool call
                    if (currentToolCall) {
                        progress.report(new vscode.LanguageModelToolCallPart(
                            currentToolCall.id,
                            currentToolCall.function.name,
                            this.parseToolArguments(currentToolCall.function.arguments)
                        ));
                        currentToolCall = null;
                    }
                    continue;
                }

                try {
                    const parsed = JSON.parse(data);
                    const choice = parsed.choices?.[0];
                    if (!choice) continue;

                    const delta = choice.delta;

                    // Handle text content
                    if (delta.content) {
                        if (Array.isArray(delta.content)) {
                            for (const contentPart of delta.content) {
                                if (typeof contentPart?.text === 'string') {
                                    progress.report(new vscode.LanguageModelTextPart(contentPart.text));
                                }
                            }
                        } else if (typeof delta.content === 'string') {
                            progress.report(new vscode.LanguageModelTextPart(delta.content));
                        }
                    }

                    // Handle tool calls
                    if (delta.tool_calls && delta.tool_calls.length > 0) {
                        for (const toolCall of delta.tool_calls) {
                            if (toolCall.function?.name) {
                                // Start of new tool call
                                if (currentToolCall) {
                                    // Finish previous tool call
                                    progress.report(new vscode.LanguageModelToolCallPart(
                                        currentToolCall.id,
                                        currentToolCall.function.name,
                                        this.parseToolArguments(currentToolCall.function.arguments)
                                    ));
                                }
                                currentToolCall = {
                                    id: toolCall.id || `call_${Date.now()}`,
                                    function: {
                                        name: toolCall.function.name,
                                        arguments: toolCall.function.arguments || ''
                                    }
                                };
                            } else if (toolCall.function?.arguments && currentToolCall) {
                                // Append to current tool call
                                currentToolCall.function.arguments += toolCall.function.arguments;
                            }
                        }
                    }

                } catch (error) {
                    console.error('Error parsing stream data:', error);
                    // Continue processing other lines
                }
            }
        }

        // Handle any remaining tool call
        if (currentToolCall) {
            progress.report(new vscode.LanguageModelToolCallPart(
                currentToolCall.id,
                currentToolCall.function.name,
                this.parseToolArguments(currentToolCall.function.arguments)
            ));
        }
    }

    private parseToolArguments(args: string) {
        try {
            return JSON.parse(args || '{}');
        } catch {
            return args || '{}';
        }
    }
}
