////
// Session storage. 1 browser tab = 1 session

import { AsyncLocalStorage } from "async_hooks"
import { send } from "./websocket"

export type Session = {
    taskId?: string
    sessionId?: string
    project?: string
    cwd?: string
    ws?: any
}

// Ensure "ALS" lives between bun's hot reloads
const g = globalThis as typeof globalThis & { __thread?: AsyncLocalStorage<Session> }
export const ALS = g.__thread ??= new AsyncLocalStorage<Session>()

const sessions: Map<string, Session> = new Map()

export async function sessionRun(sessionId: string, fn: () => any) {
    const state = sessionStore(sessionId)
    return await ALS.run(state, async () => fn())
}

export function sessionGet(key?: keyof Session): Session | any | undefined {
    const store = ALS.getStore()
    if (!store) throw "sessionGet() called outside of ALS.run"

    if (key) return store[key]

    return store
}

export function sessionSet(key: keyof Session, value: any) {
    const store = ALS.getStore()
    if (!store) throw "sessionSet() called outside of ALS.run"
    store[key] = value
    
    if (!store.ws) return
    send(store.ws, {
        type: "session:update",
        data: { [key]: value }
    })
}

export function sessionStore(sessionId: string, taskId?: string, ws?: any): Session {
    let state = sessions.get(sessionId)
    if (!state) {
        state = { sessionId: sessionId, project: "" }
        sessions.set(sessionId, state)
    }
    if (taskId)
        state.taskId = taskId
    if (ws) state.ws = ws
    return state
}