import { FmodeChatCompletion } from 'fmode-ng';
import { FlowTask, FlowTaskOptions } from '../../flow.task';
import { Subject } from 'rxjs';
import { takeUntil } from 'rxjs/operators';

export interface TextCompletionOptions extends FlowTaskOptions {
    promptTemplate: string;          // Required: Your prompt template with {{variables}}
    outputProperty: string;          // Required: Where to store the result in task.data
    inputVariables?: string[];       // Optional: List of required variables (for validation)
    modelOptions?: Record<string, any>; // Optional: Model configuration
    strictValidation?: boolean;      // Optional: Whether to fail on missing variables
}

export class TaskCompletionText extends FlowTask {
    promptTemplate: string;
    outputProperty: string;
    inputVariables: string[];
    strictPromptValidation: boolean = false;
    modelOptions: Record<string, any>;
    strictValidation: boolean;
    destroy$ = new Subject<void>();

    constructor(options: TextCompletionOptions) {
        super({
            title: options.title || 'Text Generation Task',
            output: options.output,
            initialData: options.initialData
        });

        if (!options.promptTemplate) throw new Error('promptTemplate is required');
        if (!options.outputProperty) throw new Error('outputProperty is required');

        this.promptTemplate = options.promptTemplate;
        this.outputProperty = options.outputProperty;
        this.inputVariables = options.inputVariables || this.extractPromptVariables();
        this.modelOptions = options.modelOptions || {};
        this.strictValidation = options.strictValidation ?? true;
    }


    override async handle(): Promise<void> {
        // 1. Validate all required prompt variables exist in task.data
        this.validatePromptVariables();

        // 2. Prepare the prompt with variable substitution
        const fullPrompt = this.renderPromptTemplate();

        // 3. Call the LLM for text completion
        await this.callModelCompletion(fullPrompt);
    }

    validatePromptVariables(): void {
        const requiredVariables = this.extractPromptVariables();
        const missingVariables: string[] = [];
        const undefinedVariables: string[] = [];

        requiredVariables.forEach(variable => {
            if (!(variable in this.data)) {
                missingVariables.push(variable);
            } else if (this.data[variable] === undefined) {
                undefinedVariables.push(variable);
            }
        });

        const errors: string[] = [];

        if (missingVariables.length > 0) {
            errors.push(`Missing required variables in task.data: ${missingVariables.join(', ')}`);
        }

        if (undefinedVariables.length > 0) {
            errors.push(`Variables with undefined values: ${undefinedVariables.join(', ')}`);
        }

        if (errors.length > 0 && this.strictPromptValidation) {
            throw new Error(`Prompt variable validation failed:\n${errors.join('\n')}`);
        } else if (errors.length > 0) {
            console.warn(`Prompt variable warnings:\n${errors.join('\n')}`);
        }
    }

    extractPromptVariables(): string[] {
        const matches = this.promptTemplate.match(/\{\{\w+\}\}/g) || [];
        const uniqueVariables = new Set<string>();

        matches.forEach(match => {
            const key = match.replace(/\{\{|\}\}/g, '');
            uniqueVariables.add(key);
        });

        return Array.from(uniqueVariables);
    }

    renderPromptTemplate(): string {
        let result = this.promptTemplate;
        const variables = this.extractPromptVariables();

        variables.forEach(variable => {
            if (this.data[variable] !== undefined) {
                result = result.replace(new RegExp(`\\{\\{${variable}\\}\\}`, 'g'), this.data[variable]);
            }
        });

        return result;
    }

    async callModelCompletion(prompt: string): Promise<void> {
        return new Promise((resolve, reject) => {
            const messages = [{
                role: "user",
                content: prompt
            }];

            const completion = new FmodeChatCompletion(messages);

            let accumulatedContent = '';

            completion.sendCompletion({
                ...this.modelOptions,
                onComplete: (message: any) => {
                    console.log("onComplete", message);
                }
            })
                .pipe(takeUntil(this.destroy$))
                .subscribe({
                    next: (message: any) => {
                        if (message.content && typeof message.content === 'string') {
                            accumulatedContent = message.content;
                            this.setProgress(0.3 + (accumulatedContent.length / 1000) * 0.7);
                        }

                        if (message.complete) {
                            try {
                                // Store the complete generated text in the specified property
                                this.updateData(this.outputProperty, accumulatedContent);
                                this.setProgress(1);
                                resolve();
                            } catch (error) {
                                this.handleError(error as Error);
                                reject(error);
                            }
                        }
                    },
                    error: (error) => {
                        this.handleError(error);
                        reject(error);
                    }
                });
        });
    }

    handleError(error: Error): void {
        this.updateData('error', {
            message: error.message,
            stack: error.stack,
            timestamp: new Date().toISOString()
        });
        this._status = 'failed';
    }

    onDestroy(): void {
        this.destroy$.next();
        this.destroy$.complete();
    }
}