/*
 * Runtime Web – Unified Serverless Framework for Full-Stack TypeScript Applications
 * Documentation: https://runtimeweb.com/docs/getting-started
 *
 * @template: ai
 */

/**
 * Property-based tests for Chat Handler Orchestration Flow
 *
 * Feature: ai-template, Property 2: Chat Handler Orchestration Flow
 *
 * This test verifies the chat handler orchestration flow:
 * - When a POST request is received at `/chat`, the chat handler retrieves session history from MemoryService
 * - When processing a chat message, the chat handler performs vector search via SearchService
 * - The chat handler invokes BedrockService with the message, history, and relevant context
 * - When a response is generated, the chat handler appends the conversation turn to MemoryService
 */

import { describe, expect, it, jest, beforeEach } from '@jest/globals';
import * as fc from 'fast-check';
import type { APIGatewayProxyEvent, Context } from 'aws-lambda';
import type { ConversationMessage, RelevantDocument, ConversationTurn } from '../../types/chat.types';

/**
 * Mock BedrockService for testing orchestration flow.
 */
class MockBedrockService {
    invokeCalls: Array<{ message: string; history: ConversationMessage[]; relevant: RelevantDocument[] }> = [];
    responseToReturn = 'Mock AI response';
    shouldThrow = false;

    async invoke(params: { message: string; history: ConversationMessage[]; relevant: RelevantDocument[] }): Promise<string> {
        this.invokeCalls.push(params);
        if (this.shouldThrow) {
            throw new Error('Bedrock invocation failed');
        }
        return this.responseToReturn;
    }

    reset(): void {
        this.invokeCalls = [];
        this.responseToReturn = 'Mock AI response';
        this.shouldThrow = false;
    }
}

/**
 * Mock MemoryService for testing orchestration flow.
 */
class MockMemoryService {
    getCalls: string[] = [];
    appendCalls: Array<{ sessionId: string; turn: ConversationTurn }> = [];
    historyToReturn: ConversationMessage[] = [];
    shouldThrow = false;

    async get(sessionId: string): Promise<ConversationMessage[]> {
        this.getCalls.push(sessionId);
        if (this.shouldThrow) {
            throw new Error('Memory service failed');
        }
        return this.historyToReturn;
    }

    async append(sessionId: string, turn: ConversationTurn): Promise<void> {
        this.appendCalls.push({ sessionId, turn });
        if (this.shouldThrow) {
            throw new Error('Memory service failed');
        }
    }

    reset(): void {
        this.getCalls = [];
        this.appendCalls = [];
        this.historyToReturn = [];
        this.shouldThrow = false;
    }
}

/**
 * Mock SearchService for testing orchestration flow.
 */
class MockSearchService {
    knnCalls: Array<{ query: string; k: number }> = [];
    documentsToReturn: RelevantDocument[] = [];
    shouldThrow = false;

    async knn(query: string, k: number = 5): Promise<RelevantDocument[]> {
        this.knnCalls.push({ query, k });
        if (this.shouldThrow) {
            throw new Error('Search service failed');
        }
        return this.documentsToReturn;
    }

    reset(): void {
        this.knnCalls = [];
        this.documentsToReturn = [];
        this.shouldThrow = false;
    }
}

/**
 * Creates a mock API Gateway event for testing.
 */
function createMockEvent(body: Record<string, unknown>): APIGatewayProxyEvent {
    return {
        body: JSON.stringify(body),
        headers: {},
        multiValueHeaders: {},
        httpMethod: 'POST',
        isBase64Encoded: false,
        path: '/chat',
        pathParameters: null,
        queryStringParameters: null,
        multiValueQueryStringParameters: null,
        stageVariables: null,
        requestContext: {} as APIGatewayProxyEvent['requestContext'],
        resource: '/chat',
    };
}

/**
 * Creates a mock Lambda context.
 */
function createMockContext(): Context {
    return {
        callbackWaitsForEmptyEventLoop: false,
        functionName: 'chat-handler',
        functionVersion: '1',
        invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789:function:chat-handler',
        memoryLimitInMB: '128',
        awsRequestId: 'test-request-id',
        logGroupName: '/aws/lambda/chat-handler',
        logStreamName: 'test-log-stream',
        getRemainingTimeInMillis: () => 30000,
        done: () => { },
        fail: () => { },
        succeed: () => { },
    };
}

/**
 * Arbitrary generator for valid session IDs.
 * Session IDs must be non-empty strings.
 */
const sessionIdArb = fc.string({ minLength: 1, maxLength: 100 })
    .filter(s => s.trim().length > 0);

/**
 * Arbitrary generator for valid chat messages.
 * Messages must be non-empty strings.
 */
const messageArb = fc.string({ minLength: 1, maxLength: 1000 })
    .filter(s => s.trim().length > 0);

/**
 * Arbitrary generator for conversation messages.
 */
const conversationMessageArb: fc.Arbitrary<ConversationMessage> = fc.record({
    role: fc.constantFrom('user' as const, 'assistant' as const),
    content: fc.string({ minLength: 1, maxLength: 500 }),
    timestamp: fc.date().map(d => d.toISOString()),
});

/**
 * Arbitrary generator for relevant documents.
 */
const relevantDocumentArb: fc.Arbitrary<RelevantDocument> = fc.record({
    id: fc.uuid(),
    content: fc.string({ minLength: 1, maxLength: 500 }),
    score: fc.float({ min: 0, max: 1, noNaN: true }),
    metadata: fc.option(fc.dictionary(fc.string(), fc.string()), { nil: undefined }),
});

/**
 * Arbitrary generator for conversation history.
 */
const historyArb = fc.array(conversationMessageArb, { minLength: 0, maxLength: 10 });

/**
 * Arbitrary generator for relevant documents array.
 */
const relevantDocsArb = fc.array(relevantDocumentArb, { minLength: 0, maxLength: 5 });

/**
 * Arbitrary generator for AI responses.
 */
const aiResponseArb = fc.string({ minLength: 1, maxLength: 1000 });

describe('Chat Handler Orchestration Flow - Property Tests', () => {
    let mockBedrock: MockBedrockService;
    let mockMemory: MockMemoryService;
    let mockSearch: MockSearchService;

    beforeEach(() => {
        mockBedrock = new MockBedrockService();
        mockMemory = new MockMemoryService();
        mockSearch = new MockSearchService();
    });

    /**
     * Simulates the chat handler orchestration flow.
     * This function replicates the exact flow from chat.handlers.ts for testing.
     */
    async function simulateChatHandler(
        sessionId: string,
        message: string,
        ties: { bedrock: MockBedrockService; memory: MockMemoryService; search: MockSearchService }
    ): Promise<{ statusCode: number; body: Record<string, unknown> }> {
        const { bedrock, memory, search } = ties;

        // Step 1: Get conversation history from DynamoDB
        const history = await memory.get(sessionId);

        // Step 2: Search for relevant context via vector search (RAG)
        const relevant = await search.knn(message, 5);

        // Step 3: Invoke Bedrock AI model with message, history, and context
        const response = await bedrock.invoke({
            message,
            history,
            relevant,
        });

        // Step 4: Append conversation turn to session memory
        await memory.append(sessionId, {
            user: message,
            assistant: response,
        });

        return {
            statusCode: 200,
            body: {
                sessionId,
                response,
                timestamp: new Date().toISOString(),
            },
        };
    }

    /**
     * Feature: ai-template, Property 2: Chat Handler Orchestration Flow
     *
     * WHEN a POST request is received at `/chat`,
     * THE chat handler SHALL retrieve session history from MemoryService.
     */
    describe('Requirement 2.3: Session History Retrieval', () => {
        it('should retrieve session history from MemoryService for any valid sessionId', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    historyArb,
                    async (sessionId, message, history) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockMemory.historyToReturn = history;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: MemoryService.get was called with the sessionId
                        expect(mockMemory.getCalls).toHaveLength(1);
                        expect(mockMemory.getCalls[0]).toBe(sessionId);
                    }
                ),
                { numRuns: 25 }
            );
        });

        it('should pass retrieved history to BedrockService', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    historyArb,
                    async (sessionId, message, history) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockMemory.historyToReturn = history;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: BedrockService received the history from MemoryService
                        expect(mockBedrock.invokeCalls).toHaveLength(1);
                        expect(mockBedrock.invokeCalls[0].history).toEqual(history);
                    }
                ),
                { numRuns: 25 }
            );
        });
    });

    /**
     * Feature: ai-template, Property 2: Chat Handler Orchestration Flow
     *
     * WHEN processing a chat message,
     * THE chat handler SHALL perform vector search via SearchService to find relevant context.
     */
    describe('Requirement 2.4: Vector Search Execution', () => {
        it('should perform vector search with the user message', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    relevantDocsArb,
                    async (sessionId, message, relevantDocs) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockSearch.documentsToReturn = relevantDocs;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: SearchService.knn was called with the message
                        expect(mockSearch.knnCalls).toHaveLength(1);
                        expect(mockSearch.knnCalls[0].query).toBe(message);
                        expect(mockSearch.knnCalls[0].k).toBe(5);
                    }
                ),
                { numRuns: 25 }
            );
        });

        it('should pass search results to BedrockService', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    relevantDocsArb,
                    async (sessionId, message, relevantDocs) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockSearch.documentsToReturn = relevantDocs;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: BedrockService received the relevant documents from SearchService
                        expect(mockBedrock.invokeCalls).toHaveLength(1);
                        expect(mockBedrock.invokeCalls[0].relevant).toEqual(relevantDocs);
                    }
                ),
                { numRuns: 25 }
            );
        });
    });

    /**
     * Feature: ai-template, Property 2: Chat Handler Orchestration Flow
     *
     * THE chat handler SHALL invoke BedrockService with the message, history, and relevant context.
     */
    describe('Requirement 2.5: Bedrock Invocation', () => {
        it('should invoke BedrockService with message, history, and relevant context', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    historyArb,
                    relevantDocsArb,
                    async (sessionId, message, history, relevantDocs) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockMemory.historyToReturn = history;
                        mockSearch.documentsToReturn = relevantDocs;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: BedrockService.invoke was called with correct parameters
                        expect(mockBedrock.invokeCalls).toHaveLength(1);
                        const invokeCall = mockBedrock.invokeCalls[0];
                        expect(invokeCall.message).toBe(message);
                        expect(invokeCall.history).toEqual(history);
                        expect(invokeCall.relevant).toEqual(relevantDocs);
                    }
                ),
                { numRuns: 25 }
            );
        });

        it('should invoke BedrockService exactly once per request', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    async (sessionId, message) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: BedrockService was invoked exactly once
                        expect(mockBedrock.invokeCalls).toHaveLength(1);
                    }
                ),
                { numRuns: 25 }
            );
        });
    });

    /**
     * Feature: ai-template, Property 2: Chat Handler Orchestration Flow
     *
     * WHEN a response is generated,
     * THE chat handler SHALL append the conversation turn to MemoryService.
     */
    describe('Requirement 2.6: Memory Persistence', () => {
        it('should append conversation turn to MemoryService after response generation', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    aiResponseArb,
                    async (sessionId, message, aiResponse) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockBedrock.responseToReturn = aiResponse;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: MemoryService.append was called with correct parameters
                        expect(mockMemory.appendCalls).toHaveLength(1);
                        const appendCall = mockMemory.appendCalls[0];
                        expect(appendCall.sessionId).toBe(sessionId);
                        expect(appendCall.turn.user).toBe(message);
                        expect(appendCall.turn.assistant).toBe(aiResponse);
                    }
                ),
                { numRuns: 25 }
            );
        });

        it('should preserve the exact AI response in the conversation turn', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    aiResponseArb,
                    async (sessionId, message, aiResponse) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();
                        mockBedrock.responseToReturn = aiResponse;

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: The assistant response in the turn matches the Bedrock response exactly
                        expect(mockMemory.appendCalls[0].turn.assistant).toBe(aiResponse);
                    }
                ),
                { numRuns: 25 }
            );
        });
    });

    /**
     * Feature: ai-template, Property 2: Chat Handler Orchestration Flow
     *
     * Validates the complete orchestration flow order:
     * 1. Memory.get (history retrieval)
     * 2. Search.knn (vector search)
     * 3. Bedrock.invoke (AI invocation)
     * 4. Memory.append (persistence)
     */
    describe('Complete Orchestration Flow Order', () => {
        it('should execute operations in the correct order: get → knn → invoke → append', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    historyArb,
                    relevantDocsArb,
                    aiResponseArb,
                    async (sessionId, message, history, relevantDocs, aiResponse) => {
                        // Setup with call order tracking
                        const callOrder: string[] = [];

                        const trackingMemory = {
                            getCalls: [] as string[],
                            appendCalls: [] as Array<{ sessionId: string; turn: ConversationTurn }>,
                            historyToReturn: history,
                            async get(sid: string): Promise<ConversationMessage[]> {
                                callOrder.push('memory.get');
                                this.getCalls.push(sid);
                                return this.historyToReturn;
                            },
                            async append(sid: string, turn: ConversationTurn): Promise<void> {
                                callOrder.push('memory.append');
                                this.appendCalls.push({ sessionId: sid, turn });
                            },
                        };

                        const trackingSearch = {
                            knnCalls: [] as Array<{ query: string; k: number }>,
                            documentsToReturn: relevantDocs,
                            async knn(query: string, k: number = 5): Promise<RelevantDocument[]> {
                                callOrder.push('search.knn');
                                this.knnCalls.push({ query, k });
                                return this.documentsToReturn;
                            },
                        };

                        const trackingBedrock = {
                            invokeCalls: [] as Array<{ message: string; history: ConversationMessage[]; relevant: RelevantDocument[] }>,
                            responseToReturn: aiResponse,
                            async invoke(params: { message: string; history: ConversationMessage[]; relevant: RelevantDocument[] }): Promise<string> {
                                callOrder.push('bedrock.invoke');
                                this.invokeCalls.push(params);
                                return this.responseToReturn;
                            },
                        };

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: trackingBedrock as unknown as MockBedrockService,
                            memory: trackingMemory as unknown as MockMemoryService,
                            search: trackingSearch as unknown as MockSearchService,
                        });

                        // Verify: Operations executed in correct order
                        expect(callOrder).toEqual([
                            'memory.get',
                            'search.knn',
                            'bedrock.invoke',
                            'memory.append',
                        ]);
                    }
                ),
                { numRuns: 25 }
            );
        });

        it('should use the same sessionId for both get and append operations', async () => {
            await fc.assert(
                fc.asyncProperty(
                    sessionIdArb,
                    messageArb,
                    async (sessionId, message) => {
                        // Setup
                        mockMemory.reset();
                        mockSearch.reset();
                        mockBedrock.reset();

                        // Execute
                        await simulateChatHandler(sessionId, message, {
                            bedrock: mockBedrock,
                            memory: mockMemory,
                            search: mockSearch,
                        });

                        // Verify: Same sessionId used for get and append
                        expect(mockMemory.getCalls[0]).toBe(sessionId);
                        expect(mockMemory.appendCalls[0].sessionId).toBe(sessionId);
                    }
                ),
                { numRuns: 25 }
            );
        });
    });
});
