import * as vscode from 'vscode';
import { ApiKeyManager } from './apiKeyManager';
import { ZAIAPIClient } from './apiClient';
import { ZAI_MODELS } from './models';

export class ZAIChatModelProvider implements vscode.LanguageModelChatProvider {
    private apiClient: ZAIAPIClient;

    constructor(private apiKeyManager: ApiKeyManager) {
        this.apiClient = new ZAIAPIClient(apiKeyManager);
    }

    /**
     * Refresh the provider (e.g., when API key changes)
     */
    async refresh(): Promise<void> {
        // Trigger a re-fetch of models by creating a new API client
        this.apiClient = new ZAIAPIClient(this.apiKeyManager);
    }

    /**
     * Provide information about available language models
     */
    async provideLanguageModelChatInformation(
        options: { silent: boolean },
        token: vscode.CancellationToken
    ): Promise<vscode.LanguageModelChatInformation[]> {
        
        // Check if we have an API key
        const hasKey = await this.apiKeyManager.hasApiKey();
        
        if (!hasKey) {
            if (!options.silent) {
                const setupAction = 'Setup API Key';
                const result = await vscode.window.showWarningMessage(
                    'Z.AI API key not configured. Please set up your API key to use Z.AI models.',
                    setupAction
                );

                if (result === setupAction) {
                    await vscode.commands.executeCommand('zai.manageApiKey');
                }
            }
            return [];
        }

        // Return the available models
        return ZAI_MODELS.map(model => ({
            id: model.id,
            name: model.name,
            family: model.family,
            version: model.version,
            maxInputTokens: model.maxInputTokens,
            maxOutputTokens: model.maxOutputTokens,
            tooltip: model.tooltip,
            detail: model.detail,
            capabilities: {
                imageInput: model.capabilities.imageInput,
                toolCalling: model.capabilities.toolCalling
            }
        }));
    }

    /**
     * Handle chat requests and stream responses
     */
    async provideLanguageModelChatResponse(
        model: vscode.LanguageModelChatInformation,
        messages: readonly vscode.LanguageModelChatRequestMessage[],
        options: vscode.ProvideLanguageModelChatResponseOptions,
        progress: vscode.Progress<vscode.LanguageModelResponsePart>,
        token: vscode.CancellationToken
    ): Promise<void> {
        
        // Get the API key
        const apiKey = await this.apiKeyManager.getApiKey();
        if (!apiKey) {
            throw new Error('Z.AI API key not configured');
        }

        try {
            // Make the API request and stream the response
            await this.apiClient.streamChatCompletion(
                model.id,
                messages,
                options,
                progress,
                token
            );
        } catch (error) {
            if (error instanceof Error) {
                if (error.message.includes('401') || error.message.includes('403')) {
                    vscode.window.showErrorMessage(
                        'Invalid Z.AI API key. Please check your API key configuration.',
                        'Setup API Key'
                    ).then(action => {
                        if (action === 'Setup API Key') {
                            vscode.commands.executeCommand('zai.manageApiKey');
                        }
                    });
                }
                throw error;
            }
            throw new Error(`Unexpected error: ${String(error)}`);
        }
    }

    /**
     * Count tokens in a message
     */
    async provideTokenCount(
        model: vscode.LanguageModelChatInformation,
        text: string | vscode.LanguageModelChatRequestMessage,
        token: vscode.CancellationToken
    ): Promise<number> {
        
        // Simple estimation: ~4 characters per token
        // This is a rough approximation; ideally you'd use a tokenizer
        let textToCount: string;
        
        if (typeof text === 'string') {
            textToCount = text;
        } else {
            // Convert message content to string
            textToCount = text.content
                .filter(part => part instanceof vscode.LanguageModelTextPart)
                .map(part => (part as vscode.LanguageModelTextPart).value)
                .join('');
        }

        // Approximate token count (4 chars per token is a common estimate)
        return Math.ceil(textToCount.length / 4);
    }
}
