import { RTFCast } from '../../../../../fields/Types';
import { DocumentOptions } from '../../../../documents/Documents';
import { Networking } from '../../../../Network';
import { ParametersType, ToolInfo } from '../types/tool_types';
import { Observation } from '../types/types';
import { BaseTool } from './BaseTool';
import { Upload } from '../../../../../server/SharedMediaTypes';
import { List } from '../../../../../fields/List';
import { SmartDrawHandler } from '../../../smartdraw/SmartDrawHandler';
import { FireflyImageDimensions } from '../../../smartdraw/FireflyConstants';
import { gptImageCall } from '../../../../apis/gpt/GPT';
import { ClientUtils } from '../../../../../ClientUtils';
import { Doc } from '../../../../../fields/Doc';
import { DocumentViewInternal } from '../../DocumentView';
import { OpenWhere } from '../../OpenWhere';

const imageCreationToolParams = [
    {
        name: 'image_prompt',
        type: 'string',
        description: 'The prompt for the image to be created. This should be a string that describes the image to be created in extreme detail for an AI image generator.',
        required: true,
    },
    {
        name: 'engine',
        type: 'string',
        description: 'The image generation engine to use. Options: "firefly" (default), "dalle". If not specified, defaults to "firefly".',
        required: false,
    },
    {
        name: 'aspect_ratio',
        type: 'string',
        description: 'Aspect ratio for the image (Firefly only). Options: "square" (default), "landscape", "portrait", "widescreen".',
        required: false,
    },
] as const;

type ImageCreationToolParamsType = typeof imageCreationToolParams;

const imageCreationToolInfo: ToolInfo<ImageCreationToolParamsType> = {
    name: 'imageCreationTool',
    citationRules: 'No citation needed. Cannot cite image generation for a response.',
    parameterRules: imageCreationToolParams,
    description: 'Create an image of any style, content, or design, based on a prompt. Uses Firefly by default for better quality and control. Use "dalle" engine explicitly only if requested. The prompt should be a detailed description of the image to be created.',
};

export class ImageCreationTool extends BaseTool<ImageCreationToolParamsType> {


    private _createImage: (result: Upload.FileInformation & Upload.InspectionResults, options: DocumentOptions) => void;
    constructor(createImage: (result: Upload.FileInformation & Upload.InspectionResults, options: DocumentOptions) => void) {
        super(imageCreationToolInfo);
        this._createImage = createImage;
    }

    async execute(args: ParametersType<ImageCreationToolParamsType>): Promise<Observation[]> {
        const { image_prompt, engine = 'firefly', aspect_ratio = 'square' } = args;

        console.log(`Generating image with ${engine} for prompt: ${image_prompt}`);

        try {
            if (engine.toLowerCase() === 'dalle') {
                // Use DALL-E for image generation
                return await this.generateWithDalle(image_prompt);
            } else {
                // Default to Firefly
                return await this.generateWithFirefly(image_prompt, aspect_ratio);
            }
        } catch (error) {
            console.error('ImageCreationTool error:', error);
            return [
                {
                    type: 'text',
                    text: `Error generating image: ${error}`,
                },
            ];
        }
    }

    private async generateWithDalle(prompt: string): Promise<Observation[]> {
        try {
            // Call GPT image API directly
            const imageUrls = await gptImageCall(prompt);
            
            if (imageUrls && imageUrls[0]) {
                // Upload the remote image to our server
                const uploadRes = await Networking.PostToServer('/uploadRemoteImage', { sources: [imageUrls[0]] });
                const fileInfo = (uploadRes as Upload.FileInformation[])[0];
                const source = ClientUtils.prepend(fileInfo.accessPaths.agnostic.client);
                
                // Create image document with DALL-E metadata
                this._createImage(fileInfo as Upload.FileInformation & Upload.InspectionResults, {
                    text: RTFCast(prompt),
                    ai: 'dall-e-3',
                    tags: new List<string>(['@ai']),
                    title: prompt,
                    _width: 400,
                    _height: 400,
                });
                
                return [
                    {
                        type: 'image_url',
                        image_url: { url: source },
                    },
                ];
            } else {
                return [
                    {
                        type: 'text',
                        text: 'Failed to generate image with DALL-E',
                    },
                ];
            }
        } catch (error) {
            console.error('DALL-E generation error:', error);
            throw error;
        }
    }

    private async generateWithFirefly(prompt: string, aspectRatio: string): Promise<Observation[]> {
        try {
            // Map aspect ratio string to FireflyImageDimensions enum
            const dimensionMap: Record<string, FireflyImageDimensions> = {
                'square': FireflyImageDimensions.Square,
                'landscape': FireflyImageDimensions.Landscape,
                'portrait': FireflyImageDimensions.Portrait,
                'widescreen': FireflyImageDimensions.Widescreen,
            };
            
            const dimensions = dimensionMap[aspectRatio.toLowerCase()] || FireflyImageDimensions.Square;
            
            // Use SmartDrawHandler to create Firefly image
            const doc = await SmartDrawHandler.CreateWithFirefly(prompt, dimensions);
            
            if (doc instanceof Doc) {
                // Open the document in a new tab
                DocumentViewInternal.addDocTabFunc(doc, OpenWhere.addRight);
                
                // Get the image URL from the document
                const imageUrl = doc.image || doc.url || '';
                
                return [
                    {
                        type: 'text',
                        text: `Created image with Firefly: "${prompt}". The image has been opened in a new tab.`,
                    },
                ];
            } else {
                return [
                    {
                        type: 'text',
                        text: 'Failed to generate image with Firefly',
                    },
                ];
            }
        } catch (error) {
            console.error('Firefly generation error:', error);
            throw error;
        }
    }
}