import { createClient, RedisClientOptions } from 'redis'; // Import Redis client and types
import { ProcessOutcome } from './WebSocketAPI'; // Import ProcessOutcome type from WebSocketAPI
import { ZeebeGrpcClient } from '@camunda8/sdk/dist/zeebe'; // Import Zeebe gRPC client from Camunda SDK
import { JSONDoc } from '@camunda8/sdk/dist/zeebe/lib/interfaces-1.0'; // Import JSONDoc type from Camunda SDK
import { Duration } from 'typed-duration'; // Import Duration for time-to-live settings
import { DeepPartial } from './ZeebeAffinityClient'; // Import DeepPartial type from ZeebeAffinityClient

type JSONExtended = { [key: string]: JSONDoc };
type RedisClientConnection = ReturnType<typeof createClient>;

process.on('unhandledRejection', (reason, promise) => {
    console.error('Unhandled Rejection at:', promise, 'reason:', reason);
});

export class RedisAffinity extends ZeebeGrpcClient {
    public subscriber: RedisClientConnection; // Redis client for subscribing to channels
    public publisher: RedisClientConnection; // Redis client for publishing messages
    affinityCallbacks: {
        [processInstanceKey: string]: (processOutcome: ProcessOutcome) => void; // Callbacks for process outcomes
    };
    private useLegacyMode: boolean;

    constructor(options: {
        config?: DeepPartial<{
            ZEEBE_ADDRESS: string
        }>;
    }, redisOptions: RedisClientOptions & { legacyMode?: boolean }) {
        super(options);
        this.useLegacyMode = redisOptions.legacyMode ?? false;

        const createClientOptions = this.useLegacyMode
            // Redis with legacy mode
            ? { ...redisOptions, legacyMode: true }
            // Redis without legacy mode
            : redisOptions
        ;

        this.subscriber = createClient(createClientOptions);
        this.publisher = createClient(createClientOptions);

        this.affinityCallbacks = {}; // Initialize callbacks map

        this.initializeRedisClients();
    }

    private async initializeRedisClients() {
        try {
            await this.subscriber.connect();
            console.log('Subscriber connected');
        } catch (error) {
            console.error(`Subscriber connection error: ${error}`);
            throw error; // Optionally throw the error or handle it according to your application's logic
        }

        try {
            await this.publisher.connect();
            console.log('Publisher connected');
        } catch (error) {
            console.error(`Publisher connection error: ${error}`);
            throw error; // Optionally throw the error or handle it according to your application's logic
        }

        this.subscriber.on('error', (error) => {
            console.error(`Subscriber error: ${error}`);
        });

        this.publisher.on('error', (error) => {
            console.error(`Publisher error: ${error}`);
        });
        this.subscriber.on('message', (channel: string, message: string) => {
            this.handleMessage(channel, message);
        });
        
    }

    private handleMessage(channel: string, message: string) {
        console.log(`Subscriber received message in channel ${channel}`);
        try {
            if (this.affinityCallbacks[channel]) {
                this.affinityCallbacks[channel](JSON.parse(message)); // Call registered callback
            } else {
                console.warn(`No callback registered for channel ${channel}`);
            }
        } catch (err) {
            console.error(err);
        }
    }

    async createAffinityWorker(taskType: string): Promise<void> {
        super.createWorker({
            taskType,
            taskHandler: async (job) => {
                try {
                    const stringifiedJobVars = JSON.stringify(job?.variables);

                    if (this.useLegacyMode) {
                       await this.publisher.v4.publish(
                            job.processInstanceKey,
                            stringifiedJobVars,
                        );
                    } else {
                        await this.publisher.publish(
                            job.processInstanceKey,
                            stringifiedJobVars,
                        );
                    }
                    console.log(`Publish message on channel: ${job.processInstanceKey}`);
                    return await job.complete();
                } catch (error: unknown) {
                    if (error instanceof Error) {
                        console.error(`Error while publishing message on channel: ${job.processInstanceKey}`);
                        return job.fail(error.message);
                    } else {
                        throw error;
                    }
                }
            },
        });
    }

    async createProcessInstanceWithAffinity<Variables extends JSONDoc>({
        bpmnProcessId,
        variables,
        cb,
    }: {
        bpmnProcessId: string;
        variables: Variables;
        cb: (processOutcome: ProcessOutcome) => void;
    }): Promise<void> {
        try {
            const wfi = await super.createProcessInstance({
                bpmnProcessId,
                variables,
            });

            this.affinityCallbacks[wfi.processInstanceKey] = cb; // Register callback

            await this.subscribeToChannel(wfi.processInstanceKey); // Ensure async operation
        } catch (err) {
            console.error(err);
            throw err;
        }
    }

    async publishMessageWithAffinity<Variables extends JSONExtended>({
        correlationKey,
        messageId,
        name,
        variables,
        processInstanceKey,
        cb,
    }: {
        correlationKey: string;
        messageId: string;
        name: string;
        variables: Variables;
        processInstanceKey: string;
        cb: (processOutcome: ProcessOutcome) => void;
    }): Promise<void> {
        try {
            await super.publishMessage({
                correlationKey,
                messageId,
                name,
                variables,
                timeToLive: Duration.seconds.of(10),
            });

            this.affinityCallbacks[processInstanceKey] = cb;

            await this.subscribeToChannel(processInstanceKey); // Ensure async operation
        } catch (error) {
            console.error(`Error while publishing message with affinity: ${error}`);
            throw error;
        }
    }

    private async subscribeToChannel(channel: string) {
        try {
            if (this.useLegacyMode) {
                await this.subscriber.v4.subscribe(channel, (message, channelName) => {
                    this.handleMessage(channelName, message);
                });
            } else {
                await this.subscriber.subscribe(channel, (message, channelName) => {
                    this.handleMessage(channelName, message);
                });
            }
            console.log(`Subscribed to channel ${channel}`);
        } catch (error) {
            console.error(`Subscription error for channel ${channel}: ${error}`);
            throw error;
        }
    }

    cleanup(channel: string): void {
        console.log(`Unsubscribe from channel ${channel} and removing affinity callbacks.`);
        if (this.useLegacyMode) {
            this.subscriber.v4.unsubscribe(channel);
        } else {
            this.subscriber.unsubscribe(channel);
        }
        try {
            delete this.affinityCallbacks[channel];
        } catch (err) {
            console.error(err);
        }
    }
}