/**
 * Copyright 2023 Kapeta Inc.
 * SPDX-License-Identifier: BUSL-1.1
 */

import uuid from 'node-uuid';
import { StormClient, UIPagePrompt } from './stormClient';
import { ReferenceClassification, StormEvent, StormEventPage, StormImage, UIShell } from './events';
import { EventEmitter } from 'node:events';
import PQueue from 'p-queue';

import { hasPageOnDisk, normalizePath, writeImageToDisk } from './page-utils';
import * as mimetypes from 'mime-types';
import { randomUUID } from 'node:crypto';

export interface ImagePrompt {
    name: string;
    description: string;
    source: 'local' | 'cdn' | 'example';
    title: string;
    type: 'image' | 'css' | 'javascript' | 'html';
    url: string;
    content: string;
}

type InitialPrompt = Omit<UIPagePrompt, 'shell_page'> & {
    shellType?: 'public' | 'admin' | 'user';
};
type PagePrompt = InitialPrompt & { conversationId: string; id: string };

export class PageQueue extends EventEmitter {
    private readonly queue: PQueue;
    private readonly eventQueue: PQueue;
    private readonly handle: string;
    private readonly systemId: string;
    private readonly systemPrompt: string;
    private readonly references: Map<string, boolean> = new Map();
    private readonly pages: Map<string, string> = new Map();
    private readonly images: Map<string, string> = new Map();
    private uiShells: UIShell[] = [];
    private theme = '';

    constructor(handle: string, systemId: string, systemPrompt: string, concurrency: number = 5) {
        super();
        this.handle = handle;
        this.systemId = systemId;
        this.systemPrompt = systemPrompt;
        this.queue = new PQueue({ concurrency });
        this.eventQueue = new PQueue({ concurrency: Number.MAX_VALUE });
    }

    on(event: 'error', listener: (error: unknown) => void): this;
    on(event: 'event', listener: (data: StormEvent) => void | Promise<void>): this;
    on(event: 'page', listener: (data: StormEventPage) => void | Promise<void>): this;
    on(event: 'image', listener: (data: StormImage, source: ImagePrompt) => void | Promise<void>): this;

    on(event: string, listener: (...args: any[]) => void | Promise<void>): this {
        return super.on(event, (...args) => {
            this.eventQueue
                .add(async () => listener(...args))
                // If the event queue fails, we want to emit an error
                .catch((err) => {
                    this.emit('error', err);
                });
        });
    }

    emit(type: 'error', error: unknown): boolean;
    emit(type: 'event', event: StormEvent): boolean;
    emit(type: 'page', event: StormEventPage): boolean;
    emit(type: 'image', event: StormImage, source: ImagePrompt): boolean;
    emit(eventName: string | symbol, ...args: any[]): boolean {
        return super.emit(eventName, ...args);
    }

    public addUiShell(uiShell: UIShell) {
        this.uiShells.push(uiShell);
    }

    public setUiTheme(theme: string) {
        this.theme = theme;
    }

    private hasPrompt(path: string) {
        if (this.references.has(path)) {
            //console.log('Ignoring duplicate prompt', initialPrompt.path);
            return true;
        }

        if (hasPageOnDisk(this.systemId, path, 'GET')) {
            //console.log('Ignoring prompt with existing page', initialPrompt.path);
            return true;
        }

        return false;
    }

    public addPrompt(
        initialPrompt: InitialPrompt,
        conversationId: string = uuid.v4(),
        overwrite: boolean = false,
        // Set globalEdit to true when the same prompt is being sent to multiple pages
        globalEdit: boolean = false
    ) {
        if (!overwrite && this.hasPrompt(initialPrompt.path)) {
            //console.log('Ignoring duplicate prompt', initialPrompt.path);
            return Promise.resolve();
        }
        console.log('Generating page for', initialPrompt.method, initialPrompt.path);

        const prompt: UIPagePrompt = {
            ...initialPrompt,
            shell_page: this.uiShells.find((shell) => shell.screens.includes(initialPrompt.name))?.content,
            prompt: initialPrompt.prompt,
            theme: this.theme,
            global_edit: globalEdit,
            system_prompt: this.systemPrompt,
            existing_pages: this.getExistingPages(initialPrompt.path),
            existing_images: Object.entries(this.images).map(([path, description]) => ({ path, description })),
        };

        this.references.set(prompt.path, true);
        this.pages.set(prompt.path, prompt.description);

        return this.queue.add<void>(() => this.generate(prompt, conversationId));
    }

    /**
     * Get the existing pages
     */
    private getExistingPages(excludePath?: string) {
        return (
            Object.entries(this.pages)
                // Possibly exclude one page. This is useful when we don't want to include the page
                // we're currently editing.
                .filter(([path]) => (excludePath ? path !== excludePath : true))
                .map(([path, description]) => ({ path, description }))
        );
    }

    private async processPageEventWithReferences(event: StormEventPage) {
        try {
            console.log('Processing page event', event.payload.path);
            const references = await this.resolveReferences(event.payload.content, event.payload.conversationId);
            const matchesExistingPages = (url: string) => {
                return [...this.pages.keys()].some((path) =>
                    new RegExp(`^${path.replaceAll('/*', '/[^/]+')}$`).test(url)
                );
            };
            const initialPrompts: PagePrompt[] = [];
            const resourcePromises = references.map(async (reference) => {
                if (
                    reference.url.startsWith('#') ||
                    reference.url.startsWith('javascript:') ||
                    reference.url.startsWith('http://') ||
                    reference.url.startsWith('https://') ||
                    reference.url.startsWith('data:') ||
                    reference.url.startsWith('blob:') ||
                    reference.url.startsWith('mailto:')
                ) {
                    return;
                }

                switch (reference.type) {
                    case 'image':
                        await this.addImagePrompt(
                            {
                                ...reference,
                                content: event.payload.content,
                            },
                            event.payload.conversationId
                        ).catch((err) => {
                            console.error('Failed to generate image for reference', reference.name, err);
                            this.emit('error', err);
                        });
                        break;
                    case 'css':
                    case 'javascript':
                        //console.log('Ignoring reference', reference);
                        break;
                    case 'html': {
                        const normalizedPath = normalizePath(reference.url);
                        if (matchesExistingPages(normalizedPath)) {
                            break;
                        }
                        this.pages.set(normalizedPath, reference.description);

                        initialPrompts.push({
                            conversationId: randomUUID(),
                            id: randomUUID(),
                            name: reference.name,
                            title: reference.title,
                            path: normalizedPath,
                            method: 'GET',
                            storage_prefix: this.systemId + '_',
                            prompt: `Implement a page for ${reference.name} at ${reference.url} with the following description: ${reference.description}`,
                            description: reference.description,
                            referenced_from: {
                                path: event.payload.path,
                                content: event.payload.content,
                            },
                            existing_pages: this.getExistingPages(),
                            existing_images: Object.entries(this.images).map(([path, description]) => ({
                                path,
                                description,
                            })),
                            // Only used for matching
                            filename: reference.name + '.ref.html',
                            theme: this.theme,
                        });
                        break;
                    }
                }
            });

            // Wait for resources to be generated
            await Promise.allSettled(resourcePromises);
            this.emit('page', event);

            // Emit any new pages after the current page to increase responsiveness
            initialPrompts.forEach((prompt) => {
                if (!this.hasPrompt(prompt.path)) {
                    this.emit('page', {
                        type: 'PAGE',
                        reason: 'reference',
                        created: Date.now(),
                        payload: {
                            id: prompt.id,
                            conversationId: prompt.conversationId,
                            name: prompt.name,
                            title: prompt.title,
                            filename: prompt.filename,
                            method: 'GET',
                            path: prompt.path,
                            prompt: prompt.description,
                            content: '',
                            description: prompt.description,
                        },
                    });
                }
                // Trigger but don't wait for the "bonus" pages
                this.addPrompt(prompt, prompt.conversationId).catch((err) => {
                    console.error('Failed to generate page reference', prompt.name, err);
                    this.emit('error', err);
                });
            });
        } catch (e) {
            console.error('Failed to process event', e);
            throw e;
        }
    }

    public cancel() {
        this.queue.clear();
        this.eventQueue.clear();
    }

    public async wait() {
        while (this.eventQueue.size || this.eventQueue.pending || this.queue.size || this.queue.pending) {
            await this.eventQueue.onIdle();
            await this.queue.onIdle();
        }
    }

    private async addImagePrompt(prompt: ImagePrompt, conversationId: string) {
        if (this.images.has(prompt.url)) {
            //console.log('Ignoring duplicate image prompt', prompt);
            return;
        }
        // Add safeguard to avoid generating images for nonsense URLs
        // Sometimes we get entries for Base URLs that will then cause issues on the filesystem
        // Example: https://www.kapeta.com/images/
        const mimeType = mimetypes.lookup(prompt.url);
        if (!mimeType || !mimeType.startsWith('image/')) {
            console.warn('Skipping image reference of type %s for url %s', mimeType, prompt.url);
            return;
        }

        const client = new StormClient(this.handle, this.systemId);
        this.images.set(prompt.url, prompt.description);
        const result = await client.createImage(
            `Create an image for the url "${prompt.url}" with this description: ${prompt.description}`.trim(),
            conversationId
        );

        let imageEvent: StormImage | null = null;
        result.on('data', (event: StormEvent) => {
            if (event.type === 'IMAGE') {
                imageEvent = event;
            }
        });

        await result.waitForDone();
        if (!imageEvent) {
            throw new Error('No image was generated');
        }

        await writeImageToDisk(this.systemId, imageEvent, prompt);
        this.emit('image', imageEvent, prompt);
    }

    public async generate(prompt: UIPagePrompt, conversationId: string) {
        const client = new StormClient(this.handle, this.systemId);
        const screenStream = await client.createUIPage(prompt, conversationId);
        let pageEvent: StormEventPage | null = null;
        screenStream.on('data', (event: StormEvent) => {
            if (event.type === 'PAGE') {
                event.payload.conversationId = conversationId;
                pageEvent = event;
                return;
            }
            this.emit('event', event);
        });

        await screenStream.waitForDone();
        if (!pageEvent) {
            throw new Error('No page was generated for ' + prompt.name);
        }
        await this.processPageEventWithReferences(pageEvent);
    }

    private async resolveReferences(content: string, conversationId: string): Promise<ReferenceClassification[]> {
        const client = new StormClient(this.handle, this.systemId);
        const referenceStream = await client.classifyUIReferences(content, conversationId);

        const references: ReferenceClassification[] = [];

        referenceStream.on('data', (referenceData: StormEvent) => {
            if (referenceData.type !== 'REF_CLASSIFICATION') {
                return;
            }

            //console.log('Processing reference classification', referenceData);
            references.push(referenceData.payload);
        });

        await referenceStream.waitForDone();

        return references;
    }
}