import * as net from "node:net";
import * as vscode from "vscode";

import {
    CloseAction,
    ErrorAction,
    LanguageClient,
    LanguageClientOptions,
    ServerOptions,
    State,
    StreamInfo
} from "vscode-languageclient/node";

let client: LanguageClient | undefined;
let output: vscode.OutputChannel;
let reconnectTimer: NodeJS.Timeout | undefined;
let reconnectContext: vscode.ExtensionContext | undefined;
let stoppingClient = false;
let lastAppliedThemeSignature: string | undefined;

const RECONNECT_DELAY_MS = 2000;

type PrometeuTokenStyle = {
    semanticKey: string;
    foreground: string;
    italic: boolean;
    bold: boolean;
    underline: boolean;
};

type PrometeuEditorPalette = {
    baseForeground: string;
    selectionBackground: string;
    activeHighlightBackground: string;
    lineNumberForeground: string;
    statusChipBackground: string;
    statusChipBorder: string;
    statusChipForeground: string;
};

type PrometeuVisualTheme = {
    themeId: string;
    displayName: string;
    editorPalette: PrometeuEditorPalette;
    tokenStyles: PrometeuTokenStyle[];
};

type PrometeuVisualThemesPayload = {
    frontendLanguageId: string;
    activeVisualThemeId: string;
    visualThemes: PrometeuVisualTheme[];
};

type PrometeuSemanticHostProjectionEntry = {
    semanticKey: string;
    hostTokenType: string;
    hostTokenModifiers: string[];
    fallbackTokenType: string;
    fallbackTokenModifiers: string[];
};

type PrometeuSemanticHostProjection = {
    hostId: string;
    tokenProjections: PrometeuSemanticHostProjectionEntry[];
};

type PrometeuSemanticHostProjectionsPayload = {
    frontendLanguageId: string;
    semanticKeys: string[];
    hostProjections: PrometeuSemanticHostProjection[];
};

const VSCODE_HOST_ID = "vscode";

export function activate(context: vscode.ExtensionContext): void {
    reconnectContext = context;
    output = vscode.window.createOutputChannel("Prometeu LSP");
    output.appendLine("Prometeu VS Code extension activated.");

    const connectCommand = vscode.commands.registerCommand(
        "prometeu.connectStudioLsp",
        async () => {
            await startClient(context);
        }
    );

    const restartCommand = vscode.commands.registerCommand(
        "prometeu.restartStudioLsp",
        async () => {
            await stopClient();
            await startClient(context);
        }
    );

    context.subscriptions.push(output, connectCommand, restartCommand);

    void startClient(context);
}

export async function deactivate(): Promise<void> {
    await stopClient();
}

async function startClient(context: vscode.ExtensionContext): Promise<void> {
    clearReconnectTimer();

    if (client) {
        output.appendLine("Prometeu LSP client is already running.");
        return;
    }

    const config = vscode.workspace.getConfiguration("prometeu.studio");

    const host = config.get<string>("host", "127.0.0.1");
    const port = config.get<number>("port", 7775);

    output.show(true);
    output.appendLine(`Connecting to Prometeu Studio LSP at ${host}:${port}...`);

    const serverOptions: ServerOptions = async (): Promise<StreamInfo> => {
        return new Promise<StreamInfo>((resolve, reject) => {
            const socket = net.createConnection({ host, port });

            let settled = false;

            const connectTimeout = setTimeout(() => {
                if (settled) {
                    return;
                }

                settled = true;
                socket.destroy();
                reject(new Error("Connection timed out."));
            }, 3000);

            socket.once("connect", () => {
                if (settled) {
                    return;
                }

                settled = true;
                clearTimeout(connectTimeout);

                output.appendLine("TCP connection established with Prometeu Studio.");

                resolve({
                    reader: socket,
                    writer: socket
                });
            });

            socket.once("error", (error) => {
                if (settled) {
                    return;
                }

                settled = true;
                clearTimeout(connectTimeout);

                reject(error);
            });
        });
    };

    const clientOptions: LanguageClientOptions = {
        documentSelector: [
            {
                scheme: "file",
                language: "pbs"
            }
        ],
        synchronize: {
            fileEvents: [
                vscode.workspace.createFileSystemWatcher("**/*.pbs")
            ]
        },
        initializationOptions: {
            clientKind: "vscode",
            prometeuClientVersion: context.extension.packageJSON.version
        },
        outputChannel: output,
        errorHandler: {
            error: (error) => {
                output.appendLine(`Prometeu LSP transport error: ${error.message}`);
                return { action: ErrorAction.Continue };
            },
            closed: () => {
                output.appendLine("Prometeu LSP transport closed.");
                return { action: CloseAction.DoNotRestart };
            }
        }
    };

    const nextClient = new LanguageClient(
        "prometeuStudioLsp",
        "Prometeu Studio LSP",
        serverOptions,
        clientOptions
    );
    client = nextClient;
    nextClient.onDidChangeState((event) => {
        output.appendLine(`Prometeu LSP state changed: ${State[event.oldState]} -> ${State[event.newState]}`);
        if (event.newState !== State.Stopped) {
            return;
        }
        if (client === nextClient) {
            client = undefined;
        }
        if (!stoppingClient) {
            scheduleReconnect();
        }
    });

    try {
        await nextClient.start();
        await applyPrometeuVisualTheme(nextClient.initializeResult);

        output.appendLine("Prometeu Studio LSP client started.");
        vscode.window.showInformationMessage("Connected to Prometeu Studio LSP.");
    } catch (error) {
        if (client === nextClient) {
            client = undefined;
        }

        const message = error instanceof Error ? error.message : String(error);

        output.appendLine(`Failed to start Prometeu LSP client: ${message}`);
        vscode.window.showErrorMessage(`Could not connect to Prometeu Studio LSP: ${message}`);
        scheduleReconnect();
    }
}

async function stopClient(): Promise<void> {
    clearReconnectTimer();
    if (!client) {
        return;
    }

    const currentClient = client;
    client = undefined;
    stoppingClient = true;

    output.appendLine("Stopping Prometeu Studio LSP client...");

    try {
        await currentClient.stop();
        output.appendLine("Prometeu Studio LSP client stopped.");
    } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        output.appendLine(`Failed to stop Prometeu LSP client cleanly: ${message}`);
    } finally {
        stoppingClient = false;
    }
}

function scheduleReconnect(): void {
    if (!reconnectContext) {
        return;
    }
    if (reconnectTimer || client) {
        return;
    }
    output.appendLine(`Scheduling Prometeu LSP reconnect in ${RECONNECT_DELAY_MS}ms...`);
    reconnectTimer = setTimeout(() => {
        reconnectTimer = undefined;
        void startClient(reconnectContext!);
    }, RECONNECT_DELAY_MS);
}

function clearReconnectTimer(): void {
    if (!reconnectTimer) {
        return;
    }
    clearTimeout(reconnectTimer);
    reconnectTimer = undefined;
}

async function applyPrometeuVisualTheme(initializeResult: unknown): Promise<void> {
    const visualThemesPayload = extractPrometeuVisualThemesPayload(initializeResult);
    if (!visualThemesPayload) {
        output.appendLine("Prometeu visual theme payload not present in initialize result.");
        return;
    }
    const semanticProjectionPayload = extractPrometeuSemanticHostProjectionsPayload(initializeResult);
    const vscodeProjection = selectHostProjection(semanticProjectionPayload, VSCODE_HOST_ID);
    if (!semanticProjectionPayload) {
        output.appendLine("Prometeu semantic host projection payload not present in initialize result.");
    } else if (!vscodeProjection) {
        output.appendLine("Prometeu semantic host projection payload does not declare a vscode projection.");
    }

    const activeTheme = visualThemesPayload.visualThemes.find((theme) => theme.themeId === visualThemesPayload.activeVisualThemeId);
    if (!activeTheme) {
        output.appendLine(`Prometeu visual theme '${visualThemesPayload.activeVisualThemeId}' was not declared by the server.`);
        return;
    }

    const themeSignature = JSON.stringify({
        languageId: visualThemesPayload.frontendLanguageId,
        themeId: activeTheme.themeId,
        editorPalette: activeTheme.editorPalette,
        tokenStyles: activeTheme.tokenStyles,
        semanticProjection: vscodeProjection ?? null
    });
    if (themeSignature === lastAppliedThemeSignature) {
        output.appendLine(`Prometeu visual theme '${activeTheme.themeId}' is already applied.`);
        return;
    }

    const configurationTarget = vscode.workspace.workspaceFolders?.length
        ? vscode.ConfigurationTarget.Workspace
        : vscode.ConfigurationTarget.Global;
    const configuration = vscode.workspace.getConfiguration();

    const existingSemanticCustomizations = configuration.get<Record<string, unknown>>("editor.semanticTokenColorCustomizations") ?? {};
    const semanticRules = {
        ...(readObject(existingSemanticCustomizations.rules)),
        ...buildSemanticTokenRules(visualThemesPayload.frontendLanguageId, activeTheme, vscodeProjection)
    };

    await configuration.update(
        "editor.semanticTokenColorCustomizations",
        {
            ...existingSemanticCustomizations,
            enabled: true,
            rules: semanticRules
        },
        configurationTarget
    );

    const existingWorkbenchColors = configuration.get<Record<string, unknown>>("workbench.colorCustomizations") ?? {};
    await configuration.update(
        "workbench.colorCustomizations",
        {
            ...existingWorkbenchColors,
            "editor.foreground": activeTheme.editorPalette.baseForeground,
            "editor.selectionBackground": activeTheme.editorPalette.selectionBackground,
            "editorLineNumber.foreground": activeTheme.editorPalette.lineNumberForeground
        },
        configurationTarget
    );

    lastAppliedThemeSignature = themeSignature;
    output.appendLine(`Applied Prometeu visual theme '${activeTheme.themeId}' for language '${visualThemesPayload.frontendLanguageId}'.`);
}

function extractPrometeuVisualThemesPayload(initializeResult: unknown): PrometeuVisualThemesPayload | undefined {
    if (!initializeResult || typeof initializeResult !== "object") {
        return undefined;
    }

    const capabilities = readObject((initializeResult as { capabilities?: unknown }).capabilities);
    const experimental = readObject(capabilities.experimental);
    const payload = readObject(experimental.prometeuVisualThemes);

    if (typeof payload.frontendLanguageId !== "string" || typeof payload.activeVisualThemeId !== "string") {
        return undefined;
    }
    if (!Array.isArray(payload.visualThemes)) {
        return undefined;
    }

    return payload as unknown as PrometeuVisualThemesPayload;
}

function extractPrometeuSemanticHostProjectionsPayload(initializeResult: unknown): PrometeuSemanticHostProjectionsPayload | undefined {
    if (!initializeResult || typeof initializeResult !== "object") {
        return undefined;
    }

    const capabilities = readObject((initializeResult as { capabilities?: unknown }).capabilities);
    const experimental = readObject(capabilities.experimental);
    const payload = readObject(experimental.prometeuSemanticHostProjections);

    if (typeof payload.frontendLanguageId !== "string" || !Array.isArray(payload.semanticKeys) || !Array.isArray(payload.hostProjections)) {
        return undefined;
    }

    return payload as unknown as PrometeuSemanticHostProjectionsPayload;
}

function selectHostProjection(
    payload: PrometeuSemanticHostProjectionsPayload | undefined,
    hostId: string
): PrometeuSemanticHostProjection | undefined {
    if (!payload) {
        return undefined;
    }
    return payload.hostProjections.find((projection) => projection.hostId === hostId);
}

function buildSemanticTokenRules(
    languageId: string,
    theme: PrometeuVisualTheme,
    hostProjection: PrometeuSemanticHostProjection | undefined
): Record<string, string | { foreground: string; italic?: boolean; bold?: boolean; underline?: boolean }> {
    const rules: Record<string, string | { foreground: string; italic?: boolean; bold?: boolean; underline?: boolean }> = {};
    const projectionBySemanticKey = new Map<string, PrometeuSemanticHostProjectionEntry>(
        (hostProjection?.tokenProjections ?? []).map((projection) => [projection.semanticKey, projection])
    );
    for (const tokenStyle of theme.tokenStyles) {
        const styleValue = !tokenStyle.italic && !tokenStyle.bold && !tokenStyle.underline
            ? tokenStyle.foreground
            : {
                foreground: tokenStyle.foreground,
                ...(tokenStyle.italic ? { italic: true } : {}),
                ...(tokenStyle.bold ? { bold: true } : {}),
                ...(tokenStyle.underline ? { underline: true } : {})
            };
        rules[`${tokenStyle.semanticKey}:${languageId}`] = styleValue;

        const projection = projectionBySemanticKey.get(tokenStyle.semanticKey);
        if (!projection) {
            continue;
        }

        const projectedSelector = buildSemanticSelector(languageId, projection.hostTokenType, projection.hostTokenModifiers);
        if (!(projectedSelector in rules)) {
            rules[projectedSelector] = styleValue;
        }

        const fallbackSelector = buildSemanticSelector(languageId, projection.fallbackTokenType, projection.fallbackTokenModifiers);
        if (!(fallbackSelector in rules)) {
            rules[fallbackSelector] = styleValue;
        }
    }
    return rules;
}

function buildSemanticSelector(
    languageId: string,
    tokenType: string,
    modifiers: string[]
): string {
    const normalizedModifiers = [...modifiers].sort();
    if (normalizedModifiers.length === 0) {
        return `${tokenType}:${languageId}`;
    }
    return `${tokenType}.${normalizedModifiers.join(".")}:${languageId}`;
}

function readObject(value: unknown): Record<string, unknown> {
    if (!value || typeof value !== "object" || Array.isArray(value)) {
        return {};
    }
    return value as Record<string, unknown>;
}