/**
 * March Agent SDK - Gateway Client
 * Port of Python march_agent/gateway_client.py
 * 
 * gRPC client for communicating with the Agent Gateway.
 * Includes automatic reconnection with exponential backoff.
 */

import * as grpc from '@grpc/grpc-js'
import * as protoLoader from '@grpc/proto-loader'
import { fileURLToPath } from 'url'
import { dirname, join } from 'path'
import { GatewayError, ReconnectionError } from './exceptions.js'
import {
    DEFAULT_RECONNECTION_OPTIONS,
    type KafkaMessage,
    type ProduceAck,
    type ReconnectionOptions,
    type ConnectionState,
    type ConnectionStateListener,
} from './types.js'

const __filename = fileURLToPath(import.meta.url)
const __dirname = dirname(__filename)

// Load proto file
const PROTO_PATH = join(__dirname, 'proto', 'gateway.proto')

const packageDefinition = protoLoader.loadSync(PROTO_PATH, {
    keepCase: true,
    longs: String,
    enums: String,
    defaults: true,
    oneofs: true,
})

const gatewayProto = grpc.loadPackageDefinition(packageDefinition) as unknown as {
    gateway: {
        AgentGateway: new (
            address: string,
            credentials: grpc.ChannelCredentials,
            options?: grpc.ChannelOptions
        ) => AgentGatewayClient
    }
}

interface AgentGatewayClient {
    AgentStream(): grpc.ClientDuplexStream<ClientMessage, ServerMessage>
}

interface ClientMessage {
    auth?: { api_key: string; agent_names: string[] }
    subscribe?: { agent_name: string }
    unsubscribe?: { agent_name: string }
    produce?: {
        topic: string
        key: string
        headers: Record<string, string>
        body: Buffer
        correlation_id: string
    }
    ping?: { timestamp: string }
}

interface ServerMessage {
    auth_response?: { connection_id: string; subscribed_topics: string[] }
    message?: {
        topic: string
        partition: number
        offset: string
        key: string
        headers: Record<string, string>
        body: Buffer
        timestamp: string
    }
    produce_ack?: {
        topic: string
        partition: number
        offset: string
        correlation_id: string
    }
    subscribe_ack?: { topic: string }
    unsubscribe_ack?: { agent_name: string }
    pong?: { client_timestamp: string; server_timestamp: string }
    error?: { code: string; message: string; correlation_id: string }
}

/**
 * Client for communicating with the Agent Gateway.
 * Provides gRPC bidirectional streaming for Kafka consume/produce.
 * Includes automatic reconnection with exponential backoff.
 */
export class GatewayClient {
    private readonly gatewayUrl: string
    private readonly apiKey: string
    private readonly secure: boolean
    private readonly reconnectionOptions: Required<ReconnectionOptions>

    private client?: AgentGatewayClient
    private channel?: grpc.Channel
    private stream?: grpc.ClientDuplexStream<ClientMessage, ServerMessage>
    private _connectionId?: string
    private messageQueue: KafkaMessage[] = []
    private pendingProduceAcks: Map<string, (ack: ProduceAck) => void> = new Map()
    private pendingSubscribeAcks: Map<string, (topic: string) => void> = new Map()
    private messageHandlers: Map<string, (msg: KafkaMessage) => void> = new Map()
    private correlationCounter: number = 0
    private connected: boolean = false

    // Reconnection state
    private connectionState: ConnectionState = 'disconnected'
    private stateListeners: ConnectionStateListener[] = []
    private agentNames: string[] = []
    private reconnectAttempts: number = 0
    private reconnectTimer?: ReturnType<typeof setTimeout>
    private isReconnecting: boolean = false
    private manualClose: boolean = false
    private lastError?: Error

    constructor(
        gatewayUrl: string,
        apiKey: string,
        secure: boolean = false,
        reconnectionOptions?: ReconnectionOptions
    ) {
        this.gatewayUrl = gatewayUrl
        this.apiKey = apiKey
        this.secure = secure
        this.reconnectionOptions = {
            ...DEFAULT_RECONNECTION_OPTIONS,
            ...reconnectionOptions,
        }
    }

    /**
     * Get current connection state.
     */
    getConnectionState(): ConnectionState {
        return this.connectionState
    }

    /**
     * Register a listener for connection state changes.
     * Returns an unsubscribe function.
     */
    onConnectionStateChange(listener: ConnectionStateListener): () => void {
        this.stateListeners.push(listener)
        return () => {
            const index = this.stateListeners.indexOf(listener)
            if (index !== -1) {
                this.stateListeners.splice(index, 1)
            }
        }
    }

    /**
     * Update connection state and notify listeners.
     */
    private setConnectionState(state: ConnectionState, error?: Error): void {
        const previousState = this.connectionState
        this.connectionState = state

        if (previousState !== state) {
            console.log(`Gateway connection state: ${previousState} -> ${state}`)
            for (const listener of this.stateListeners) {
                try {
                    listener(state, error)
                } catch (e) {
                    console.error('Error in connection state listener:', e)
                }
            }
        }
    }

    /**
     * Calculate backoff delay using exponential backoff with jitter.
     */
    private calculateBackoff(): number {
        const { initialDelayMs, maxDelayMs, backoffMultiplier } = this.reconnectionOptions
        const exponentialDelay = initialDelayMs * Math.pow(backoffMultiplier, this.reconnectAttempts)
        const cappedDelay = Math.min(exponentialDelay, maxDelayMs)
        // Add jitter (0-25% of the delay)
        const jitter = cappedDelay * Math.random() * 0.25
        return Math.floor(cappedDelay + jitter)
    }

    /**
     * HTTP URL for AI Inventory service via proxy.
     */
    get aiInventoryUrl(): string {
        const protocol = this.secure ? 'https' : 'http'
        return `${protocol}://${this.gatewayUrl}/s/ai-inventory`
    }

    /**
     * HTTP URL for Conversation Store service via proxy.
     */
    get conversationStoreUrl(): string {
        const protocol = this.secure ? 'https' : 'http'
        return `${protocol}://${this.gatewayUrl}/s/conversation-store`
    }

    /**
     * HTTP URL for AI Memory service via proxy.
     */
    get aiMemoryUrl(): string {
        const protocol = this.secure ? 'https' : 'http'
        return `${protocol}://${this.gatewayUrl}/s/ai-memory`
    }

    /**
     * HTTP URL for Attachment service via proxy.
     */
    get attachmentUrl(): string {
        const protocol = this.secure ? 'https' : 'http'
        return `${protocol}://${this.gatewayUrl}/s/attachment`
    }

    /**
     * Register a handler for a topic.
     */
    registerHandler(topic: string, handler: (msg: KafkaMessage) => void): void {
        this.messageHandlers.set(topic, handler)
    }

    /**
     * Connect to the gateway and authenticate.
     */
    async connect(agentNames: string[]): Promise<string[]> {
        // Store agent names for reconnection
        this.agentNames = agentNames
        this.manualClose = false
        this.setConnectionState('connecting')

        const credentials = this.secure
            ? grpc.credentials.createSsl()
            : grpc.credentials.createInsecure()

        // gRPC channel options with keepalive settings (matching Python SDK)
        const channelOptions: grpc.ChannelOptions = {
            'grpc.keepalive_time_ms': 30000,           // Send keepalive ping every 30s
            'grpc.keepalive_timeout_ms': 10000,        // Wait 10s for ping ack
            'grpc.keepalive_permit_without_calls': 1,  // Allow pings without active calls
            'grpc.http2.min_time_between_pings_ms': 10000,  // Min time between pings
            'grpc.http2.max_pings_without_data': 0,    // Unlimited pings without data
        }

        this.client = new gatewayProto.gateway.AgentGateway(
            this.gatewayUrl,
            credentials,
            channelOptions
        )

        this.stream = this.client.AgentStream()

        // Set up message handling
        this.stream.on('data', (msg: ServerMessage) => {
            this.handleServerMessage(msg)
        })

        this.stream.on('error', (err: Error) => {
            console.error('Gateway stream error:', err)
            this.lastError = err
            this.connected = false
            this.handleDisconnect(err)
        })

        this.stream.on('end', () => {
            console.log('Gateway stream ended')
            this.connected = false
            this.handleDisconnect(new GatewayError('Stream ended unexpectedly'))
        })

        // Send auth request
        return new Promise((resolve, reject) => {
            const timeout = setTimeout(() => {
                this.setConnectionState('disconnected', new GatewayError('Authentication timeout'))
                reject(new GatewayError('Authentication timeout'))
            }, 10000)

            const authHandler = (msg: ServerMessage) => {
                if (msg.auth_response) {
                    clearTimeout(timeout)
                    this._connectionId = msg.auth_response.connection_id
                    this.connected = true
                    this.reconnectAttempts = 0  // Reset on successful connection
                    this.setConnectionState('connected')
                    resolve(msg.auth_response.subscribed_topics)
                } else if (msg.error) {
                    clearTimeout(timeout)
                    const error = new GatewayError(`Authentication failed: ${msg.error.message}`)
                    this.setConnectionState('disconnected', error)
                    reject(error)
                }
            }

            // Temporarily override handler for auth response
            const originalHandler = this.handleServerMessage.bind(this)
            this.handleServerMessage = (msg: ServerMessage) => {
                if (msg.auth_response || msg.error) {
                    authHandler(msg)
                    this.handleServerMessage = originalHandler
                } else {
                    originalHandler(msg)
                }
            }

            this.stream!.write({
                auth: {
                    api_key: this.apiKey,
                    agent_names: agentNames,
                },
            })
        })
    }

    /**
     * Handle disconnection - triggers reconnection if enabled.
     */
    private handleDisconnect(error: Error): void {
        // Don't reconnect if manually closed
        if (this.manualClose) {
            this.setConnectionState('disconnected')
            return
        }

        // Don't start another reconnection if already reconnecting
        if (this.isReconnecting) {
            return
        }

        // Check if auto-reconnect is enabled
        if (!this.reconnectionOptions.autoReconnect) {
            this.setConnectionState('disconnected', error)
            return
        }

        this.scheduleReconnect(error)
    }

    /**
     * Schedule a reconnection attempt with exponential backoff.
     */
    private scheduleReconnect(error?: Error): void {
        // Check if we've exceeded max retries
        if (this.reconnectAttempts >= this.reconnectionOptions.maxRetries) {
            const reconnectError = new ReconnectionError(
                `Failed to reconnect after ${this.reconnectAttempts} attempts`,
                this.reconnectAttempts,
                this.lastError
            )
            this.setConnectionState('disconnected', reconnectError)
            this.isReconnecting = false
            return
        }

        this.isReconnecting = true
        this.setConnectionState('reconnecting', error)

        const delay = this.calculateBackoff()
        console.log(
            `Scheduling reconnection attempt ${this.reconnectAttempts + 1}/${this.reconnectionOptions.maxRetries} in ${delay}ms`
        )

        this.reconnectTimer = setTimeout(async () => {
            this.reconnectAttempts++
            try {
                await this.reconnect()
            } catch (e) {
                console.error(`Reconnection attempt ${this.reconnectAttempts} failed:`, e)
                this.lastError = e instanceof Error ? e : new Error(String(e))
                // Schedule next attempt
                this.scheduleReconnect(this.lastError)
            }
        }, delay)
    }

    /**
     * Attempt to reconnect to the gateway.
     */
    async reconnect(): Promise<string[]> {
        // Clean up existing connection
        this.cleanupConnection()

        console.log(`Attempting reconnection (attempt ${this.reconnectAttempts})...`)

        try {
            const topics = await this.connect(this.agentNames)
            console.log('Reconnected successfully to gateway')
            this.isReconnecting = false
            return topics
        } catch (error) {
            throw error
        }
    }

    /**
     * Force a reconnection attempt, resetting the attempt counter.
     * Useful when you want to manually trigger reconnection.
     */
    async forceReconnect(): Promise<string[]> {
        this.reconnectAttempts = 0
        this.isReconnecting = false
        this.cancelReconnect()
        return this.reconnect()
    }

    /**
     * Cancel any pending reconnection attempt.
     */
    cancelReconnect(): void {
        if (this.reconnectTimer) {
            clearTimeout(this.reconnectTimer)
            this.reconnectTimer = undefined
        }
        this.isReconnecting = false
    }

    /**
     * Clean up the current connection without triggering reconnection.
     */
    private cleanupConnection(): void {
        if (this.stream) {
            try {
                this.stream.removeAllListeners()
                this.stream.end()
            } catch (e) {
                // Ignore cleanup errors
            }
            this.stream = undefined
        }
        this.client = undefined
        this.connected = false
        this._connectionId = undefined
    }

    /**
     * Handle incoming server messages.
     */
    private handleServerMessage(msg: ServerMessage): void {
        if (msg.message) {
            const kafkaMsg: KafkaMessage = {
                topic: msg.message.topic,
                partition: msg.message.partition,
                offset: parseInt(msg.message.offset, 10),
                key: msg.message.key,
                headers: msg.message.headers,
                body: JSON.parse(msg.message.body.toString()),
                timestamp: parseInt(msg.message.timestamp, 10),
            }

            // Check for registered handler
            const handler = this.messageHandlers.get(kafkaMsg.topic)
            if (handler) {
                handler(kafkaMsg)
            } else {
                // Queue message if no handler
                this.messageQueue.push(kafkaMsg)
            }
        } else if (msg.produce_ack) {
            const callback = this.pendingProduceAcks.get(msg.produce_ack.correlation_id)
            if (callback) {
                callback({
                    topic: msg.produce_ack.topic,
                    partition: msg.produce_ack.partition,
                    offset: parseInt(msg.produce_ack.offset, 10),
                    correlationId: msg.produce_ack.correlation_id,
                })
                this.pendingProduceAcks.delete(msg.produce_ack.correlation_id)
            }
        } else if (msg.subscribe_ack) {
            const callback = this.pendingSubscribeAcks.get(msg.subscribe_ack.topic)
            if (callback) {
                callback(msg.subscribe_ack.topic)
                this.pendingSubscribeAcks.delete(msg.subscribe_ack.topic)
            }
        } else if (msg.error) {
            console.error('Gateway error:', msg.error.message)
        }
    }

    /**
     * Subscribe to an additional agent's topic.
     */
    async subscribe(agentName: string): Promise<string> {
        if (!this.stream || !this.connected) {
            throw new GatewayError('Not connected')
        }

        return new Promise((resolve, reject) => {
            const timeout = setTimeout(() => {
                reject(new GatewayError('Subscribe timeout'))
            }, 5000)

            const expectedTopic = `${agentName}.inbox`
            this.pendingSubscribeAcks.set(expectedTopic, (topic) => {
                clearTimeout(timeout)
                resolve(topic)
            })

            this.stream!.write({
                subscribe: { agent_name: agentName },
            })
        })
    }

    /**
     * Unsubscribe from an agent's topic.
     */
    unsubscribe(agentName: string): void {
        if (!this.stream || !this.connected) {
            throw new GatewayError('Not connected')
        }

        this.stream.write({
            unsubscribe: { agent_name: agentName },
        })
    }

    /**
     * Produce a message to Kafka via the gateway.
     */
    produce(
        topic: string,
        key: string,
        headers: Record<string, string>,
        body: Record<string, unknown>,
        correlationId?: string
    ): void {
        if (!this.stream || !this.connected) {
            throw new GatewayError('Not connected')
        }

        const corrId = correlationId ?? `${++this.correlationCounter}`

        this.stream.write({
            produce: {
                topic,
                key,
                headers,
                body: Buffer.from(JSON.stringify(body)),
                correlation_id: corrId,
            },
        })
    }

    /**
     * Produce a message and wait for acknowledgment.
     */
    async produceAndWait(
        topic: string,
        key: string,
        headers: Record<string, string>,
        body: Record<string, unknown>
    ): Promise<ProduceAck> {
        if (!this.stream || !this.connected) {
            throw new GatewayError('Not connected')
        }

        const correlationId = `${++this.correlationCounter}`

        return new Promise((resolve, reject) => {
            const timeout = setTimeout(() => {
                this.pendingProduceAcks.delete(correlationId)
                reject(new GatewayError('Produce timeout'))
            }, 10000)

            this.pendingProduceAcks.set(correlationId, (ack) => {
                clearTimeout(timeout)
                resolve(ack)
            })

            this.stream!.write({
                produce: {
                    topic,
                    key,
                    headers,
                    body: Buffer.from(JSON.stringify(body)),
                    correlation_id: correlationId,
                },
            })
        })
    }

    /**
     * Consume a single message (polling from queue).
     */
    consumeOne(_timeout: number = 1000): KafkaMessage | null {
        if (this.messageQueue.length > 0) {
            return this.messageQueue.shift()!
        }
        return null
    }

    /**
     * Send a ping to the gateway.
     */
    ping(): void {
        if (!this.stream || !this.connected) {
            throw new GatewayError('Not connected')
        }

        this.stream.write({
            ping: { timestamp: String(Date.now()) },
        })
    }

    /**
     * Make a sync POST request (used for registration).
     */
    async httpPost(
        service: string,
        path: string,
        body: unknown
    ): Promise<Response> {
        const protocol = this.secure ? 'https' : 'http'
        const url = `${protocol}://${this.gatewayUrl}/s/${service}${path}`

        return fetch(url, {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
                'X-API-Key': this.apiKey,
            },
            body: JSON.stringify(body),
        })
    }

    /**
     * Check if connected.
     */
    isConnected(): boolean {
        return this.connected
    }

    /**
     * Get the last error that occurred.
     */
    getLastError(): Error | undefined {
        return this.lastError
    }

    /**
     * Get the current reconnection attempt count.
     */
    getReconnectAttempts(): number {
        return this.reconnectAttempts
    }

    /**
     * Check if currently attempting to reconnect.
     */
    isReconnectingNow(): boolean {
        return this.isReconnecting
    }

    /**
     * Close the gateway connection.
     * This will prevent automatic reconnection.
     */
    close(): void {
        this.manualClose = true
        this.cancelReconnect()
        this.cleanupConnection()
        this.setConnectionState('disconnected')
    }
}