import * as vscode from 'vscode';

export interface InlineRequest {
    prompt: string;
    suffix: string;
    language: string;
    filename: string;
    max_tokens?: number;
}

export interface InlineResponse {
    completion: string;
    finish_reason: string;
    suggestionId?: string;
    similarityToTraining?: number;
    nearestNeighborRepo?: string | null;
    confidence?: 'original' | 'derivative' | 'verbatim_fragment';
}

export interface ChatMessage {
    role: 'user' | 'assistant' | 'system';
    content: string;
}

export interface ChatRequest {
    messages: ChatMessage[];
    model?: string;
    stream?: boolean;
    temperature?: number;
    max_tokens?: number;
}

export interface ChatResponse {
    message: ChatMessage;
    finish_reason: string;
    usage?: {
        prompt_tokens: number;
        completion_tokens: number;
        total_tokens: number;
    };
}

export interface AuditEntry {
    timestamp: number;
    type: 'inline' | 'chat';
    suggestionId?: string;
    similarityToTraining?: number;
    nearestNeighborRepo?: string | null;
    confidence?: 'original' | 'derivative' | 'verbatim_fragment';
    accepted?: boolean;
    blockedGeo?: boolean;
    geoCountry?: string;
    error?: string;
}

export interface AbysiusConfig {
    apiKey: string;
    chatEndpoint: string;
    inlineEndpoint: string;
    model: string;
    dataResidency?: 'auto' | 'us' | 'eu' | 'asia';
    enableAuditLog?: boolean;
    geoBlockRestrictedRegions?: boolean;
    anonymizeFilePaths?: boolean;
}

export class AbysiusApi {
    private config: AbysiusConfig;
    private auditLog: AuditEntry[] = [];
    private readonly RESTRICTED_COUNTRIES = ['CU', 'IR', 'KP', 'RU', 'SY'];

    constructor(config: AbysiusConfig) {
        this.config = config;
    }

    updateConfig(config: Partial<AbysiusConfig>) {
        this.config = { ...this.config, ...config };
    }

    private getHeaders(): Record<string, string> {
        const headers: Record<string, string> = {
            'Content-Type': 'application/json',
            'User-Agent': 'AbysiusCodium/0.1.0'
        };
        if (this.config.apiKey) {
            headers['Authorization'] = `Bearer ${this.config.apiKey}`;
        }
        if (this.config.dataResidency && this.config.dataResidency !== 'auto') {
            headers['X-Data-Residency'] = this.config.dataResidency;
        }
        return headers;
    }

    private async geoCheck(): Promise<{ blocked: boolean; country?: string }> {
        if (this.config.geoBlockRestrictedRegions === false) {
            return { blocked: false };
        }
        // In production, derive from API response or client IP context.
        // Returning blocked=false here; server implements actual geo-blocking.
        return { blocked: false };
    }

    private writeAudit(entry: AuditEntry): void {
        if (!this.config.enableAuditLog) return;
        this.auditLog.push(entry);
        // Trim audit log to last 10,000 entries to prevent unbounded growth
        if (this.auditLog.length > 10000) {
            this.auditLog = this.auditLog.slice(-10000);
        }
        // In production, optionally persist to file or stream to a compliance endpoint
    }

    getAuditLog(): ReadonlyArray<AuditEntry> {
        return this.auditLog;
    }

    clearAuditLog(): void {
        this.auditLog = [];
    }

    private anonymizeFilename(filename: string): string {
        if (this.config.anonymizeFilePaths !== false) {
            return filename.split(/[\\/]/).pop() || filename;
        }
        return filename;
    }

    async getInlineCompletion(request: InlineRequest, signal?: AbortSignal): Promise<InlineResponse | null> {
        try {
            const config = vscode.workspace.getConfiguration('abysius');
            const maxLength = config.get<number>('inlineCompletionMaxLength', 200);

            const response = await fetch(this.config.inlineEndpoint, {
                method: 'POST',
                headers: this.getHeaders(),
                body: JSON.stringify({
                    ...request,
                    model: this.config.model,
                    max_tokens: Math.min(request.max_tokens || maxLength, maxLength)
                }),
                signal
            });

            if (!response.ok) {
                const error = await response.text();
                console.error('[Abysius API] Inline completion error:', error);
                return null;
            }

            return await response.json() as InlineResponse;
        } catch (err) {
            if (err instanceof Error && err.name === 'AbortError') {
                return null;
            }
            console.error('[Abysius API] Inline completion failed:', err);
            return null;
        }
    }

    async *streamChat(request: ChatRequest, signal?: AbortSignal): AsyncGenerator<string, ChatResponse | null, unknown> {
        try {
            const response = await fetch(this.config.chatEndpoint, {
                method: 'POST',
                headers: this.getHeaders(),
                body: JSON.stringify({
                    ...request,
                    model: this.config.model,
                    stream: true
                }),
                signal
            });

            if (!response.ok) {
                const error = await response.text();
                console.error('[Abysius API] Chat error:', error);
                return null;
            }

            const reader = response.body?.getReader();
            if (!reader) return null;

            const decoder = new TextDecoder();
            let buffer = '';

            try {
                while (true) {
                    const { done, value } = await reader.read();
                    if (done) break;

                    buffer += decoder.decode(value, { stream: true });
                    const lines = buffer.split('\n');
                    buffer = lines.pop() || '';

                    for (const line of lines) {
                        const trimmed = line.trim();
                        if (!trimmed || !trimmed.startsWith('data: ')) continue;
                        
                        const data = trimmed.slice(6);
                        if (data === '[DONE]') continue;

                        try {
                            const chunk = JSON.parse(data);
                            const delta = chunk.choices?.[0]?.delta?.content;
                            if (delta) {
                                yield delta;
                            }
                        } catch {
                            // ignore malformed JSON
                        }
                    }
                }
            } finally {
                reader.releaseLock();
            }

            return null;
        } catch (err) {
            if (err instanceof Error && err.name === 'AbortError') {
                return null;
            }
            console.error('[Abysius API] Chat stream failed:', err);
            return null;
        }
    }

    async sendChat(request: ChatRequest, signal?: AbortSignal): Promise<ChatResponse | null> {
        try {
            const response = await fetch(this.config.chatEndpoint, {
                method: 'POST',
                headers: this.getHeaders(),
                body: JSON.stringify({
                    ...request,
                    model: this.config.model,
                    stream: false
                }),
                signal
            });

            if (!response.ok) {
                const error = await response.text();
                console.error('[Abysius API] Chat error:', error);
                return null;
            }

            return await response.json() as ChatResponse;
        } catch (err) {
            if (err instanceof Error && err.name === 'AbortError') {
                return null;
            }
            console.error('[Abysius API] Chat failed:', err);
            return null;
        }
    }
}