import NeDBManager from '@/lib/database/nedb.js';

// 生成任务表数据模型
export interface IGenerationTask {
    task_id: string;              // 任务唯一标识，主键
    task_type: 'image' | 'video'; // 任务类型
    server_id: string;            // 分配的服务器ID（外部分配，对应SERVICE_ID）
    
    // 原始请求参数
    original_params: {
        model?: string;           // 模型名称
        prompt: string;           // 提示词
        negative_prompt?: string; // 负向提示词
        width?: number;           // 宽度
        height?: number;          // 高度
        sample_strength?: number; // 采样强度
        images?: Array<{          // 图片参数（视频生成用）
            url: string;
            width: number;
            height: number;
        }>;
        is_pro?: boolean;         // 是否专业版
        duration?: number;        // 时长（毫秒）
        ratio?: string;           // 比例
        response_format?: string; // 响应格式
    };
    
    // 生成过程中的内部参数
    internal_params: {
        refresh_token: string;    // 认证令牌
        component_id?: string;    // 组件ID
        history_id?: string;      // 即梦平台历史记录ID
        mapped_model?: string;    // 映射后的模型名
        submit_id?: string;       // 提交ID
    };
    
    // 任务状态和控制
    status: 'pending' | 'processing' | 'polling' | 'processing_success' | 'uploading' | 'completed' | 'failed';
    retry_count: number;          // 当前重试次数
    max_retries: number;          // 最大重试次数（默认3次）
    upload_retry_count: number;   // TOS上传重试次数
    max_upload_retries: number;   // 最大TOS上传重试次数（默认5次）
    
    // 轮询控制
    next_poll_at?: number;        // 下次轮询时间戳（用于控制轮询间隔）
    poll_interval: number;        // 轮询间隔（秒，默认10秒）
    task_timeout: number;         // 任务超时时间（秒，图片1小时，视频24小时）
    
    // 时间戳
    created_at: number;           // 创建时间戳（秒）
    updated_at: number;           // 更新时间戳（秒）
    started_at?: number;          // 开始处理时间戳（秒）
    completed_at?: number;        // 完成时间戳（秒）
    
    // 错误信息
    error_message?: string;       // 错误描述
    fail_code?: string;           // 即梦平台返回的失败代码
}

// NeDB 数据操作类
export class GenerationTask {
    private static dbName = 'generation_tasks';

    // 创建任务
    static async create(taskData: Omit<IGenerationTask, '_id'>): Promise<IGenerationTask> {
        const currentTime = Math.floor(Date.now() / 1000);
        const task = {
            ...taskData,
            created_at: taskData.created_at || currentTime,
            updated_at: taskData.updated_at || currentTime,
            status: taskData.status || 'pending',
            retry_count: taskData.retry_count || 0,
            max_retries: taskData.max_retries || 3,
            upload_retry_count: taskData.upload_retry_count || 0,
            max_upload_retries: taskData.max_upload_retries || 5,
            poll_interval: taskData.poll_interval || 10
        };

        return await NeDBManager.insert(this.dbName, task);
    }

    // 查找单个任务
    static async findOne(query: any): Promise<IGenerationTask | null> {
        return await NeDBManager.findOne(this.dbName, query);
    }

    // 查找多个任务
    static async find(query: any, sort?: any, limit?: number): Promise<IGenerationTask[]> {
        return await NeDBManager.find(this.dbName, query, sort, limit);
    }

    // 更新任务
    static async updateOne(query: any, update: any, options: any = {}): Promise<number> {
        // 自动更新 updated_at 时间戳
        if (update.$set) {
            update.$set.updated_at = Math.floor(Date.now() / 1000);
        } else {
            update.updated_at = Math.floor(Date.now() / 1000);
        }
        
        return await NeDBManager.update(this.dbName, query, update, { ...options, multi: false });
    }

    // 批量更新任务
    static async updateMany(query: any, update: any, options: any = {}): Promise<number> {
        // 自动更新 updated_at 时间戳
        if (update.$set) {
            update.$set.updated_at = Math.floor(Date.now() / 1000);
        } else {
            update.updated_at = Math.floor(Date.now() / 1000);
        }
        
        return await NeDBManager.update(this.dbName, query, update, { ...options, multi: true });
    }

    // 删除任务
    static async deleteOne(query: any): Promise<number> {
        return await NeDBManager.remove(this.dbName, query, { multi: false });
    }

    // 批量删除任务
    static async deleteMany(query: any): Promise<number> {
        return await NeDBManager.remove(this.dbName, query, { multi: true });
    }

    // 计数
    static async countDocuments(query: any): Promise<number> {
        return await NeDBManager.count(this.dbName, query);
    }

    // 原子更新任务状态
    static async atomicUpdateTaskStatus(
        taskId: string,
        fromStatus: string,
        toStatus: string,
        additionalUpdate: any = {}
    ): Promise<boolean> {
        const updateData = {
            $set: {
                status: toStatus,
                updated_at: Math.floor(Date.now() / 1000),
                ...additionalUpdate
            }
        };

        const numUpdated = await this.updateOne(
            { task_id: taskId, status: fromStatus },
            updateData
        );

        return numUpdated > 0;
    }

    // 获取待处理的任务批次
    static async getBatchPendingTasks(
        serverId: string,
        limit: number = 10
    ): Promise<IGenerationTask[]> {
        const currentTime = Math.floor(Date.now() / 1000);
        
        return await this.find(
            {
                server_id: serverId,
                status: 'pending',
                $or: [
                    { next_poll_at: { $exists: false } },
                    { next_poll_at: { $lte: currentTime } }
                ]
            },
            { created_at: 1 }, // 按创建时间升序
            limit
        );
    }

    // 获取需要轮询的任务
    static async getPollingTasks(
        serverId: string,
        limit: number = 10
    ): Promise<IGenerationTask[]> {
        const currentTime = Math.floor(Date.now() / 1000);
        
        return await this.find(
            {
                server_id: serverId,
                status: 'polling',
                next_poll_at: { $lte: currentTime }
            },
            { next_poll_at: 1 }, // 按轮询时间升序
            limit
        );
    }

    // 获取超时的任务
    static async getTimeoutTasks(): Promise<IGenerationTask[]> {
        const currentTime = Math.floor(Date.now() / 1000);
        
        // NeDB 不支持 $expr，先查询所有活跃任务，然后在内存中过滤超时任务
        const activeTasks = await this.find({
            status: { $in: ['processing', 'polling'] }
        });
        
        return activeTasks.filter(task => {
            const timeoutAt = task.created_at + task.task_timeout;
            return timeoutAt < currentTime;
        });
    }
}

export default GenerationTask;