import { WebSocketMessageResponses } from '../../../agents/constants';
import { BaseController } from '../baseController';
import { generateId } from '../../../utils/idGenerator';
import { AgentState } from '../../../agents/core/state';
import { BehaviorType, ProjectType } from '../../../agents/core/types';
import { getBehaviorTypeForProject } from '../../../agents/core/features';
import { getAgentStub, getTemplateForQuery } from '../../../agents';
import {
    AgentConnectionData,
    AgentPreviewResponse,
    CodeGenArgs,
    MAX_AGENT_QUERY_LENGTH,
} from './types';
import { SecurityError, SecurityErrorType } from 'shared/types/errors';
import { ApiResponse, ControllerResponse } from '../types';
import { RouteContext } from '../../types/route-context';
import { AppService, ModelConfigService } from '../../../database';
import { ModelConfig, credentialsToRuntimeOverrides } from '../../../agents/inferutils/config.types';
import { RateLimitService } from '../../../services/rate-limit/rateLimits';
import { validateWebSocketOrigin } from '../../../middleware/security/websocket';
import { createLogger } from '../../../logger';
import { getPreviewDomain } from 'worker/utils/urls';
import { ImageType, uploadImage } from 'worker/utils/images';
import { ProcessedImageAttachment } from 'worker/types/image-attachment';
import { getTemplateImportantFiles } from 'worker/services/sandbox/utils';
import { hasTicketParam } from '../../../middleware/auth/ticketAuth';
import { checkUsageAndBalance, getUserGateway } from '../../../services/rate-limit/usageChecker';
import { readTokenCookie } from '../../../utils/oauthCookie';
import { UsageLimitExceededError } from 'shared/types/errors';

const defaultCodeGenArgs: Partial<CodeGenArgs> = {
    language: 'typescript',
    frameworks: ['react', 'vite'],
    selectedTemplate: 'auto',
};

const resolveBehaviorType = (body: CodeGenArgs): BehaviorType => {
    if (body.behaviorType) return body.behaviorType;
    // Defer to the feature-definitions registry (DEFAULT_FEATURE_DEFINITIONS)
    // so the single source of truth for "what engine drives this project
    // type" lives alongside the rest of the feature config.
    return getBehaviorTypeForProject(body.projectType ?? 'app');
};

const resolveProjectType = (body: CodeGenArgs): ProjectType | 'auto' => {
    return body.projectType || 'auto';
};


/**
 * CodingAgentController to handle all code generation related endpoints
 */
export class CodingAgentController extends BaseController {
    static logger = createLogger('CodingAgentController');
    /**
     * Start the incremental code generation process
     */
    static async startCodeGeneration(request: Request, env: Env, _: ExecutionContext, context: RouteContext): Promise<Response> {
        try {
            this.logger.info('Starting code generation process');

            const url = new URL(request.url);
            const hostname = url.hostname === 'localhost' ? `localhost:${url.port}`: getPreviewDomain(env);
            // Parse the query from the request body
            let body: CodeGenArgs;
            try {
                body = await request.json() as CodeGenArgs;
            } catch (error) {
                return CodingAgentController.createErrorResponse(`Invalid JSON in request body: ${JSON.stringify(error, null, 2)}`, 400);
            }

            const query = body.query;
            if (typeof query !== 'string' || query.trim().length === 0) {
                return CodingAgentController.createErrorResponse('Missing "query" field in request body', 400);
            }
            if (query.length > MAX_AGENT_QUERY_LENGTH) {
                return CodingAgentController.createErrorResponse(
                    new SecurityError(
                        SecurityErrorType.INVALID_INPUT,
                        `Prompt too large (${query.length} characters). Maximum allowed is ${MAX_AGENT_QUERY_LENGTH} characters.`,
                        413,
                    ),
                    413,
                );
            }
            const { readable, writable } = new TransformStream({
                transform(chunk, controller) {
                    if (chunk === "terminate") {
                        controller.terminate();
                    } else {
                        const encoded = new TextEncoder().encode(JSON.stringify(chunk) + '\n');
                        controller.enqueue(encoded);
                    }
                }
            });
            const writer = writable.getWriter();
            // Check if user is authenticated (required for app creation)
            const user = context.user!;
            try {
                await RateLimitService.enforceAppCreationRateLimit(env, context.config.security.rateLimit, user, request);
            } catch (error) {
                if (error instanceof Error) {
                    return CodingAgentController.createErrorResponse(error, 429);
                } else {
                    this.logger.error('Unknown error in enforceAppCreationRateLimit', error);
                    return CodingAgentController.createErrorResponse(JSON.stringify(error), 429);
                }
            }

            const agentId = generateId();
            const modelConfigService = new ModelConfigService(env);
            const projectType = resolveProjectType(body);
            const behaviorType = resolveBehaviorType(body);

            this.logger.info(`Resolved behaviorType: ${behaviorType}, projectType: ${projectType} for agent ${agentId}`);
                                
            // Fetch all user model configs, api keys and agent instance at once
            const userConfigsRecord = await modelConfigService.getUserModelConfigs(user.id);
                                
            // Extract only user-overridden configs, stripping metadata fields
            const userModelConfigs: Record<string, ModelConfig> = {};
            for (const [actionKey, mergedConfig] of Object.entries(userConfigsRecord)) {
                if (mergedConfig.isUserOverride) {
                    const { isUserOverride, userConfigId, ...modelConfig } = mergedConfig;
                    userModelConfigs[actionKey] = modelConfig;
                }
            }

            const runtimeOverrides = credentialsToRuntimeOverrides(body.credentials);
            // Check usage limits BEFORE making AI calls.
            // The encrypted blob lives in an HttpOnly cookie; read it once here so we
            // can pass it down to the inference pipeline for BYOK decryption.
            const userToken = readTokenCookie(request, env);
            const limitResult = await checkUsageAndBalance(env, user.id, request, userToken);

            if (!limitResult.allowed) {
                this.logger.warn('Request blocked by usage check', {
                    userId: user.id,
                    reason: limitResult.reason,
                    withinLimits: limitResult.withinLimits,
                    hasUserToken: limitResult.hasUserToken,
                    balance: limitResult.balance,
                });
                
                const limit = limitResult.limit;
                const current = Number.isFinite(limit)
                    ? limit - limitResult.remaining
                    : 0;
                const percentUsed = Number.isFinite(limit) && limit > 0
                    ? Math.min(100, (current / limit) * 100)
                    : 0;

                throw new UsageLimitExceededError(
                    limitResult.reason || 'Usage limits exceeded',
                    [{
                        type: 'credits',
                        window: limitResult.windowKind ?? 'rolling',
                        current,
                        max: limit,
                        percentUsed,
                        resetAt: limitResult.resetAt,
                    }],
                    limitResult.hasUserToken
                );
            }

            // Get user's AI Gateway if using BYOK
            let userGateway: { accountId: string; gatewaySlug: string } | null = null;
            if (userToken) {
                userGateway = await getUserGateway(env, user.id);
            }

            // If the limit check transparently refreshed the token, use the fresh blob
            // downstream so inference decrypts to an unexpired access token.
            const effectiveUserToken = limitResult.refreshedBlob ?? userToken;

            const inferenceContext = {
                metadata: {
                    agentId: agentId,
                    userId: user.id,
                    shouldUseUserKey: limitResult.shouldUseByok,
                    userApiToken: effectiveUserToken,
                    userGateway,
                },
                userModelConfigs,
                runtimeOverrides,
                enableRealtimeCodeFix: false, // This costs us too much, so disabled it for now
                enableFastSmartCodeFix: false,
                shouldUseUserKey: limitResult.shouldUseByok, // Use BYOK if needed
                userApiToken: effectiveUserToken, // Encrypted blob from HttpOnly cookie (for BYOK)
                userGateway, // User's AI Gateway for BYOK
            }
                                
            this.logger.info(`Initialized inference context for user ${user.id}`, {
                modelConfigsCount: Object.keys(userModelConfigs).length,
                shouldUseUserKey: limitResult.shouldUseByok,
                balance: limitResult.balance,
                hasUserGateway: !!userGateway,
            });
            this.logger.info(`Creating project of type: ${projectType}`);

            // Think bypasses VibeSDK's template catalog entirely — its
            // companion SpaceDO seeds files via agent-driven writes, and
            // `ThinkCodingBehavior` ignores `templateInfo` regardless.
            // Skip the scratch placeholder + selection synthesis on this path.
            const isThink = behaviorType === 'think';
            const templateResult = isThink
                ? undefined
                : await getTemplateForQuery(env, inferenceContext, query, projectType, body.images, this.logger, body.selectedTemplate);
            const finalProjectType: Exclude<ProjectType, never> = isThink
                ? ((projectType === 'auto' ? 'app' : projectType) as Exclude<ProjectType, never>)
                : templateResult!.projectType;

            const websocketUrl = `${url.protocol === 'https:' ? 'wss:' : 'ws:'}//${url.host}/api/agent/${agentId}/ws`;
            const httpStatusUrl = `${url.origin}/api/agent/${agentId}`;

            let uploadedImages: ProcessedImageAttachment[] = [];
            if (body.images) {
                uploadedImages = await Promise.all(body.images.map(async (image) => {
                    return uploadImage(env, image, ImageType.UPLOADS);
                }));
            }

            writer.write({
                message: 'Code generation started',
                agentId: agentId,
                websocketUrl,
                httpStatusUrl,
                behaviorType,
                projectType: finalProjectType,
                template: isThink
                    ? { name: 'think', files: [] }
                    : {
                        name: templateResult!.templateDetails.name,
                        files: getTemplateImportantFiles(templateResult!.templateDetails),
                    },
            });
            const agentInstance = await getAgentStub(env, agentId, { behaviorType, projectType: finalProjectType });

            const baseInitArgs = {
                query,
                language: body.language || defaultCodeGenArgs.language,
                frameworks: body.frameworks || defaultCodeGenArgs.frameworks,
                hostname,
                inferenceContext,
                images: uploadedImages,
                onBlueprintChunk: (chunk: string) => {
                    writer.write({chunk});
                },
            } as const;

            const initArgs = isThink
                ? baseInitArgs
                : { ...baseInitArgs, templateInfo: { templateDetails: templateResult!.templateDetails, selection: templateResult!.selection } };

            const agentPromise = agentInstance.initialize(initArgs) as Promise<AgentState>;
            void (async () => {
                try {
                    await agentPromise;
                    this.logger.info(`Agent ${agentId} initialized successfully`);
                } catch (error) {
                    const message = error instanceof Error ? error.message : String(error);
                    this.logger.error(`Agent ${agentId} initialization failed`, error);
                    await writer.write({ error: { message } }).catch(() => undefined);
                } finally {
                    await writer.write("terminate").catch(() => undefined);
                    await writer.close().catch(() => undefined);
                }
            })();

            this.logger.info(`Agent ${agentId} init launched successfully`);
            
            const streamHeaders = new Headers({
                // Use SSE content-type to ensure Cloudflare disables buffering,
                // while the payload remains NDJSON lines consumed by the client.
                'Content-Type': 'text/event-stream; charset=utf-8',
                // Prevent intermediary caches/proxies from buffering or transforming
                'Cache-Control': 'no-cache, no-store, must-revalidate, no-transform',
                'Pragma': 'no-cache',
                'Connection': 'keep-alive',
            });
            if (limitResult.refreshedCookie) {
                streamHeaders.append('Set-Cookie', limitResult.refreshedCookie);
            }
            return new Response(readable, {
                status: 200,
                headers: streamHeaders,
            });
        } catch (error) {
            this.logger.error('Error starting code generation', error);
            return CodingAgentController.handleError(error, 'start code generation');
        }
    }

    /**
     * Handle WebSocket connections for code generation
     * This routes the WebSocket connection directly to the Agent
     * 
     * Supports two authentication methods:
     * 1. Ticket-based auth (SDK): ?ticket=tk_xxx in URL
     * 2. JWT-based auth (Browser): Cookie/Header with origin validation
     */
    static async handleWebSocketConnection(
        request: Request,
        env: Env,
        _: ExecutionContext,
        context: RouteContext
    ): Promise<Response> {
        try {
            const agentId = context.pathParams.agentId;
            if (!agentId) {
                return CodingAgentController.createErrorResponse('Missing agent ID parameter', 400);
            }

            // Ensure the request is a WebSocket upgrade request
            if (request.headers.get('Upgrade') !== 'websocket') {
                return new Response('Expected WebSocket upgrade', { status: 426 });
            }

            // User already authenticated via ticket OR JWT by middleware
            const user = context.user;
            if (!user) {
                return CodingAgentController.createErrorResponse('Authentication required', 401);
            }

            // Origin validation only for non-ticket auth (ticket auth is origin-agnostic)
            const isTicketAuth = hasTicketParam(request);
            if (!isTicketAuth && !validateWebSocketOrigin(request, env)) {
                return new Response('Forbidden: Invalid origin', { status: 403 });
            }

            this.logger.info('WebSocket connection authorized', {
                agentId,
                userId: user.id,
                authMethod: isTicketAuth ? 'ticket' : 'jwt',
            });

            try {
                // Get the agent instance to handle the WebSocket connection
                const agentInstance = await getAgentStub(env, agentId);

                // Let the agent handle the WebSocket connection directly
                return agentInstance.fetch(request);
            } catch (error) {
                this.logger.error(`Failed to get agent instance with ID ${agentId}:`, error);
                // Return an appropriate WebSocket error response
                const { 0: client, 1: server } = new WebSocketPair();

                server.accept();
                server.send(JSON.stringify({
                    type: WebSocketMessageResponses.ERROR,
                    error: `Failed to get agent instance: ${error instanceof Error ? error.message : String(error)}`
                }));

                server.close(1011, 'Agent instance not found');

                return new Response(null, {
                    status: 101,
                    webSocket: client
                });
            }
        } catch (error) {
            this.logger.error('Error handling WebSocket connection', error);
            return CodingAgentController.handleError(error, 'handle WebSocket connection');
        }
    }

    /**
     * Connect to an existing agent instance
     * Returns connection information for an already created agent
     */
    static async connectToExistingAgent(
        request: Request,
        env: Env,
        _: ExecutionContext,
        context: RouteContext
    ): Promise<ControllerResponse<ApiResponse<AgentConnectionData>>> {
        try {
            const agentId = context.pathParams.agentId;
            if (!agentId) {
                return CodingAgentController.createErrorResponse<AgentConnectionData>('Missing agent ID parameter', 400);
            }

            this.logger.info(`Connecting to existing agent: ${agentId}`);

            try {
                // Verify the agent instance exists
                const agentInstance = await getAgentStub(env, agentId);
                if (!agentInstance || !(await agentInstance.isInitialized())) {
                    return CodingAgentController.createErrorResponse<AgentConnectionData>('Agent instance not found or not initialized', 404);
                }
                this.logger.info(`Successfully connected to existing agent: ${agentId}`);

                // Construct WebSocket URL
                const url = new URL(request.url);
                const websocketUrl = `${url.protocol === 'https:' ? 'wss:' : 'ws:'}//${url.host}/api/agent/${agentId}/ws`;

                const responseData: AgentConnectionData = {
                    websocketUrl,
                    agentId,
                };

                return CodingAgentController.createSuccessResponse(responseData);
            } catch (error) {
                this.logger.error(`Failed to connect to agent ${agentId}:`, error);
                return CodingAgentController.createErrorResponse<AgentConnectionData>(`Agent instance not found or unavailable: ${error instanceof Error ? error.message : String(error)}`, 404);
            }
        } catch (error) {
            this.logger.error('Error connecting to existing agent', error);
            return CodingAgentController.handleError(error, 'connect to existing agent') as ControllerResponse<ApiResponse<AgentConnectionData>>;
        }
    }

    static async deployPreview(
        _request: Request,
        env: Env,
        _: ExecutionContext,
        context: RouteContext
    ): Promise<ControllerResponse<ApiResponse<AgentPreviewResponse>>> {
        try {
            const agentId = context.pathParams.agentId;
            if (!agentId) {
                return CodingAgentController.createErrorResponse<AgentPreviewResponse>('Missing agent ID parameter', 400);
            }

            const appService = new AppService(env);
            const appResult = await appService.getAppDetails(agentId);

            if (!appResult) {
                return CodingAgentController.createErrorResponse<AgentPreviewResponse>('App not found', 404);
            }

            // Visibility gate (intentional design — not an authorization gap):
            //  - PRIVATE apps: owner-only (any non-owner gets 403 below).
            //  - PUBLIC apps: any authenticated user may spin up an ephemeral
            //    *sandbox* preview. This is deliberate: publishing is decoupled
            //    from production deploy, and viewer-triggered previews are the
            //    intended way to preview public apps that were never deployed.
            //
            // This action is SANDBOX-ONLY: deployToSandbox() never mutates the
            // owner's production deploy state (deploymentId / dispatch worker),
            // and the request passes through the global + per-user auth rate
            // limiters. The credential-theft chain historically associated with
            // this endpoint (reading .dev.vars via the sandbox control plane on
            // port 3000) is closed in worker/services/sandbox/request-handler.ts,
            // which blocks control-plane ports and verifies per-port tokens.
            if(appResult.visibility !== 'public') {
                // If user is logged in and is the owner, allow preview deployment
                const user = context.user;
                if (!user || user.id !== appResult.userId) {
                    return CodingAgentController.createErrorResponse<AgentPreviewResponse>('App is not public. Preview deployment is only available for public apps.', 403);
                }
            }
            this.logger.info(`Deploying preview for agent: ${agentId}`);

            try {
                // Get the agent instance
                const agentInstance = await getAgentStub(env, agentId);
                
                // Deploy the preview
                const preview = await agentInstance.deployToSandbox();
                if (!preview) {
                    return CodingAgentController.createErrorResponse<AgentPreviewResponse>('Failed to deploy preview', 500);
                }
                this.logger.info('Preview deployed successfully', {
                    agentId,
                    previewUrl: preview.previewURL
                });

                return CodingAgentController.createSuccessResponse(preview);
            } catch (error) {
                this.logger.error('Failed to deploy preview', { agentId, error });
                return CodingAgentController.createErrorResponse<AgentPreviewResponse>('Failed to deploy preview', 500);
            }
        } catch (error) {
            this.logger.error('Error deploying preview', error);
            const appError = CodingAgentController.handleError(error, 'deploy preview') as ControllerResponse<ApiResponse<AgentPreviewResponse>>;
            return appError;
        }
    }
}