import { BaseTool } from './BaseTool';
import { Observation } from '../types/types';
import { ParametersType, ToolInfo } from '../types/tool_types';
import { AgentDocumentManager } from '../utils/AgentDocumentManager';
import { gptAPICall, GPTCallType, DescStart, DescEnd } from '../../../../apis/gpt/GPT';
import { ChatSortField } from '../../../collections/CollectionSubView';
import { v4 as uuidv4 } from 'uuid';
import { DocumentView } from '../../DocumentView';
import { docSortings } from '../../../collections/CollectionSubView';
import { Doc } from '../../../../../fields/Doc';

const parameterRules = [
    {
        name: 'sortCriteria',
        type: 'string',
        description: 'Criteria provided by the user to sort the documents.',
        required: true,
    },
] as const;

const toolInfo: ToolInfo<typeof parameterRules> = {
    name: 'sortDocs',
    description: 'Sorts documents within the current Dash environment based on user-specified criteria.',
    parameterRules,
    citationRules: 'No citation needed for sorting operations.',
};

export class SortDocsTool extends BaseTool<typeof parameterRules> {
    private _docManager: AgentDocumentManager;
    private _collectionView: DocumentView;
    private _documentDescriptions: Promise<string> | undefined;

    constructor(docManager: AgentDocumentManager, collectionView: DocumentView) {
        super(toolInfo);
        // Grab the parent collection’s DocumentView (the ChatBox container)
        // We assume the ChatBox itself is currently selected in its parent view.
        this._collectionView = collectionView;
        this._docManager = docManager;
        this._docManager.initializeDocuments();
    }

    get TextToDocMap() {
        // Use any type to avoid complex type checking while maintaining runtime safety
        const childDocs = this._collectionView?.ComponentView?.hasChildDocs?.();
        if (childDocs) {
            const textToDocMap = new Map<string, Doc>();
            try {
                this._documentDescriptions = Promise.all(
                    childDocs.map((doc: Doc) =>
                        Doc.getDescription(doc).then(text => {
                            const cleanText = text.replace(/\n/g, ' ').trim();
                            textToDocMap.set(cleanText, doc);
                            return `${DescStart}${cleanText}${DescEnd}`;
                        })
                    )
                ).then(docDescriptions => docDescriptions.join(''));
                return textToDocMap;
            } catch (error) {
                console.warn('[SortDocsTool] Error initializing document context:', error);
            }
        }
        return undefined;
    }

    async execute(args: ParametersType<typeof parameterRules>): Promise<Observation[]> {
        const chunkId = uuidv4();
        let textToDocMap = await this.TextToDocMap;
        await this._documentDescriptions;
        let chunks: string;

        if (textToDocMap && textToDocMap.size > 0 && this._documentDescriptions) {
            console.log('[SortDocsTool] Using pre-computed document descriptions');
            chunks = await this._documentDescriptions;
        } else {
            // Method 2: Build descriptions from scratch using docManager
            console.log('[SortDocsTool] Building document descriptions from docManager');
            textToDocMap = new Map<string, Doc>();
            const blocks: string[] = [];

            for (const id of this._docManager.docIds) {
                const descRaw = await this._docManager.getDocDescription(id);
                const desc = descRaw.replace(/\n/g, ' ').trim();

                if (!desc) {
                    console.warn(`[SortDocsTool] Skipping document ${id} with empty description`);
                    continue;
                }

                const doc = this._docManager.getDocument(id);
                if (doc) {
                    textToDocMap.set(desc, doc);
                    blocks.push(`${DescStart}${desc}${DescEnd}`);
                }
            }

            chunks = blocks.join('');
        }
        try {
            // 2) call GPT to sort those chunks
            const gptResponse = await gptAPICall(args.sortCriteria, GPTCallType.SORTDOCS, chunks);
            console.log('GPT RESP:', gptResponse);

            // 3) parse & map back to IDs
            const sortedIds = gptResponse
                .split(DescEnd)
                .filter(s => s.trim() !== '')
                .map(s => s.replace(DescStart, '').replace(/\n/g, ' ').trim())
                .map(s => textToDocMap.get(s)) // lookup in our map
                .filter(doc => doc)
                .map(doc => doc!);

            // 4) write back the ordering
            sortedIds.forEach((doc, idx) => {
                doc[ChatSortField] = idx;
            });

            const fieldKey = this._collectionView.ComponentView!.fieldKey;
            this._collectionView.Document[`${fieldKey}_sort`] = docSortings.Chat;

            return [
                {
                    type: 'text',
                    text: `<chunk chunk_id="${chunkId}" chunk_type="sort_status">
Successfully sorted ${sortedIds.length} documents by "${args.sortCriteria}".
</chunk>`,
                },
            ];
        } catch (err) {
            return [
                {
                    type: 'text',
                    text: `<chunk chunk_id="${chunkId}" chunk_type="error">
Sorting failed: ${err instanceof Error ? err.message : err}
</chunk>`,
                },
            ];
        }
    }
}