import { Observation } from '../types/types';
import { ParametersType, ToolInfo } from '../types/tool_types';
import { BaseTool } from './BaseTool';

const dataAnalysisToolParams = [
    {
        name: 'csv_file_names',
        type: 'string[]',
        description: 'List of names of the CSV files to analyze',
        required: true,
        max_inputs: 3,
    },
] as const;

type DataAnalysisToolParamsType = typeof dataAnalysisToolParams;

const dataAnalysisToolInfo: ToolInfo<DataAnalysisToolParamsType> = {
    name: 'dataAnalysis',
    description: 'Provides the full CSV file text for your analysis based on the user query and the available CSV file(s).',
    citationRules: 'No citation needed.',
    parameterRules: dataAnalysisToolParams,
};

export class DataAnalysisTool extends BaseTool<DataAnalysisToolParamsType> {
    private csv_files_function: () => { filename: string; id: string; text: string }[];

    constructor(csv_files: () => { filename: string; id: string; text: string }[]) {
        super(dataAnalysisToolInfo);
        this.csv_files_function = csv_files;
    }

    getFileContent(filename: string): string | undefined {
        const files = this.csv_files_function();
        const file = files.find(f => f.filename === filename);
        return file?.text;
    }

    getFileID(filename: string): string | undefined {
        const files = this.csv_files_function();
        const file = files.find(f => f.filename === filename);
        return file?.id;
    }

    async execute(args: ParametersType<DataAnalysisToolParamsType>): Promise<Observation[]> {
        const filenames = args.csv_file_names;
        const results: Observation[] = [];

        for (const filename of filenames) {
            const fileContent = this.getFileContent(filename);
            const fileID = this.getFileID(filename);

            if (fileContent && fileID) {
                results.push({
                    type: 'text',
                    text: `<chunk chunk_id="${fileID}" chunk_type="csv">${fileContent}</chunk>`,
                });
            } else {
                results.push({
                    type: 'text',
                    text: `File not found: ${filename}`,
                });
            }
        }

        return results;
    }
}