import { BaseMessage } from "@langchain/core/messages";

export interface SessionItem {
    message: BaseMessage;
    category: string;
}

export interface SessionData {
    items: SessionItem[];
    entityStore: Record<string, any>;
    modelProvider: 'openai' | 'gemini';
}

export class SessionManager {
    private sessions: Map<string, SessionData> = new Map();

    createSession(socketId: string) {
        this.sessions.set(socketId, {
            items: [],
            entityStore: {},
            modelProvider: 'gemini' // Default
        });
        console.log(`Session created for ${socketId}`);
    }

    deleteSession(socketId: string) {
        this.sessions.delete(socketId);
        console.log(`Session deleted for ${socketId}`);
    }

    getSession(socketId: string): SessionData {
        let session = this.sessions.get(socketId);
        if (!session) {
            this.createSession(socketId);
            return this.sessions.get(socketId)!;
        }
        return session;
    }

    getValidHistory(socketId: string): BaseMessage[] {
        const session = this.getSession(socketId);
        return session.items
            .filter(item => item.category !== 'OutOfScope')
            .map(item => item.message);
    }

    updateSession(socketId: string, userMsg: BaseMessage, aiMsg: BaseMessage | undefined, category: string, entityStore: Record<string, any>) {
        const session = this.getSession(socketId);

        // Add User Message
        session.items.push({ message: userMsg, category });

        // Add AI Message if exists
        if (aiMsg) {
            session.items.push({ message: aiMsg, category });
        }

        // Update Entity Store
        session.entityStore = entityStore;

        this.sessions.set(socketId, session);
    }

    setModelProvider(socketId: string, provider: 'openai' | 'gemini') {
        const session = this.getSession(socketId);
        session.modelProvider = provider;
        this.sessions.set(socketId, session); // Ensure map is updated if needed (though object ref works)
        console.log(`Updated model provider for ${socketId} to ${provider}`);
    }

    getModelProvider(socketId: string): 'openai' | 'gemini' {
        return this.getSession(socketId).modelProvider;
    }
}