import { useDrizzle } from '~~/server/utils/drizzle';
import { generations, messages as messages_drizzle, messagesRelations } from '~~/db/schema';
import { type GenerationStreamEvent, type ChatMessage, type GenerationStatus } from '~~/server/types/chat';
import { eq } from 'drizzle-orm';

/**
 * Streaming generation state - only stores active stream controllers
 * All persistent state lives in the database
 */
interface ActiveGenerationStream {
    userId: string;
    topicId: string;
    clients: Set<ReadableStreamDefaultController<Uint8Array>>;
    isGenerating: boolean;
}

const activeGenerationStreams = new Map<string, ActiveGenerationStream>();
const db = useDrizzle();

/**
 * Add a client connection to an active generation stream
 */
export const addClientToGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): boolean => {
    const stream = activeGenerationStreams.get(generationId);
    if (!stream) {
        return false;
    }
    stream.clients.add(controller);
    return true;
};

/**
 * Remove a client connection from an active generation stream
 */
export const removeClientFromGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): void => {
    const stream = activeGenerationStreams.get(generationId);
    if (stream) {
        stream.clients.delete(controller);
        // Clean up if no clients left and generation is complete
        if (stream.clients.size === 0 && !stream.isGenerating) {
            activeGenerationStreams.delete(generationId);
        }
    }
};

/**
 * Send an event to all connected clients for a generation
 */
export const sendToClients = (generationId: string, event: GenerationStreamEvent): void => {
    const stream = activeGenerationStreams.get(generationId);
    if (!stream) return;

    for (const client of stream.clients) {
        sendToClient(client, event);
    }
};

/**
 * Send an event to a single client
 */
export const sendToClient = (client: ReadableStreamDefaultController<Uint8Array>, event: GenerationStreamEvent): void => {
    const data = JSON.stringify(event);
    const encoder = new TextEncoder();
    try {
        client.enqueue(encoder.encode(`${data}\n`));
    } catch (error) {
        console.error('Failed to send to client:', error);
    }
};

/**
 * Build a prompt from chat messages
 */
const buildPrompt = (messages: ChatMessage[]): string => {
    return messages
        .map((msg: ChatMessage) => {
            const roleMap: Record<typeof msg.type, string> = {
                system: 'System',
                user: 'User',
                agent: 'Assistant'
            };
            return `${roleMap[msg.type]}: ${msg.message}`;
        })
        .join('\n\n');
};

/**
 * Create a new pending generation in the database
 * Returns the generation ID
 */
export const createPendingGeneration = async (userId: string, topicId: string, messages: ChatMessage[], regeneratesFrom?: string): Promise<string> => {
    const generationValues: any = {
        userId,
        topicId,
        status: 'pending' as GenerationStatus,
    };

    if (regeneratesFrom) {
        generationValues.regeneratesFrom = regeneratesFrom;
    }

    const [generation] = await db
        .insert(generations)
        .values(generationValues)
        .returning();

    return generation.id;
};

/**
 * Get current generation status and content from database
 */
export const getGenerationStatus = async (generationId: string) => {
    const [generation] = await db
        .select()
        .from(generations)
        .where(eq(generations.id, generationId));

    return generation || null;
};

/**
 * Start a generation: update status to active and begin streaming
 * This is called when a client connects to the stream
 */
export const startGeneration = async (
    generationId: string,
    userId: string,
    topicId: string,
    messages: ChatMessage[],
    controller: ReadableStreamDefaultController<Uint8Array>
): Promise<void> => {
    try {
        // Get current generation from database
        const generation = await getGenerationStatus(generationId);
        if (!generation) {
            throw new Error('Generation not found');
        }

        // Create active stream tracking
        activeGenerationStreams.set(generationId, {
            userId,
            topicId,
            clients: new Set([controller]),
            isGenerating: true,
        });

        // Update status to active
        await db
            .update(generations)
            .set({
                status: 'active' as GenerationStatus,
                startedAt: new Date(),
            })
            .where(eq(generations.id, generationId));

        sendToClients(generationId, {
            type: 'start',
            data: null,
        });

        const prompt = buildPrompt(messages);
        const dummyResponse = generateDummyResponse(prompt, messages);
        const tokens = dummyResponse.split(' ');

        // Simulate token streaming
        let accumulatedContent = '';
        for (const token of tokens) {
            await new Promise(resolve => setTimeout(resolve, 50));

            accumulatedContent += token + ' ';

            sendToClients(generationId, {
                type: 'token',
                data: token + ' ',
            });
        }

        const finalContent = accumulatedContent.trim();

        // Create message record for this generation
        const messageValues: any = {
            topicId,
            userId,
            content: finalContent,
            isUser: false,
        };

        // If this is a regeneration, set the regeneratedFromId
        if (generation.regeneratesFrom) {
            messageValues.regeneratedFromId = generation.regeneratesFrom;
            messageValues.isRegenerated = true;
        }

        let [message] = await db
            .insert(messages_drizzle)
            .values(messageValues)
            .returning();

        // Update generation as completed
        await db
            .update(generations)
            .set({
                status: 'completed' as GenerationStatus,
                completedAt: new Date(),
                messageId: message.id,
            })
            .where(eq(generations.id, generationId));

        let fmessage = await db.select().from(messages_drizzle).where(eq(messages_drizzle.userId, userId)).leftJoin(generations, eq(messages_drizzle.id, generations.messageId))
        console.log(fmessage);

        // Mark stream as no longer generating
        const stream = activeGenerationStreams.get(generationId);
        if (stream) {
            stream.isGenerating = false;
        }

        sendToClients(generationId, {
            type: 'complete',
            data: fmessage,
        });

        // Close all client connections
        const finalStream = activeGenerationStreams.get(generationId);
        if (finalStream) {
            for (const client of finalStream.clients) {
                try {
                    client.close();
                } catch (error) {
                    // Client already closed
                }
            }
            activeGenerationStreams.delete(generationId);
        }
    } catch (error) {
        console.error('Generation failed:', error);

        const errorMessage = error instanceof Error ? error.message : 'Unknown error';

        // Update generation as failed
        await db
            .update(generations)
            .set({
                status: 'failed' as GenerationStatus,
                error: errorMessage,
                completedAt: new Date(),
            })
            .where(eq(generations.id, generationId));

        sendToClients(generationId, {
            type: 'error',
            data: { error: errorMessage },
        });

        // Close all client connections
        const stream = activeGenerationStreams.get(generationId);
        if (stream) {
            for (const client of stream.clients) {
                try {
                    client.close();
                } catch (error) {
                    // Client already closed
                }
            }
            activeGenerationStreams.delete(generationId);
        }
    }
};

/**
 * Check if a generation is currently being streamed
 */
export const isGenerationStreaming = (generationId: string): boolean => {
    return activeGenerationStreams.has(generationId);
};

/**
 * Generate a dummy response for testing
 */
const generateDummyResponse = (prompt: string, messages: ChatMessage[]): string => {
    const responses = [
        "This is a simulated response to your prompt. In a real implementation, this would be generated by an AI model like GPT-4 or Claude."
        + " I'm processing your message about: " + prompt.substring(0, 50) + "... "
        + "This dummy generation demonstrates the streaming and background save functionality.",

        "I understand your query. This is a placeholder response that simulates AI-generated content."
        + " The system will continue generating this response even if you close the tab, and it will"
        + " automatically save to the database when complete.",

        "Here's a simulated AI response. This demonstrates two key features:"
        + " 1) The generation continues in the background even if you disconnect,"
        + " 2) The complete response is automatically saved to the database without requiring"
        + " a separate update request from the client."
    ];

    const lastUserMessage = messages[messages.length - 1]?.message.toLowerCase() || '';

    if (lastUserMessage.includes('hello') || lastUserMessage.includes('hi')) {
        return "Hello! I'm a dummy AI assistant. This is a simulated response to your greeting."
            + " In production, this would be replaced with actual AI-generated content from an LLM provider.";
    }

    return responses[Math.floor(Math.random() * responses.length)];
};