import { infer, InferError, InferResponseString, InferResponseObject, AbortError, CompletionConfig } from './core';
import { createAssistantMessage, createUserMessage, Message } from './common';
import z from 'zod';
// import { CodeEnhancementOutput, CodeEnhancementOutputType } from '../codegen/phasewiseGenerator';
import { SchemaFormat } from './schemaFormatters';
import type { ReasoningEffort } from './config.types';
import { AgentActionKey, AIModels, InferenceContext, ModelConfig } from './config.types';
import { AGENT_CONFIG } from './config';
import { createLogger } from '../../logger';
import { RateLimitExceededError, SecurityError } from 'shared/types/errors';
import { ToolDefinition } from '../tools/types';
import { validateAgentConstraints } from 'worker/api/controllers/modelConfig/constraintHelper';
import { isValidAIModel } from './config.types';

const logger = createLogger('InferenceUtils');

const responseRegenerationPrompts = `
The response you provided was either in an incorrect/unparsable format or was incomplete.
Please provide a valid response that matches the expected output format exactly.
`;

/**
 * Resolves model configuration with field-by-field merge
 *
 * Precedence: userConfig > AGENT_CONFIG defaults
 * Each field is resolved independently (first defined value wins)
 */
function resolveModelConfig(
    agentActionName: AgentActionKey,
    userConfig?: ModelConfig,
): ModelConfig {
    const defaultConfig = AGENT_CONFIG[agentActionName];

    const merged: ModelConfig = {
        name: userConfig?.name ?? defaultConfig.name,
        reasoning_effort: userConfig?.reasoning_effort ?? defaultConfig.reasoning_effort,
        max_tokens: userConfig?.max_tokens ?? defaultConfig.max_tokens,
        temperature: userConfig?.temperature ?? defaultConfig.temperature,
        fallbackModel: userConfig?.fallbackModel ?? defaultConfig.fallbackModel,
    };

    // Validate model name - try userConfig first, then default
    const modelCandidates = [userConfig?.name, defaultConfig.name]
        .filter((n): n is AIModels | string => n !== undefined);

    let validModelName: AIModels | string | undefined;
    for (const candidate of modelCandidates) {
        if (!isValidAIModel(candidate)) {
            logger.warn(`Model ${candidate} not valid, trying next`);
            continue;
        }
        const check = validateAgentConstraints(agentActionName, candidate);
        if (check.constraintEnabled && !check.valid) {
            logger.warn(`Model ${candidate} violates constraints for ${agentActionName}`);
            continue;
        }
        validModelName = candidate;
        break;
    }

    if (!validModelName) {
        logger.warn(`No valid model found for ${agentActionName}, using default`);
        validModelName = defaultConfig.name;
    }
    merged.name = validModelName;

    // Validate fallback model
    if (merged.fallbackModel) {
        const fallbackCheck = validateAgentConstraints(agentActionName, merged.fallbackModel);
        if (fallbackCheck.constraintEnabled && !fallbackCheck.valid) {
            logger.warn(`Fallback ${merged.fallbackModel} violates constraints, using default`);
            merged.fallbackModel = defaultConfig.fallbackModel;
        }
    }

    logger.info(`Resolved config for ${agentActionName}: model=${merged.name}, fallback=${merged.fallbackModel}`);
    return merged;
}

/**
 * Helper function to execute AI inference with consistent error handling
 * @param params Parameters for the inference operation
 * @returns The inference result or null if error
 */

interface InferenceParamsBase {
    env: Env;
    messages: Message[];
    maxTokens?: number;
    temperature?: number;
    modelName?: AIModels | string;
    retryLimit?: number;
    agentActionName: AgentActionKey;
    tools?: ToolDefinition<any, any>[];
    stream?: {
        chunk_size: number;
        onChunk: (chunk: string) => void;
        onReasoning?: (delta: string) => void;
    };
    reasoning_effort?: ReasoningEffort;
    context: InferenceContext;
    onAssistantMessage?: (message: Message) => Promise<void>;
    completionConfig?: CompletionConfig;
    onUsageConsumed?: () => void;
}

interface InferenceParamsStructured<T extends z.ZodObject> extends InferenceParamsBase {
    schema: T;
    format?: SchemaFormat;
}

export async function executeInference<T extends z.ZodObject>(
    params: InferenceParamsStructured<T>
): Promise<InferResponseObject<T>>;

export async function executeInference(
    params: InferenceParamsBase
): Promise<InferResponseString>;
    

export async function executeInference<T extends z.ZodObject>(   {
    env,
    messages,
    temperature,
    maxTokens,
    retryLimit = 5,
    stream,
    tools,
    reasoning_effort,
    schema,
    agentActionName,
    format,
    modelName,
    context,
    onAssistantMessage,
    completionConfig,
}: InferenceParamsBase &    {
    schema?: T;
    format?: SchemaFormat;
}): Promise<InferResponseString | InferResponseObject<T> | null> {
    // Resolve config with clear precedence: userConfig > defaults
    const resolvedConfig = resolveModelConfig(
        agentActionName,
        context?.userModelConfigs?.[agentActionName],
    );

    modelName = modelName || resolvedConfig.name;
    temperature = temperature ?? resolvedConfig.temperature ?? 0.2;
    maxTokens = maxTokens || resolvedConfig.max_tokens || 16000;
    reasoning_effort = reasoning_effort || resolvedConfig.reasoning_effort;

    // Exponential backoff for retries
    const backoffMs = (attempt: number) => Math.min(500 * Math.pow(2, attempt), 10000);

    let useCheaperModel = false;

    for (let attempt = 0; attempt < retryLimit; attempt++) {
        try {
            logger.info(`Starting ${agentActionName} operation with model ${modelName} (attempt ${attempt + 1}/${retryLimit})`);

            const result = schema ? await infer<T>({
                env,
                metadata: context.metadata,
                messages,
                schema,
                schemaName: agentActionName,
                actionKey: agentActionName,
                format,
                maxTokens,
                modelName: useCheaperModel ? AIModels.GEMINI_2_5_FLASH : modelName,
                formatOptions: {
                    debug: false,
                },
                tools,
                stream,
                reasoning_effort: useCheaperModel ? undefined : reasoning_effort,
                temperature,
                abortSignal: context.abortSignal,
                onAssistantMessage,
                completionConfig,
                runtimeOverrides: context.runtimeOverrides,
                onUsageConsumed: context.onUsageConsumed,
                shouldUseUserKey: context.shouldUseUserKey,
                userApiToken: context.userApiToken,
                userGateway: context.userGateway,
            }) : await infer({
                env,
                metadata: context.metadata,
                messages,
                maxTokens,
                modelName: useCheaperModel ? AIModels.GEMINI_2_5_FLASH: modelName,
                tools,
                stream,
                actionKey: agentActionName,
                reasoning_effort: useCheaperModel ? undefined : reasoning_effort,
                temperature,
                abortSignal: context.abortSignal,
                onAssistantMessage,
                completionConfig,
                runtimeOverrides: context.runtimeOverrides,
                onUsageConsumed: context.onUsageConsumed,
                shouldUseUserKey: context.shouldUseUserKey,
                userApiToken: context.userApiToken,
                userGateway: context.userGateway,
            });
            logger.info(`Successfully completed ${agentActionName} operation`);
            // console.log(result);
            return result;
        } catch (error) {
            if (error instanceof RateLimitExceededError || error instanceof SecurityError) {
                throw error;
            }
            
            // Check if cancellation - don't retry, propagate immediately
            if (error instanceof InferError && error.message.includes('cancelled')) {
                logger.info(`${agentActionName} operation cancelled by user, not retrying`);
                throw error;
            }
            
            const isLastAttempt = attempt === retryLimit - 1;
            logger.error(
                `Error during ${agentActionName} operation (attempt ${attempt + 1}/${retryLimit}):`,
                error
            );

            if (error instanceof InferError && !(error instanceof AbortError)) {
                // If its an infer error and not an abort error, we can append the partial response to the list of messages and ask a cheaper model to retry the generation
                if (error.response && error.response.length > 1000) {
                    messages.push(createAssistantMessage(error.response));
                    messages.push(createUserMessage(responseRegenerationPrompts));
                    useCheaperModel = true;
                }
            } else {
                // Switch to fallback model if available
                if (resolvedConfig.fallbackModel && resolvedConfig.fallbackModel !== modelName) {
                    logger.info(`Switching to fallback model: ${resolvedConfig.fallbackModel}`);
                    modelName = resolvedConfig.fallbackModel;
                }
            }

            if (!isLastAttempt) {
                // Wait with exponential backoff before retrying
                const delay = backoffMs(attempt);
                logger.info(`Retrying in ${delay}ms...`);
                await new Promise(resolve => setTimeout(resolve, delay));
            }
        }
    }
    return null;
}

/**
 * Creates a file enhancement request message
 * @param filePath Path to the file being enhanced
 * @param fileContents Contents of the file to enhance
 * @returns A message for the AI model to enhance the file
 */
export function createFileEnhancementRequestMessage(filePath: string, fileContents: string): Message {
    const fileExtension = filePath.split('.').pop() || '';
    const codeBlock = fileExtension ?
        `\`\`\`${fileExtension}\n${fileContents}\n\`\`\`` :
        `\`\`\`\n${fileContents}\n\`\`\``;

    return createUserMessage(`
<FILE_ENHANCEMENT_REQUEST>
Please review the following file and identify any potential issues:
- Syntax errors
- Missing variable declarations
- Incorrect imports
- Incorrect usage of libraries or APIs
- Unicode or special characters that shouldn't be there
- Inconsistent indentation or formatting
- Logic errors
- Any other issues that could cause runtime errors

If you find any issues:
1. Fix them directly in the code
2. Return the full enhanced code with all issues fixed
3. Provide a list of issues that were fixed with clear descriptions

If no issues are found, simply indicate this without modifying the code.

File Path: ${filePath}

${codeBlock}
</FILE_ENHANCEMENT_REQUEST>
`);
}

/**
 * Creates a response message about a generated file
 */
export function createFileGenerationResponseMessage(filePath: string, fileContents: string, explanation: string, nextFile?: { path: string, purpose: string }): Message {
    // Format the message in a focused way to reduce token usage
    const fileExtension = filePath.split('.').pop() || '';
    const codeBlock = fileExtension ?
        `\`\`\`${fileExtension}\n${fileContents}\n\`\`\`` :
        `\`\`\`\n${fileContents}\n\`\`\``;

    return {
        role: 'assistant',
        content: `
<GENERATED FILE: "${filePath}">
${codeBlock}

Explanation: ${explanation}
Next file to generate: ${nextFile ? `Path: ${nextFile.path} | Purpose: (${nextFile.purpose})` : "None"}
`};
}