// src/pages/api/ai/query.ts

import type { APIRoute } from 'astro';
import { withAPIAuth } from '../../../utils/auth.js';
import { apiError, apiServerError, createAuthErrorResponse } from '../../../utils/api.js';
import { enqueueApiCall } from '../../../utils/rateLimitedQueue.js';   
import { pipelineOrchestrator } from '../../../services/ai/pipelineOrchestrator.js';
import { config } from '../../../config/appConfig.js';
import { logger } from '../../../services/logger.js';

export const prerender = false;

interface RateLimitData {
  count: number;
  resetTime: number;
  microTaskCount: number; 
}

const rateLimitStore = new Map<string, RateLimitData>();

function sanitizeInput(input: string): string {
  let sanitized = input
    .replace(/```[\s\S]*?```/g, '[CODE_BLOCK_REMOVED]')
    .replace(/\<\/?[^>]+(>|$)/g, '')
    .replace(/\b(system|assistant|user)\s*[:]/gi, '[ROLE_REMOVED]')
    .replace(/\b(ignore|forget|disregard)\s+(previous|all|your)\s+(instructions?|context|rules?)/gi, '[INSTRUCTION_REMOVED]')
    .trim();
  
  sanitized = sanitized.slice(0, 2000).replace(/\s+/g, ' ');
  return sanitized;
}

function checkRateLimit(userId: string): { allowed: boolean; reason?: string; microTasksRemaining?: number } {
  const now = Date.now();
  const userLimit = rateLimitStore.get(userId);
  
  if (!userLimit || now > userLimit.resetTime) {
    rateLimitStore.set(userId, { 
      count: 1, 
      resetTime: now + config.rateLimit.window,
      microTaskCount: 0 
    });
    return { 
      allowed: true, 
      microTasksRemaining: config.rateLimit.microTaskTotalLimit 
    };
  }
  
  if (userLimit.count >= config.rateLimit.maxRequests) {
    return { 
      allowed: false, 
      reason: `Main rate limit exceeded. Max ${config.rateLimit.maxRequests} requests per minute.`
    };
  }
  
  if (userLimit.microTaskCount >= config.rateLimit.microTaskTotalLimit) {
    return { 
      allowed: false, 
      reason: `Micro-task limit exceeded. Max ${config.rateLimit.microTaskTotalLimit} AI calls per minute.`
    };
  }
  
  userLimit.count++;
  
  return { 
    allowed: true, 
    microTasksRemaining: config.rateLimit.microTaskTotalLimit - userLimit.microTaskCount
  };
}

function incrementMicroTaskCount(userId: string, aiCallsMade: number): void {
  const userLimit = rateLimitStore.get(userId);
  if (userLimit) {
    userLimit.microTaskCount += aiCallsMade;
    logger.debug('rate-limit', `User micro-task usage updated`, {
      userId,
      currentCount: userLimit.microTaskCount,
      limit: config.rateLimit.microTaskTotalLimit
    });
  }
}

function cleanupExpiredRateLimits() {
  const now = Date.now();
  const maxStoreSize = 1000; 
  
  for (const [userId, limit] of rateLimitStore.entries()) {
    if (now > limit.resetTime) {
      rateLimitStore.delete(userId);
    }
  }
  
  if (rateLimitStore.size > maxStoreSize) {
    const entries = Array.from(rateLimitStore.entries());
    entries.sort((a, b) => a[1].resetTime - b[1].resetTime);
    
    const toRemove = entries.slice(0, entries.length - maxStoreSize);
    toRemove.forEach(([userId]) => rateLimitStore.delete(userId));
    
    logger.info('rate-limit', `Cleanup completed`, {
      removed: toRemove.length,
      remaining: rateLimitStore.size
    });
  }
}

setInterval(cleanupExpiredRateLimits, 5 * 60 * 1000);

export const POST: APIRoute = async ({ request }) => {
  try {
    const authResult = await withAPIAuth(request, 'ai');
    if (!authResult.authenticated) {
      return createAuthErrorResponse();
    }
    
    const userId = authResult.userId;

    const rateLimitResult = checkRateLimit(userId);
    if (!rateLimitResult.allowed) {
      return apiError.rateLimit(rateLimitResult.reason || 'Rate limit exceeded');
    }

    const body = await request.json();
    const { query, mode = 'workflow', taskId: clientTaskId } = body;

    if (!query || typeof query !== 'string') {
      logger.api('POST', '/api/ai/query', 400, { error: 'Invalid query', userId });
      return apiError.badRequest('Query required');
    }

    if (!['workflow', 'tool'].includes(mode)) {
      logger.api('POST', '/api/ai/query', 400, { error: 'Invalid mode', userId, mode });
      return apiError.badRequest('Invalid mode. Must be "workflow" or "tool"');
    }

    const sanitizedQuery = sanitizeInput(query);
    if (sanitizedQuery.includes('[FILTERED]')) {
      logger.api('POST', '/api/ai/query', 400, { error: 'Filtered input', userId });
      return apiError.badRequest('Invalid input detected');
    }

    const taskId = clientTaskId || `ai_${userId}_${Date.now()}_${Math.random().toString(36).substr(2, 6)}`;
    
    logger.api('POST', '/api/ai/query', 200, {
      taskId,
      mode,
      queryLength: sanitizedQuery.length,
      userId,
      microTasksRemaining: rateLimitResult.microTasksRemaining
    });
    
    const result = await enqueueApiCall(() => 
      pipelineOrchestrator.processQuery(sanitizedQuery, mode)
    , taskId);

    if (!result || !result.recommendation) {
      return apiServerError.unavailable('No response from AI pipeline');
    }

    const stats = result.processingStats;
    const estimatedAICallsMade = stats.microTasksCompleted + stats.microTasksFailed;
    incrementMicroTaskCount(userId, estimatedAICallsMade);

    logger.api('POST', '/api/ai/query', 200, {
      taskId,
      mode,
      userId,
      queryLength: sanitizedQuery.length,
      processingTime: `${stats.processingTimeMs}ms`,
      microTasksCompleted: stats.microTasksCompleted,
      microTasksFailed: stats.microTasksFailed,
      estimatedAICallsMade,
      embeddingsUsed: stats.embeddingsUsed,
      finalItems: stats.finalSelectedItems
    });

    const currentLimit = rateLimitStore.get(userId);
    const remainingMicroTasks = currentLimit ? 
      config.rateLimit.microTaskTotalLimit - currentLimit.microTaskCount : config.rateLimit.microTaskTotalLimit;

    return new Response(JSON.stringify({
      success: true,
      mode,
      taskId,
      recommendation: result.recommendation,
      query: sanitizedQuery,
      processingStats: {
        ...result.processingStats,
        pipelineType: 'micro-task',
        microTasksSuccessRate: stats.microTasksCompleted / (stats.microTasksCompleted + stats.microTasksFailed),
        averageTaskTime: stats.processingTimeMs / (stats.microTasksCompleted + stats.microTasksFailed),
        estimatedAICallsMade
      },
      rateLimitInfo: {
        mainRequestsRemaining: config.rateLimit.maxRequests - (currentLimit?.count || 0),
        microTaskCallsRemaining: remainingMicroTasks,
        resetTime: Date.now() + config.rateLimit.window
      }
    }), {
      status: 200,
      headers: { 'Content-Type': 'application/json' }
    });

  } catch (error) {
    logger.error('api', 'Pipeline processing failed', error as Error, { 
      endpoint: '/api/ai/query'
    });
    
    if (error.message.includes('embeddings')) {
      return apiServerError.unavailable('Embeddings service error - using AI fallback');
    } else if (error.message.includes('micro-task')) {
      return apiServerError.unavailable('Micro-task pipeline error - some analysis steps failed');
    } else if (error.message.includes('selector')) {
      return apiServerError.unavailable('AI selector service error');
    } else if (error.message.includes('rate limit')) {
      return apiError.rateLimit('AI service rate limits exceeded during micro-task processing');
    } else {
      return apiServerError.internal('Micro-task AI pipeline error');
    }
  }
};