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

import { Definition } from '@kapeta/local-cluster-config';
import {
    AIFileTypes,
    BlockCodeGenerator,
    CodeWriter,
    GeneratedFile,
    GeneratedResult,
    MODE_CREATE_ONLY,
    MODE_WRITE_NEVER,
} from '@kapeta/codegen';

import { BlockDefinition } from '@kapeta/schemas';
import { codeGeneratorManager } from '../codeGeneratorManager';
import { STORM_ID, StormClient } from './stormClient';
import {
    StormEvent,
    StormEventBlockStatusType,
    StormEventFileChunk,
    StormEventFileDone,
    StormEventFileLogical,
    StormEventScreen,
} from './events';
import { BlockDefinitionInfo, StormEventParser } from './event-parser';
import { StormFileImplementationPrompt, StormFileInfo, StormStream } from './stream';
import { KapetaURI, parseKapetaUri } from '@kapeta/nodejs-utils';
import { writeFile } from 'fs/promises';
import path from 'path';
import Path, { join } from 'path';
import os from 'node:os';
import YAML from 'yaml';
import { PREDEFINED_BLOCKS } from './predefined';
import { Archetype } from './archetype';
import _ from 'lodash';
import { getSystemBaseImplDir } from './page-utils';

type ImplementationGenerator<T = StormFileImplementationPrompt> = (
    prompt: T,
    conversationId?: string
) => Promise<StormStream>;

const SIMULATED_DELAY = 1000;
const ENABLE_SIMULATED_DELAY = false;
class SimulatedFileDelay {
    private readonly file: StormEventFileDone;
    public readonly stream;
    constructor(file: StormEventFileDone, stream: StormStream) {
        this.file = file;
        this.stream = stream;
    }

    public async start() {
        const commonPayload = {
            filename: this.file.payload.filename,
            path: this.file.payload.path,
            blockName: this.file.payload.blockName,
            blockRef: this.file.payload.blockRef,
            instanceId: this.file.payload.instanceId,
        };
        this.stream.emit('data', {
            type: 'FILE_START',
            created: Date.now(),
            reason: 'File start',
            payload: commonPayload,
        } satisfies StormEventFileLogical);

        const lines = this.file.payload.content.split('\n');
        const delayPerLine = SIMULATED_DELAY / lines.length;
        let lineNumber = 0;
        for (const line of lines) {
            await new Promise<void>((resolve) => {
                setTimeout(() => {
                    this.stream.emit('data', {
                        type: 'FILE_CHUNK',
                        created: Date.now(),
                        reason: 'File chunk',
                        payload: {
                            ...commonPayload,
                            content: line,
                            lineNumber: lineNumber++,
                        },
                    } satisfies StormEventFileChunk);
                    resolve();
                }, delayPerLine);
            });
        }

        this.stream.emit('data', this.file);
    }
}

export class StormCodegen {
    private readonly userPrompt: string;
    private readonly blocks: BlockDefinitionInfo[];
    private readonly out = new StormStream();
    private readonly events: StormEvent[];
    private tmpDir: string;
    private readonly conversationId: string;
    private readonly uiSystemId?: string;

    constructor(
        conversationId: string,
        userPrompt: string,
        blocks: BlockDefinitionInfo[],
        events: StormEvent[],
        uiSystemId?: string
    ) {
        this.userPrompt = userPrompt;
        this.blocks = blocks;
        this.events = events;
        this.tmpDir = Path.join(os.tmpdir(), conversationId);
        this.conversationId = conversationId;
        this.uiSystemId = uiSystemId;
    }

    public setTmpDir(tmpDir: string) {
        this.tmpDir = tmpDir;
    }

    public async process() {
        const promises = this.blocks.map((block) => {
            if (block.archetype) {
                return this.cloneArchetype(block);
            } else {
                return this.processBlockCode(block);
            }
        });
        await Promise.all(promises);
        this.out.end();
    }

    isAborted() {
        return this.out.isAborted();
    }

    public getStream() {
        return this.out;
    }

    private handleTemplateFileOutput(
        blockUri: KapetaURI,
        aiName: string,
        data: StormEvent
    ): StormEventFileDone | undefined {
        if (this.handleError(data)) {
            return;
        }

        if (this.handleFileEvents(blockUri, aiName, data)) {
            return;
        }

        switch (data.type) {
            case 'FILE_DONE':
                return this.handleFileDoneOutput(blockUri, aiName, data);
        }
    }

    private handleError(data: StormEvent): boolean {
        if (data.type === 'ERROR_INTERNAL') {
            this.out.emit('data', data);
            return true;
        }

        return false;
    }

    private handleUiOutput(blockUri: KapetaURI, blockName: string, data: StormEvent): StormEventFileDone | undefined {
        const blockRef = blockUri.toNormalizedString();
        const instanceId = StormEventParser.toInstanceIdFromRef(blockRef);

        if (this.handleError(data)) {
            return;
        }

        if (this.handleFileEvents(blockUri, blockName, data)) {
            return;
        }

        switch (data.type) {
            case 'SCREEN':
                this.out.emit('data', {
                    type: 'SCREEN',
                    reason: data.reason,
                    created: Date.now(),
                    payload: {
                        ...data.payload,
                        blockName,
                        blockRef,
                        instanceId,
                    },
                });
                break;
            case 'FILE_DONE':
                return this.handleFileDoneOutput(blockUri, blockName, data);

            default:
                console.warn('Unknown UI event type', data);
                break;
        }
    }

    private handleFileEvents(blockUri: KapetaURI, blockName: string, data: StormEvent) {
        const blockRef = blockUri.toNormalizedString();
        const instanceId = StormEventParser.toInstanceIdFromRef(blockRef);
        const basePath = this.getBasePath(blockUri.fullName);
        switch (data.type) {
            case 'FILE_START':
            case 'FILE_CHUNK_RESET':
                this.out.emit('data', {
                    ...data,
                    payload: {
                        ...data.payload,
                        path: join(basePath, data.payload.filename),
                        blockName,
                        blockRef,
                        instanceId,
                    },
                });
                return true;
            case 'FILE_CHUNK':
                this.out.emit('data', {
                    ...data,
                    payload: {
                        ...data.payload,
                        path: join(basePath, data.payload.filename),
                        blockName,
                        blockRef,
                        instanceId,
                    },
                });
                return true;
            case 'FILE_STATE':
                this.out.emit('data', {
                    ...data,
                    payload: {
                        ...data.payload,
                        path: join(basePath, data.payload.filename),
                        blockName,
                        blockRef,
                        instanceId,
                    },
                });
                return true;
            case 'FILE_FAILED':
                this.out.emit('data', {
                    ...data,
                    payload: {
                        ...data.payload,
                        path: join(basePath, data.payload.filename),
                        blockName,
                        blockRef,
                        instanceId,
                    },
                });
        }

        return false;
    }

    private handleFileDoneOutput(
        blockUri: KapetaURI,
        aiName: string,
        data: StormEvent
    ): StormEventFileDone | undefined {
        switch (data.type) {
            case 'FILE_DONE':
                const ref = blockUri.toNormalizedString();
                this.emitFile(blockUri, aiName, data.payload.filename, data.payload.content, data.reason);
                return {
                    type: 'FILE_DONE',
                    created: Date.now(),
                    payload: {
                        filename: data.payload.filename,
                        path: join(this.getBasePath(blockUri.fullName), data.payload.filename),
                        content: data.payload.content.replace(/\\n/g, '\n'),
                        blockRef: ref,
                        instanceId: StormEventParser.toInstanceIdFromRef(ref),
                    },
                } as StormEventFileDone;
        }
    }

    private getBasePath(blockName: string) {
        return path.join(this.tmpDir, blockName);
    }

    /**
     * Generates the code for a block and sends it to the AI
     */
    private async processBlockCode(block: BlockDefinitionInfo) {
        if (this.isAborted()) {
            return;
        }

        const kapetaYaml = YAML.stringify(block.content);
        const blockDefinition = block.content;
        const basePath = this.getBasePath(blockDefinition.metadata.name);

        // Generate the code for the block using the standard codegen templates
        const generatedResult = await this.generateBlock(blockDefinition);
        if (!generatedResult) {
            return;
        }

        const kapetaYmlPath = join(basePath, 'kapeta.yml');
        await writeFile(kapetaYmlPath, kapetaYaml);

        const allFiles = this.toStormFiles(generatedResult);

        // Send all the non-ai files to the stream
        await this.emitStaticFiles(parseKapetaUri(block.uri), block.aiName, allFiles);

        if (this.isAborted()) {
            return;
        }

        const blockUri = parseKapetaUri(block.uri);

        const relevantFiles: StormFileInfo[] = allFiles.filter(
            (file) =>
                file.type !== AIFileTypes.IGNORE &&
                file.type !== AIFileTypes.WEB_SCREEN &&
                file.type !== AIFileTypes.WEB_ROUTER
        );
        const uiTemplates = allFiles.filter((file) => file.type === AIFileTypes.WEB_SCREEN);

        const screenFiles: StormEventFileDone[] = [];

        let filteredEvents = [] as StormEvent[];
        for (const event of this.events) {
            filteredEvents.push(event);
            if (event.type === 'PLAN_RETRY') {
                filteredEvents = [];
            }
        }

        const failedEvents = filteredEvents.filter((event) => event.type.endsWith('_FAILED'));
        if (failedEvents.length > 0) {
            console.warn('Codegen encountered failed plan events. Plan might be invalid', failedEvents);
        }
        // TODO: remove this when we have a better way to handle failed events
        // Skip failed events - they are not relevant to the materializer
        filteredEvents = filteredEvents.filter(
            (event) =>
                !event.type.endsWith('_FAILED') &&
                !event.type.endsWith('ERROR_INTERNAL') &&
                !event.type.endsWith('INVALID_RESPONSE')
        );

        const screenEvents: StormEventScreen[] = [];
        const getScreenEventsFile = () => ({
            content: JSON.stringify(screenEvents),
            filename: '<screens>.json',
            type: AIFileTypes.CONFIG,
            mode: MODE_WRITE_NEVER,
            permissions: '0644',
        });
        const uiEvents = [];

        const stormClient = new StormClient(blockUri.handle, this.uiSystemId);
        // generate screens
        if (uiTemplates.length) {
            const screenStream = await stormClient.listScreens({
                events: filteredEvents,
                templates: uiTemplates,
                context: relevantFiles,
                blockName: block.aiName,
                prompt: this.userPrompt,
            });
            screenStream.on('data', (evt) => {
                if (evt.type === 'SCREEN') {
                    screenEvents.push(evt);
                }
                this.handleUiOutput(blockUri, block.aiName, evt);
            });

            this.out.on('aborted', () => {
                screenStream.abort();
            });

            await screenStream.waitForDone();

            // screenfiles
            const screenTemplates = screenEvents
                .map((screenEvent) => ({
                    ...uiTemplates.find((template) => template.filename.endsWith(screenEvent.payload.template)),
                    filename: screenEvent.payload.filename,
                }))
                .filter((tpl): tpl is StormFileInfo => !!tpl.content);

            await Promise.all(
                screenTemplates.map(async (template) => {
                    const payload = {
                        events: filteredEvents,
                        blockName: block.aiName,
                        filename: template.filename,
                        template: template,
                        context: relevantFiles.concat([getScreenEventsFile()]),
                        prompt: this.userPrompt,
                    };

                    const uiStream = await stormClient.createUIImplementation(payload);

                    uiStream.on('data', (evt) => {
                        const uiFile = this.handleUiOutput(blockUri, block.aiName, evt);
                        if (uiFile != undefined) {
                            screenFiles.push(uiFile);
                        }
                        uiEvents.push(evt);
                    });

                    this.out.on('aborted', () => {
                        uiStream.abort();
                    });

                    await uiStream.waitForDone();
                })
            );
        }

        if (this.isAborted()) {
            return;
        }

        const screenFilesConverted = screenFiles.map((screenFile) => {
            return {
                filename: screenFile.payload.filename,
                content: screenFile.payload.content,
                mode: MODE_CREATE_ONLY,
                permissions: '0644',
                type: AIFileTypes.WEB_SCREEN,
            };
        });
        allFiles.push(...screenFilesConverted);

        let webRouters = allFiles.filter((file) => file.type === AIFileTypes.WEB_ROUTER);
        webRouters = await this.processTemplates(
            blockUri,
            block.aiName,
            stormClient.generateCode.bind(stormClient),
            webRouters,
            screenFilesConverted.concat([getScreenEventsFile()])
        );

        // Gather the context files for implementation. These will be all be passed to the AI
        const contextFiles: StormFileInfo[] = relevantFiles.filter(
            (file) => ![AIFileTypes.SERVICE, AIFileTypes.WEB_SCREEN, AIFileTypes.WEB_ROUTER].includes(file.type)
        );

        // Send the service and UI templates to the AI. These will be sent one-by-one in addition to the context files
        let serviceFiles: StormFileInfo[] = allFiles.filter((file) => file.type === AIFileTypes.SERVICE);
        if (serviceFiles.length > 0) {
            serviceFiles = await this.processTemplates(
                blockUri,
                block.aiName,
                stormClient.createServiceImplementation.bind(stormClient),
                serviceFiles,
                contextFiles
            );
        }

        if (this.isAborted()) {
            return;
        }

        for (const contextFile of contextFiles) {
            const filePath = join(basePath, contextFile.filename);
            await writeFile(filePath, contextFile.content);
        }

        for (const screenFile of screenFilesConverted) {
            const filePath = join(basePath, screenFile.filename);
            await writeFile(filePath, screenFile.content);
        }

        // this might decide to overwrite a file that is also a context file
        for (const serviceFile of serviceFiles) {
            const filePath = join(basePath, serviceFile.filename);
            await writeFile(filePath, serviceFile.content);
        }

        for (const webRouterFile of webRouters) {
            const filePath = join(basePath, webRouterFile.filename);
            await writeFile(filePath, webRouterFile.content);
        }

        const blockRef = block.uri;

        this.emitBlockStatus(blockUri, block.aiName, StormEventBlockStatusType.QA);

        /* TODO: temporarily disabled - enable again when codegen is more stable
                const filesToBeFixed = serviceFiles.concat(contextFiles).concat(screenFilesConverted);
                const codeGenerator = new BlockCodeGenerator(blockDefinition);
*/

        this.emitBlockStatus(blockUri, block.aiName, StormEventBlockStatusType.BUILDING);
        //        await this.verifyAndFixCode(blockUri, block.aiName, codeGenerator, basePath, filesToBeFixed, allFiles);
        this.emitBlockStatusDone(basePath, block, blockRef);
    }

    private emitBlockStatusDone(basePath: string, block: BlockDefinitionInfo, blockRef: string) {
        this.out.emit('data', {
            type: 'BLOCK_READY',
            reason: 'Block ready',
            created: Date.now(),
            payload: {
                path: basePath,
                blockName: block.aiName,
                blockRef,
                instanceId: StormEventParser.toInstanceIdFromRef(blockRef),
            },
        } satisfies StormEvent);
    }

    private emitBlockStatus(blockUri: KapetaURI, blockName: string, status: StormEventBlockStatusType) {
        this.out.emit('data', {
            type: 'BLOCK_STATUS',
            reason: status,
            created: Date.now(),
            payload: {
                status,
                blockName,
                blockRef: blockUri.toNormalizedString(),
                instanceId: StormEventParser.toInstanceIdFromRef(blockUri.toNormalizedString()),
            },
        } satisfies StormEvent);
    }

    removePrefix(prefix: string, str: string): string {
        if (str.startsWith(prefix)) {
            return str.slice(prefix.length);
        }
        return str;
    }

    /**
     * Emits the text-based files to the stream
     */
    private async emitStaticFiles(uri: KapetaURI, aiName: string, files: StormFileInfo[]) {
        const promises = files.map((file) => {
            if (!file.content || typeof file.content !== 'string') {
                return;
            }

            if (file.type === AIFileTypes.SERVICE) {
                // Don't send the service files to the stream yet
                // They will need to be implemented by the AI
                return;
            }

            if ([AIFileTypes.WEB_ROUTER, AIFileTypes.WEB_SCREEN].includes(file.type)) {
                // Don't send the web screen files to the stream yet
                // They will need to be implemented by the AI
                return;
            }

            const basePath = this.getBasePath(uri.fullName);
            const ref = uri.toNormalizedString();
            const fileEvent: StormEventFileDone = {
                type: 'FILE_DONE',
                reason: 'File generated',
                created: Date.now(),
                payload: {
                    filename: file.filename,
                    path: join(basePath, file.filename),
                    content: file.content,
                    blockName: aiName,
                    blockRef: ref,
                    instanceId: StormEventParser.toInstanceIdFromRef(ref),
                },
            };

            if (ENABLE_SIMULATED_DELAY) {
                // Simulate a delay when sending the file
                return new SimulatedFileDelay(fileEvent, this.out).start();
            } else {
                this.out.emit('data', fileEvent);
            }
        });

        return Promise.all(promises);
    }

    private emitFile(
        uri: KapetaURI,
        blockName: string,
        filename: string,
        content: string,
        reason: string = 'File generated'
    ) {
        const basePath = this.getBasePath(uri.fullName);
        const ref = uri.toNormalizedString();
        this.out.emit('data', {
            type: 'FILE_DONE',
            reason,
            created: Date.now(),
            payload: {
                filename: filename,
                path: join(basePath, filename),
                content: content,
                blockName,
                blockRef: ref,
                instanceId: StormEventParser.toInstanceIdFromRef(ref),
            },
        } satisfies StormEventFileDone);
    }

    /**
     * Sends the template to the AI and processes the response
     */
    private async processTemplates(
        blockUri: KapetaURI,
        aiName: string,
        generator: ImplementationGenerator,
        templates: StormFileInfo[],
        contextFiles: StormFileInfo[]
    ): Promise<StormFileInfo[]> {
        const promises = templates.map(async (templateFile) => {
            let changedFiles: StormEventFileDone[] = [];
            const stream = await generator({
                context: contextFiles,
                template: templateFile,
                prompt: this.userPrompt,
            });

            this.out.on('aborted', () => {
                stream.abort();
            });

            stream.on('data', (evt) => {
                let changedFile = this.handleTemplateFileOutput(blockUri, aiName, evt);
                if (changedFile) {
                    changedFiles.push(...changedFiles, changedFile);
                }
            });

            await stream.waitForDone();
            return changedFiles;
        });

        const changedFiles: StormEventFileDone[][] = await Promise.all(promises);
        let allFiles = templates.concat(contextFiles);
        let result: StormFileInfo[] = [];
        for (const changedFile of changedFiles.flat()) {
            const find = allFiles.find((file) => file.filename === changedFile.payload.filename);
            if (find) {
                result.push({
                    type: find.type,
                    filename: find.filename,
                    mode: find.mode,
                    permissions: find.permissions,
                    content: changedFile.payload.content,
                });
            } else {
                console.warn(
                    "processTemplates: AI changed a file that wasn't in the input [" +
                        changedFile.payload.filename +
                        ']'
                );
            }
        }
        return result;
    }

    /**
     * Converts the generated files to a format that can be sent to the AI
     */
    private toStormFiles(generatedResult: GeneratedResult) {
        const allFiles: StormFileInfo[] = generatedResult.files.map((file: GeneratedFile) => {
            if (!file.content) {
                return {
                    ...file,
                    type: AIFileTypes.IGNORE,
                };
            }
            if (typeof file.content !== 'string') {
                return {
                    ...file,
                    type: AIFileTypes.IGNORE,
                };
            }

            const rx = /\/\/AI-TYPE:([a-z0-9- _]+)\n/gi;

            const match = rx.exec(file.content);
            if (!match) {
                return {
                    ...file,
                    type: AIFileTypes.IGNORE,
                };
            }
            const type = match[1].trim() as AIFileTypes;
            file.content = file.content.replace(rx, '');

            return {
                ...file,
                type,
            };
        });

        return allFiles;
    }

    /**
     * Generates the code using codegen for a given block.
     */
    private async generateBlock(yamlContent: Definition): Promise<GeneratedResult | undefined> {
        if (this.isAborted()) {
            return;
        }
        if (!yamlContent.spec.target?.kind) {
            //Not all block types have targets
            return;
        }

        if (!(await codeGeneratorManager.ensureTarget(yamlContent.spec.target?.kind))) {
            return;
        }
        const basePath = this.getBasePath(yamlContent.metadata.name);

        const codeGenerator = new BlockCodeGenerator(yamlContent as BlockDefinition);
        codeGenerator.withOption('AIContext', STORM_ID);

        if (this.uiSystemId) {
            codeGenerator.withOption('AIStaticFiles', getSystemBaseImplDir(this.uiSystemId));
        }

        const generatedResult = await codeGenerator.generate();
        new CodeWriter(basePath).write(generatedResult);
        return generatedResult;
    }

    private async cloneArchetype(block: BlockDefinitionInfo): Promise<void> {
        const predefinedBlock = PREDEFINED_BLOCKS.get(block.archetype!)!;

        let blockDefinition = await predefinedBlock.getBlockDefinition();
        const kapetaURI = new KapetaURI(block.uri);
        _.set(blockDefinition!, ['metadata', 'name'], kapetaURI.fullName);
        _.set(blockDefinition!, ['metadata', 'title'], kapetaURI.name);

        const basePath = this.getBasePath(blockDefinition.metadata.name);

        let archetype = new Archetype();

        const generatedResult = await archetype.cloneRepository(predefinedBlock, this.tmpDir, basePath);

        await this.emitStaticFiles(kapetaURI, block.aiName, this.toStormFiles(generatedResult));

        const blockRef = block.uri;
        this.emitBlockStatusDone(basePath, block, blockRef);
    }

    abort() {
        this.out.abort();
    }
}