import { v4 as uuidv4 } from 'uuid';
import { BaseTool } from './BaseTool';
import { Observation } from '../types/types';
import { ParametersType, Parameter, ToolInfo } from '../types/tool_types';
import { DocumentOptions, Docs } from '../../../../documents/Documents';

/**
 * List of supported document types that can be created via text LLM.
 */
type supportedDocumentTypesType = 'text' | 'html' | 'equation' | 'function_plot' | 'dataviz' | 'note_taking' | 'rtf' | 'message' | 'mermaid_diagram' | 'script';
const supportedDocumentTypes: supportedDocumentTypesType[] = ['text', 'html', 'equation', 'function_plot', 'dataviz', 'note_taking', 'rtf', 'message', 'mermaid_diagram', 'script'];

/**
 * Description of document options and data field for each type.
 */
const documentTypesInfo = {
    text: {
        options: ['title', 'backgroundColor', 'fontColor', 'text_align', 'layout'],
        dataDescription: 'The text content of the text document. Should contain all the text content.',
    },
    html: {
        options: ['title', 'backgroundColor', 'layout'],
        dataDescription: 'The HTML-formatted text content of the document.',
    },
    equation: {
        options: ['title', 'backgroundColor', 'fontColor', 'layout'],
        dataDescription: 'The equation content as a string.',
    },
    function_plot: {
        options: ['title', 'backgroundColor', 'layout', 'function_definition'],
        dataDescription: 'The function definition(s) for plotting. Provide as a string or array of function definitions.',
    },
    dataviz: {
        options: ['title', 'backgroundColor', 'layout', 'chartType'],
        dataDescription: 'A string of comma-separated values representing the CSV data.',
    },
    note_taking: {
        options: ['title', 'backgroundColor', 'layout'],
        dataDescription: 'The initial content or structure for note-taking.',
    },
    rtf: {
        options: ['title', 'backgroundColor', 'layout'],
        dataDescription: 'The rich text content in RTF format.',
    },
    message: {
        options: ['title', 'backgroundColor', 'layout'],
        dataDescription: 'The message content of the document.',
    },
    mermaid_diagram: {
        options: ['title', 'backgroundColor', 'layout'],
        dataDescription: 'The Mermaid diagram content.',
    },
    script: {
        options: ['title', 'backgroundColor', 'layout'],
        dataDescription: 'The compilable JavaScript code. Use this for creating scripts.',
    },
};

const createAnyDocumentToolParams = [
    {
        name: 'document_type',
        type: 'string',
        description: `The type of the document to create. Supported types are: ${supportedDocumentTypes.join(', ')}`,
        required: true,
    },
    {
        name: 'data',
        type: 'string',
        description: 'The content or data of the document. The exact format depends on the document type.',
        required: true,
    },
    {
        name: 'options',
        type: 'string',
        description: `A JSON string representing the document options. Available options depend on the document type. For example:
${supportedDocumentTypes
    .map(
        docType => `
- For '${docType}' documents, options include: ${documentTypesInfo[docType].options.join(', ')}`
    )
    .join('\n')}`,
        required: false,
    },
] as const;

type CreateAnyDocumentToolParamsType = typeof createAnyDocumentToolParams;

const createAnyDocToolInfo: ToolInfo<CreateAnyDocumentToolParamsType> = {
    name: 'createAnyDocument',
    description: `Creates any type of document (in Dash) with the provided options and data. Supported document types are: ${supportedDocumentTypes.join(', ')}. dataviz is a csv table tool, so for CSVs, use dataviz. Here are the options for each type:
            <supported_document_types>
        ${supportedDocumentTypes
            .map(
                docType => `
        <document_type name="${docType}">
          <data_description>${documentTypesInfo[docType].dataDescription}</data_description>
          <options>
            ${documentTypesInfo[docType].options.map(option => `<option>${option}</option>`).join('\n')}
          </options>
        </document_type>
        `
            )
            .join('\n')}
      </supported_document_types>`,
    parameterRules: createAnyDocumentToolParams,
    citationRules: 'No citation needed.',
};

export class CreateAnyDocumentTool extends BaseTool<CreateAnyDocumentToolParamsType> {
    private _addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void;

    constructor(addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void) {
        super(createAnyDocToolInfo);
        this._addLinkedDoc = addLinkedDoc;
    }

    async execute(args: ParametersType<CreateAnyDocumentToolParamsType>): Promise<Observation[]> {
        try {
            const documentType: supportedDocumentTypesType = args.document_type.toLowerCase() as supportedDocumentTypesType;
            let options: DocumentOptions = {};

            if (!supportedDocumentTypes.includes(documentType)) {
                throw new Error(`Unsupported document type: ${documentType}. Supported types are: ${supportedDocumentTypes.join(', ')}.`);
            }

            if (!args.data) {
                throw new Error(`Data is required for ${documentType} documents. ${documentTypesInfo[documentType].dataDescription}`);
            }

            if (args.options) {
                try {
                    options = JSON.parse(args.options as string) as DocumentOptions;
                } catch (e) {
                    throw new Error('Options must be a valid JSON string.');
                }
            }

            const data = args.data as string;
            const id = uuidv4();

            // Set default options if not provided
            options.title = options.title || `New ${documentType.charAt(0).toUpperCase() + documentType.slice(1)} Document`;

            // Call the function to add the linked document
            this._addLinkedDoc(documentType, data, options, id);

            return [
                {
                    type: 'text',
                    text: `Created ${documentType} document with ID ${id}.`,
                },
            ];
        } catch (error) {
            return [
                {
                    type: 'text',
                    text: 'Error creating document: ' + (error as Error).message,
                },
            ];
        }
    }
}