/**
 * Secure Completion Client
 * Main client class for encrypted communication with NOMYO router
 * 
 * Port of Python's SecureCompletionClient with full API compatibility
 */

import { ClientConfig } from '../types/client';
import { EncryptedPackage } from '../types/crypto';
import { KeyManager } from './crypto/keys';
import { AESEncryption } from './crypto/encryption';
import { RSAOperations } from './crypto/rsa';
import { createHttpClient, HttpClient } from './http/client';
import { createSecureMemory, SecureByteContext } from './memory/secure';
import { SecurityError, APIConnectionError } from '../errors';
import {
    arrayBufferToBase64,
    base64ToArrayBuffer,
    stringToArrayBuffer,
    arrayBufferToString,
} from './crypto/utils';

export class SecureCompletionClient {
    private routerUrl: string;
    private allowHttp: boolean;
    private secureMemory: boolean;
    private keyManager: KeyManager;
    private aes: AESEncryption;
    private rsa: RSAOperations;
    private httpClient: HttpClient;
    private secureMemoryImpl = createSecureMemory();

    constructor(config: ClientConfig = { routerUrl: 'https://api.nomyo.ai:12434' }) {
        const {
            routerUrl = 'https://api.nomyo.ai:12434',
            allowHttp = false,
            secureMemory = true,
            keySize = 4096,
        } = config;

        this.routerUrl = routerUrl.replace(/\/$/, ''); // Remove trailing slash
        this.allowHttp = allowHttp;
        this.secureMemory = secureMemory;

        // Validate HTTPS for security
        if (!this.routerUrl.startsWith('https://')) {
            if (!allowHttp) {
                console.warn(
                    '⚠️  WARNING: Using HTTP instead of HTTPS. ' +
                    'This is INSECURE and should only be used for local development. ' +
                    'Man-in-the-middle attacks are possible!'
                );
            } else {
                console.log('HTTP mode enabled for local development (INSECURE)');
            }
        }

        // Initialize components
        this.keyManager = new KeyManager();
        this.aes = new AESEncryption();
        this.rsa = new RSAOperations();
        this.httpClient = createHttpClient();

        // Log memory protection info
        const protectionInfo = this.secureMemoryImpl.getProtectionInfo();
        console.log(`Memory protection: ${protectionInfo.method} (${protectionInfo.details})`);
    }

    /**
     * Generate RSA key pair
     */
    async generateKeys(options: {
        saveToFile?: boolean;
        keyDir?: string;
        password?: string;
    } = {}): Promise<void> {
        await this.keyManager.generateKeys({
            keySize: 4096,
            ...options,
        });
    }

    /**
     * Load existing keys from files (Node.js only)
     */
    async loadKeys(
        privateKeyPath: string,
        publicKeyPath?: string,
        password?: string
    ): Promise<void> {
        await this.keyManager.loadKeys(
            { privateKeyPath, publicKeyPath },
            password
        );
    }

    /**
     * Ensure keys are loaded, generate if necessary
     */
    private async ensureKeys(): Promise<void> {
        if (this.keyManager.hasKeys()) {
            return;
        }

        // Try to load keys from default location (Node.js only)
        if (typeof window === 'undefined') {
            try {
                const fs = require('fs').promises;
                const path = require('path');

                const privateKeyPath = path.join('client_keys', 'private_key.pem');
                const publicKeyPath = path.join('client_keys', 'public_key.pem');

                // Check if keys exist
                await fs.access(privateKeyPath);
                await fs.access(publicKeyPath);

                // Load keys
                await this.loadKeys(privateKeyPath, publicKeyPath);
                console.log('Loaded existing keys from client_keys/');
                return;
            } catch (error) {
                // Keys don't exist, generate new ones
                console.log('No existing keys found, generating new keys...');
            }
        }

        // Generate new keys
        await this.generateKeys({
            saveToFile: typeof window === 'undefined', // Only save in Node.js
            keyDir: 'client_keys',
        });
    }

    /**
     * Fetch server's public key from /pki/public_key endpoint
     */
    async fetchServerPublicKey(): Promise<string> {
        console.log("Fetching server's public key...");

        // Security check: Ensure HTTPS is used unless HTTP explicitly allowed
        if (!this.routerUrl.startsWith('https://')) {
            if (!this.allowHttp) {
                throw new SecurityError(
                    'Server public key must be fetched over HTTPS to prevent MITM attacks. ' +
                    'For local development, initialize with allowHttp=true: ' +
                    'new SecureChatCompletion({ baseUrl: "http://localhost:12434", allowHttp: true })'
                );
            } else {
                console.warn('Fetching key over HTTP (local development mode)');
            }
        }

        const url = `${this.routerUrl}/pki/public_key`;

        try {
            const response = await this.httpClient.get(url, { timeout: 60000 });

            if (response.statusCode === 200) {
                const serverPublicKey = arrayBufferToString(response.body);

                // Validate it's a valid PEM key
                try {
                    await this.rsa.importPublicKey(serverPublicKey);
                } catch (error) {
                    throw new Error('Server returned invalid public key format');
                }

                if (this.routerUrl.startsWith('https://')) {
                    console.log("Server's public key fetched securely over HTTPS");
                } else {
                    console.warn("Server's public key fetched over HTTP (INSECURE)");
                }

                return serverPublicKey;
            } else {
                throw new Error(`Failed to fetch server's public key: HTTP ${response.statusCode}`);
            }
        } catch (error) {
            if (error instanceof SecurityError) {
                throw error;
            }
            if (error instanceof Error) {
                throw new APIConnectionError(`Failed to fetch server's public key: ${error.message}`);
            }
            throw error;
        }
    }

    /**
     * Encrypt a payload using hybrid encryption (AES-256-GCM + RSA-OAEP)
     */
    async encryptPayload(payload: object): Promise<ArrayBuffer> {
        console.log('Encrypting payload...');

        // Validate payload
        if (!payload || typeof payload !== 'object') {
            throw new Error('Payload must be an object');
        }

        // Ensure keys are loaded
        await this.ensureKeys();

        // Serialize payload to JSON
        const payloadJson = JSON.stringify(payload);
        const payloadBytes = stringToArrayBuffer(payloadJson);

        // Validate payload size (prevent DoS)
        const MAX_PAYLOAD_SIZE = 10 * 1024 * 1024; // 10MB limit
        if (payloadBytes.byteLength > MAX_PAYLOAD_SIZE) {
            throw new Error(`Payload too large: ${payloadBytes.byteLength} bytes (max: ${MAX_PAYLOAD_SIZE})`);
        }

        console.log(`Payload size: ${payloadBytes.byteLength} bytes`);

        // Use secure memory context if enabled
        if (this.secureMemory) {
            const context = new SecureByteContext(payloadBytes, true);
            return await context.use(async (protectedPayload) => {
                return await this.performEncryption(protectedPayload);
            });
        } else {
            return await this.performEncryption(payloadBytes);
        }
    }

    /**
     * Perform the actual encryption (separated for secure memory context)
     */
    private async performEncryption(payloadBytes: ArrayBuffer): Promise<ArrayBuffer> {
        // Generate AES key
        const aesKey = await this.aes.generateKey();
        const aesKeyBytes = await this.aes.exportKey(aesKey);

        // Protect AES key in memory
        const aesContext = new SecureByteContext(aesKeyBytes, this.secureMemory);
        return await aesContext.use(async (protectedAesKey) => {
            // Encrypt payload with AES-GCM
            const { ciphertext, nonce } = await this.aes.encrypt(payloadBytes, aesKey);

            // Fetch server's public key
            const serverPublicKeyPem = await this.fetchServerPublicKey();
            const serverPublicKey = await this.rsa.importPublicKey(serverPublicKeyPem);

            // Encrypt AES key with server's RSA public key
            const encryptedAesKey = await this.rsa.encryptKey(protectedAesKey, serverPublicKey);

            // Create encrypted package (matching Python format)
            const encryptedPackage = {
                version: '1.0',
                algorithm: 'hybrid-aes256-rsa4096',
                encrypted_payload: {
                    ciphertext: arrayBufferToBase64(ciphertext),
                    nonce: arrayBufferToBase64(nonce),
                    // Note: GCM tag is included in ciphertext in Web Crypto API
                },
                encrypted_aes_key: arrayBufferToBase64(encryptedAesKey),
                key_algorithm: 'RSA-OAEP-SHA256',
                payload_algorithm: 'AES-256-GCM',
            };

            // Serialize package to JSON
            const packageJson = JSON.stringify(encryptedPackage);
            const packageBytes = stringToArrayBuffer(packageJson);

            console.log(`Encrypted package size: ${packageBytes.byteLength} bytes`);

            return packageBytes;
        });
    }

    /**
     * Decrypt a response from the secure endpoint
     */
    async decryptResponse(encryptedResponse: ArrayBuffer, payloadId: string): Promise<object> {
        console.log('Decrypting response...');

        // Validate input
        if (!encryptedResponse || encryptedResponse.byteLength === 0) {
            throw new Error('Empty encrypted response');
        }

        // Parse encrypted package
        let packageData: any;
        try {
            const packageJson = arrayBufferToString(encryptedResponse);
            packageData = JSON.parse(packageJson);
        } catch (error) {
            throw new Error('Invalid encrypted package format: malformed JSON');
        }

        // Validate package structure
        const requiredFields = ['version', 'algorithm', 'encrypted_payload', 'encrypted_aes_key'];
        for (const field of requiredFields) {
            if (!(field in packageData)) {
                throw new Error(`Missing required field in encrypted package: ${field}`);
            }
        }

        // Validate encrypted_payload structure
        const payloadRequired = ['ciphertext', 'nonce'];
        for (const field of payloadRequired) {
            if (!(field in packageData.encrypted_payload)) {
                throw new Error(`Missing field in encrypted_payload: ${field}`);
            }
        }

        // Decrypt AES key with private key
        try {
            const encryptedAesKey = base64ToArrayBuffer(packageData.encrypted_aes_key);
            const privateKey = this.keyManager.getPrivateKey();
            const aesKeyBytes = await this.rsa.decryptKey(encryptedAesKey, privateKey);

            // Use secure memory context for AES key
            const aesContext = new SecureByteContext(aesKeyBytes, this.secureMemory);
            const response = await aesContext.use(async (protectedAesKey) => {
                // Import AES key
                const aesKey = await this.aes.importKey(protectedAesKey);

                // Decrypt payload with AES-GCM
                const ciphertext = base64ToArrayBuffer(packageData.encrypted_payload.ciphertext);
                const nonce = base64ToArrayBuffer(packageData.encrypted_payload.nonce);

                const plaintext = await this.aes.decrypt(ciphertext, nonce, aesKey);

                // Use secure memory for plaintext
                const plaintextContext = new SecureByteContext(plaintext, this.secureMemory);
                return await plaintextContext.use(async (protectedPlaintext) => {
                    // Parse decrypted response
                    const responseJson = arrayBufferToString(protectedPlaintext);
                    return JSON.parse(responseJson);
                });
            });

            // Add metadata
            if (!response._metadata) {
                response._metadata = {};
            }
            response._metadata = {
                ...response._metadata,
                payload_id: payloadId,
                processed_at: packageData.processed_at,
                is_encrypted: true,
                encryption_algorithm: packageData.algorithm,
            };

            console.log('Response decrypted successfully');
            return response;
        } catch (error) {
            // Don't leak specific decryption errors (timing attacks)
            throw new SecurityError('Decryption failed: integrity check or authentication failed');
        }
    }

    /**
     * Send a secure chat completion request to the router
     */
    async sendSecureRequest(
        payload: object,
        payloadId: string,
        apiKey?: string
    ): Promise<object> {
        console.log('Sending secure chat completion request...');

        // Ensure keys are loaded
        await this.ensureKeys();

        // Step 1: Encrypt the payload
        const encryptedPayload = await this.encryptPayload(payload);

        // Step 2: Prepare headers
        const publicKeyPem = await this.keyManager.getPublicKeyPEM();
        const headers: Record<string, string> = {
            'X-Payload-ID': payloadId,
            'X-Public-Key': encodeURIComponent(publicKeyPem),
            'Content-Type': 'application/octet-stream',
        };

        // Add Authorization header if api_key is provided
        if (apiKey) {
            headers['Authorization'] = `Bearer ${apiKey}`;
        }

        // Step 3: Send request to router
        const url = `${this.routerUrl}/v1/chat/secure_completion`;
        console.log(`Target URL: ${url}`);

        try {
            const response = await this.httpClient.post(url, {
                headers,
                body: encryptedPayload,
                timeout: 60000,
            });

            console.log(`HTTP Status: ${response.statusCode}`);

            if (response.statusCode === 200) {
                // Step 4: Decrypt the response
                const decryptedResponse = await this.decryptResponse(response.body, payloadId);
                return decryptedResponse;
            } else {
                // Handle error responses
                const { handleErrorResponse } = await import('../errors');
                throw this.handleErrorResponse(response);
            }
        } catch (error) {
            if (error instanceof Error) {
                if (error.message === 'Request timeout') {
                    throw new APIConnectionError('Connection to server timed out');
                }
                throw new APIConnectionError(`Failed to connect to router: ${error.message}`);
            }
            throw error;
        }
    }

    /**
     * Handle error HTTP responses
     */
    private handleErrorResponse(response: { statusCode: number; body: ArrayBuffer }): Error {
        const {
            AuthenticationError,
            InvalidRequestError,
            RateLimitError,
            ServerError,
            APIError,
        } = require('../errors');

        let errorData: any = {};
        try {
            const errorJson = arrayBufferToString(response.body);
            errorData = JSON.parse(errorJson);
        } catch (e) {
            // Ignore JSON parse errors
        }

        const detail = errorData.detail || 'Unknown error';

        switch (response.statusCode) {
            case 400:
                return new InvalidRequestError(`Bad request: ${detail}`, 400, errorData);
            case 401:
                return new AuthenticationError(`Invalid API key or authentication failed: ${detail}`, 401, errorData);
            case 404:
                return new APIError(`Endpoint not found: ${detail}`, 404, errorData);
            case 429:
                return new RateLimitError(`Rate limit exceeded: ${detail}`, 429, errorData);
            case 500:
                return new ServerError(`Server error: ${detail}`, 500, errorData);
            default:
                return new APIError(`Unexpected status code: ${response.statusCode}`, response.statusCode, errorData);
        }
    }
}