import { WebSocket, WebSocketServer } from 'ws';
import { z } from 'zod';
import jwt from 'jsonwebtoken';

import server from "../../ui/_/code/bridge/serverFunctions.js"
import ForumHandler from "./handlers/ForumHandler.js"
import MessagesHandler from "./handlers/MessagesHandler.js"

export default class Socket {
    wss;

    functionCallSchema = z.object({
        id: z.string(),
        name: z.string(),
        args: z.array(z.any()),

        data: z.union([
            z.object({}).passthrough(), // allows any object
            z.array(z.any())            // allows any array
        ]).optional()
    })
    .strict()

    appOperationSchema = z.object({
        id: z.string(),
        app: z.string().optional(),
        operation: z.string().optional(),

        msg: z.union([
            z.object({}).passthrough(), // allows any object
            z.array(z.any())            // allows any array
        ]).optional()
    })
    .superRefine((data, ctx) => {
        if (data.operation !== "GET" && data.msg === undefined) {
            ctx.addIssue({
                code: z.ZodIssueCode.custom,
                path: ["msg"],
                message: "msg is required when operation is not GET"
            })
        }
    })
    .strict()

    constructor(server) {
        this.wss = new WebSocketServer({ server });

        this.wss.on('connection', (ws, req) => {
            console.log('✅ New WebSocket client connected');

            function parseCookies(cookieHeader = "") {
                return Object.fromEntries(
                    cookieHeader.split(";").map(c => {
                        const [key, ...v] = c.trim().split("=");
                        return [key, decodeURIComponent(v.join("="))];
                    })
                );
            }

            const cookies = parseCookies(req.headers.cookie);
            const token = cookies.auth_token;
            if (!token) throw new Error("No auth token");
            let payload;
            try {
                payload = jwt.verify(token, process.env.JWT_SECRET);
            } catch(e) {
                console.error("error: jwt is expired ", e)
            }
            ws.userEmail = payload.email;
            
            ws.on('message', (msg) => {
                this.handleMessage(msg, ws);
            });

            ws.on('close', () => {
                console.log('Client disconnected');
            });
        });

        console.log('WebSocket server initialized');
    }

    // Build a system where the ws obj is updated every time on navigate, so it already has context
    // this way, we can only send broadcast messages to clients that actually have that app / subapp open
    handleMessage = async (msg, ws) => {
        
        try {
            const text = msg.toString();
            const req = JSON.parse(text);
            console.log(req)

            if(this.appOperationSchema.safeParse(req).success) {
                this.handleAppOperation(req, ws)
            } else if(this.functionCallSchema.safeParse(req).success) {
                this.handleFunction(req, ws)
            } else {
                throw new Error("Socket.handleMessage: Incoming ws message has incorrect format!")
            }

        } catch (e) {
            console.error("Invalid WS message:", e);
        }
    }

    async handleAppOperation(req, ws) {
        let responseData;
        switch (req.app) {
            case "FORUM":
                responseData = await ForumHandler.handle(req.operation, req.msg, ws)
                break;

            case "MESSAGES":
                responseData = MessagesHandler.handle(req.operation, req.msg, ws)
                break;

            default:
                console.log("unknown ws message")
        }

        let response = {
            ...req
        }
        response.msg = responseData

        if(!this.appOperationSchema.safeParse(response).success) throw new Error("Socket.handleMessage: Outgoing ws message has incorrect format!")
        ws.send(JSON.stringify(response))
    }

    async handleFunction(req, ws) {
        console.log("func call: ", req.name, req.args)
        let responseData = await server[req.name](...req.args)

        let response = {
            ...req
        }
        response.data = responseData

        console.log(response)

        if(!this.functionCallSchema.safeParse(response).success) throw new Error("Socket.handleMessage: Outgoing ws message has incorrect format!")
        ws.send(JSON.stringify(response))
    }

    broadcast(event) {
        if (!this.wss) return;

        let message = JSON.stringify(event)

        this.wss.clients.forEach(ws => {
                if (ws.readyState === WebSocket.OPEN) {
                ws.send(message);
            }
        });
    }
}
