import { z } from 'zod';
import { AIContextType, AIAgentType, AIContextData } from '../entities/task-ai-context.entity';

/**
 * Validation schema for AI context data based on context type
 */
const PromptContextDataSchema = z.object({
  prompt: z.string().min(1),
  systemPrompt: z.string().optional(),
  parameters: z.record(z.any()).optional(),
  modelUsed: z.string().optional(),
  temperature: z.number().min(0).max(2).optional(),
  maxTokens: z.number().positive().optional(),
  timestamp: z.string().optional(),
  sessionId: z.string().optional(),
  requestId: z.string().optional(),
  metadata: z.record(z.any()).optional(),
});

const ResultContextDataSchema = z.object({
  response: z.string(),
  completionTokens: z.number().nonnegative().optional(),
  promptTokens: z.number().nonnegative().optional(),
  totalTokens: z.number().nonnegative().optional(),
  modelUsed: z.string().optional(),
  timestamp: z.string().optional(),
  sessionId: z.string().optional(),
  requestId: z.string().optional(),
  metadata: z.record(z.any()).optional(),
});

const FeedbackContextDataSchema = z.object({
  feedbackType: z.enum(['positive', 'negative', 'correction']),
  feedbackMessage: z.string(),
  userRating: z.number().min(1).max(5).optional(),
  timestamp: z.string().optional(),
  sessionId: z.string().optional(),
  metadata: z.record(z.any()).optional(),
});

const ErrorContextDataSchema = z.object({
  errorCode: z.string().optional(),
  errorMessage: z.string(),
  stackTrace: z.string().optional(),
  retryCount: z.number().nonnegative().optional(),
  timestamp: z.string().optional(),
  sessionId: z.string().optional(),
  requestId: z.string().optional(),
  metadata: z.record(z.any()).optional(),
});

/**
 * Schema for creating a new task AI context
 */
export const CreateTaskAIContextSchema = z.object({
  task_id: z.string().uuid('Invalid task ID format'),
  ai_agent_type: z.nativeEnum(AIAgentType),
  context_type: z.nativeEnum(AIContextType),
  context_data: z.any(), // Will be validated based on context_type
  execution_timestamp: z.date().optional(),
  success: z.boolean().optional(),
  tokens_used: z.number().nonnegative().nullable().optional(),
  execution_time_ms: z.number().nonnegative().nullable().optional(),
});

/**
 * Schema for updating a task AI context
 */
export const UpdateTaskAIContextSchema = z.object({
  context_data: z.any().optional(),
  success: z.boolean().optional(),
  tokens_used: z.number().nonnegative().nullable().optional(),
  execution_time_ms: z.number().nonnegative().nullable().optional(),
});

/**
 * Validates context data based on context type
 * @param contextType The type of context
 * @param contextData The context data to validate
 * @returns Validated context data
 * @throws ZodError if validation fails
 */
export function validateContextData(
  contextType: AIContextType,
  contextData: any,
): AIContextData {
  switch (contextType) {
    case AIContextType.PROMPT:
      return PromptContextDataSchema.parse(contextData);
    
    case AIContextType.RESULT:
      return ResultContextDataSchema.parse(contextData);
    
    case AIContextType.FEEDBACK:
      return FeedbackContextDataSchema.parse(contextData);
    
    case AIContextType.ERROR:
      return ErrorContextDataSchema.parse(contextData);
    
    default:
      throw new Error(`Unknown context type: ${contextType}`);
  }
}

/**
 * Type for creating a task AI context
 */
export type CreateTaskAIContextDto = z.infer<typeof CreateTaskAIContextSchema>;

/**
 * Type for updating a task AI context
 */
export type UpdateTaskAIContextDto = z.infer<typeof UpdateTaskAIContextSchema>;

/**
 * Validates and prepares data for creating a task AI context
 * @param data Raw input data
 * @returns Validated and prepared data
 */
export function prepareCreateTaskAIContext(data: any): CreateTaskAIContextDto {
  // First validate the basic structure
  const validated = CreateTaskAIContextSchema.parse(data);
  
  // Then validate the context data based on context type
  if (validated.context_data) {
    validated.context_data = validateContextData(
      validated.context_type,
      validated.context_data,
    );
  }
  
  // Set default values
  if (!validated.execution_timestamp) {
    validated.execution_timestamp = new Date();
  }
  
  if (validated.success === undefined) {
    validated.success = true;
  }
  
  return validated;
}

/**
 * Calculates total tokens from context data
 * @param contextData The context data
 * @returns Total token count or null
 */
export function calculateTotalTokens(contextData: AIContextData): number | null {
  if (contextData.totalTokens) {
    return contextData.totalTokens;
  }
  
  if (contextData.promptTokens && contextData.completionTokens) {
    return contextData.promptTokens + contextData.completionTokens;
  }
  
  return null;
}

/**
 * Estimates cost based on token usage and model
 * @param tokens Number of tokens used
 * @param model Model name
 * @returns Estimated cost in USD
 */
export function estimateCost(tokens: number, model: string = 'gpt-4-turbo'): number {
  // Pricing as of 2024 (per 1K tokens)
  const pricing: Record<string, { input: number; output: number }> = {
    'gpt-4-turbo': { input: 0.01, output: 0.03 },
    'gpt-4': { input: 0.03, output: 0.06 },
    'gpt-3.5-turbo': { input: 0.0005, output: 0.0015 },
  };
  
  const modelPricing = pricing[model] || pricing['gpt-4-turbo'];
  // Using average of input and output pricing for estimation
  const avgPrice = (modelPricing.input + modelPricing.output) / 2;
  
  return (tokens / 1000) * avgPrice;
}

/**
 * Validates if an AI agent type is supported
 * @param agentType The agent type to check
 * @returns True if supported
 */
export function isValidAgentType(agentType: string): agentType is AIAgentType {
  return Object.values(AIAgentType).includes(agentType as AIAgentType);
}

/**
 * Validates if a context type is valid
 * @param contextType The context type to check
 * @returns True if valid
 */
export function isValidContextType(contextType: string): contextType is AIContextType {
  return Object.values(AIContextType).includes(contextType as AIContextType);
}