import GObject, { register, property, signal } from "astal/gobject";
import {
  WorkspaceResponse,
  OutputResponse,
  SwayMessage,
  SwayEvent,
  WorkspaceEvent,
  OutputEvent,
  TreeNode,
  CommandResponse,
} from "./sway_types";
import Gio from "gi://Gio";
import GLib from "gi://GLib";
import * as utils from "@/utils";
import { Mutex } from "@/synchronization";
import { CompositorConnection } from "../connection";
import { ActionError } from "../errors";
import {
  EventType,
  OutputConfig,
  VSOutputAdapter,
  VSOutputEvent,
  VSWorkspaceAdapter,
  VSWorkspaceEvent,
} from "../types";

class SwayCommandError
  extends ActionError
  implements Required<CommandResponse>
{
  public readonly success: false = false;
  public readonly error: string;
  public readonly parse_error: boolean;

  constructor(
    response: Required<CommandResponse>,
    public readonly command?: string,
  ) {
    super(`Command ${command ?? "<unknown>"} failed: ${response.error}`);
    this.error = response.error;
    this.parse_error = response.parse_error;
  }

  /**
   * Run the given commands and assert that they were successful, raising a SwayCommandError if they were not.
   * When multiple commands are given, the first error is what will be raised.
   */
  public static async assertCommands(...commands: string[]) {
    const ipc = SwayIPC.instance;
    const results = await ipc.command(commands.join("; "));
    results.forEach((response, idx) => {
      if (!response.success) {
        throw new this(response as Required<CommandResponse>, commands[idx]);
      }
    });
  }
}

/**
 * Sway implementation of the CompositorConnection API
 */
@register({
  GTypeName: "SwayCompositorConnection",
})
export class SwayConnection extends CompositorConnection {
  private ipc: SwayIPC;

  private lastOutputs = new Map<string, VSOutputAdapter>();
  private lastWorkspaces = new Map<string, VSWorkspaceAdapter>();

  constructor() {
    super();
    this.ipc = SwayIPC.instance;
    this.onWorkspaceEvent();
    this.onOutputEvent();
    this.ipc.connect("workspace", () => {
      this.onWorkspaceEvent().catch(console.error);
    });
    this.ipc.connect("output", () => {
      this.onOutputEvent().catch(console.error);
    });
  }

  /**
   * Call when we receive a workspace event from Sway
   */
  private async onWorkspaceEvent() {
    const workspaces = await this.getWorkspaces();
    // Emit new events for all new outputs. No need to check set, this is the first emission.
    workspaces
      .filter((workspace) => !this.lastWorkspaces.has(workspace.id))
      .forEach((workspace) => {
        this.workspaceChange(
          new VSWorkspaceEvent({
            type: EventType.CREATE,
            workspace: workspace.id,
            adapter: workspace,
          }),
        );
      });
    // Emit changed events for all outputs that have changed. No need to check set, as this and previous loops are guaranteed disjoint.
    workspaces
      .filter((workspace) => this.lastWorkspaces.has(workspace.id))
      .filter(
        (workspace) => !workspace.equals(this.lastWorkspaces.get(workspace.id)),
      )
      .forEach((workspace) => {
        this.workspaceChange(
          new VSWorkspaceEvent({
            type: EventType.CHANGE,
            workspace: workspace.id,
            adapter: workspace,
          }),
        );
      });
    // Emit destroy events for all outputs that do not now exist.
    const workspaceNames = new Set(workspaces.map((workspace) => workspace.id));
    Array.from(this.lastWorkspaces.values())
      .filter((workspace) => !workspaceNames.has(workspace.id))
      .forEach((workspace) => {
        this.workspaceChange(
          new VSWorkspaceEvent({
            type: EventType.DESTROY,
            workspace: workspace.id,
            adapter: null,
          }),
        );
      });
    this.lastWorkspaces = new Map();
    workspaces.forEach((workspace) =>
      this.lastWorkspaces.set(workspace.id, workspace),
    );
  }

  /**
   * Call when we receive an output event from Sway
   *
   * NOTE: Currently, Sway does not provide any useful information with this event, so we need to get all outputs
   * ourselves and determine what changed.
   */
  private async onOutputEvent() {
    const outputs = await this.getOutputs();
    // Emit new events for all new outputs. No need to check set, this is the first emission.
    outputs
      .filter((output) => !this.lastOutputs.has(output.name))
      .forEach((output) => {
        this.outputChange(
          new VSOutputEvent({
            type: EventType.CREATE,
            output: output.name,
            adapter: output,
          }),
        );
      });
    // Emit changed events for all outputs that have changed. No need to check set, as this and previous loops are guaranteed disjoint.
    outputs
      .filter((output) => this.lastOutputs.has(output.name))
      .filter((output) => !output.equals(this.lastOutputs.get(output.name)))
      .forEach((output) => {
        this.outputChange(
          new VSOutputEvent({
            type: EventType.CHANGE,
            output: output.name,
            adapter: output,
          }),
        );
      });
    // Emit destroy events for all outputs that do not now exist.
    const outputNames = new Set(outputs.map((output) => output.name));
    Array.from(this.lastOutputs.values())
      .filter((output) => !outputNames.has(output.name))
      .forEach((output) => {
        this.outputChange(
          new VSOutputEvent({
            type: EventType.DESTROY,
            output: output.name,
            adapter: null,
          }),
        );
      });
    this.lastOutputs = new Map();
    outputs.forEach((output) => this.lastOutputs.set(output.name, output));
  }

  async focusWorkspace(workspaceId: string) {
    return SwayCommandError.assertCommands(`workspace ${workspaceId}`);
  }

  async moveWorkspace(workspaceId: string, outputName: string) {
    return SwayCommandError.assertCommands(
      `workspace ${workspaceId}`,
      `move workspace to output ${outputName}`,
    );
  }

  async focusOutput(outputName: string) {
    return SwayCommandError.assertCommands(`focus output ${outputName}`);
  }

  async moveContainer(workspaceId: string) {
    return SwayCommandError.assertCommands(
      `move container to workspace ${workspaceId}`,
    );
  }

  async configureOutput(
    outputName: string,
    { modeline, scale, position, ...extra_args }: OutputConfig,
  ) {
    const scale_setting = scale ? `scale ${scale}` : "";
    const extra_settings = Object.entries(extra_args).reduce(
      (settings, [name, arg]) => `${settings} ${name} ${arg}`,
      "",
    );
    return SwayCommandError.assertCommands(
      `output ${outputName} mode ${modeline} position ${position.join(" ")} ${scale_setting} ${extra_settings} enable`,
    );
  }

  async disableOutput(outputName: string) {
    return SwayCommandError.assertCommands(`output ${outputName} disable`);
  }

  async setOutputProperty<Prop extends string & keyof OutputConfig>(
    output: string,
    property: Prop,
    value: Exclude<OutputConfig[Prop], undefined>,
  ) {
    switch (property) {
      case "modeline":
        return SwayCommandError.assertCommands(
          `output ${output} mode ${value}`,
        );
      case "scale":
        return SwayCommandError.assertCommands(
          `output ${output} scale ${value}`,
        );
      case "position":
        return SwayCommandError.assertCommands(
          `output ${output} position ${(value as number[])[0]} ${(value as number[])[1]}`,
        );
      case "active":
        if (value) {
          return SwayCommandError.assertCommands(`output ${output} enable`);
        } else {
          return SwayCommandError.assertCommands(`output ${output} disable`);
        }
      case "powered":
        return SwayCommandError.assertCommands(
          `output ${output} power ${value ? "on" : "off"}`,
        );
      case "transform":
        return SwayCommandError.assertCommands(
          `output ${output} transform ${value}`,
        );
      case "colorProfile":
        return SwayCommandError.assertCommands(
          `output ${output} color_profile ${value}`,
        );
    }
  }

  private parseSwayOutput(output: OutputResponse): VSOutputAdapter {
    return new VSOutputAdapter({
      name: output.name,
      active: output.active,
      powered: output.power,
      rect: output.rect,
      scale: output.scale,
      transform: output.transform,
      current_workspace: output.current_workspace,
      make: output.make,
      model: output.model,
      serial: output.serial,
      availableModes: output.modes.map(({ picture_aspect_ratio, ...mode }) => ({
        aspect_ratio: picture_aspect_ratio,
        ...mode,
      })),
      mode: {
        height: output.current_mode.height,
        width: output.current_mode.width,
        refresh: output.current_mode.refresh,
        aspect_ratio: output.current_mode.picture_aspect_ratio,
      },
    });
  }

  private parseSwayWorkspace(workspace: WorkspaceResponse): VSWorkspaceAdapter {
    return new VSWorkspaceAdapter({
      id: workspace.name,
      output: workspace.output,
      focused: workspace.focused,
      visible: workspace.visible,
      urgent: workspace.urgent,
    });
  }

  async getOutputs(): Promise<VSOutputAdapter[]> {
    return (await this.ipc.getOutputs()).map((swayOutput) =>
      this.parseSwayOutput(swayOutput),
    );
  }

  async getWorkspaces(): Promise<VSWorkspaceAdapter[]> {
    return (await this.ipc.getWorkspaces()).map((swayWorkspace) =>
      this.parseSwayWorkspace(swayWorkspace),
    );
  }
}

/**
 * Low-level Sway IPC API
 */
@register({
  GTypeName: "SwayIPC",
})
class SwayIPC extends GObject.Object {
  #commandSocket: Gio.SocketConnection | null = null;
  #subscriptionSocket: Gio.SocketConnection | null = null;

  #mode: string = "";

  /**
   * Emitted when the connection to sway is established
   */
  @signal()
  declare connected: () => void;

  /**
   * Emitted when the connection to sway is lost
   */
  @signal()
  declare disconnected: () => void;

  // TODO: figure out how to give these signals a proper type. For now, Object is /fine/.

  /**
   * Emitted when we receive a subscription event
   */
  @signal(Object)
  declare subscribed: (events: SwayEvent[]) => void;

  /**
   * Emitted when we receive a workspace event
   */
  @signal(Object)
  declare workspace: (event: WorkspaceEvent) => void;

  /**
   * Emitted when we receive an output event
   */
  @signal(Object)
  declare output: (event: OutputEvent) => void;

  /**
   * The current binding mode of sway
   */
  @property(String)
  get mode(): string {
    return this.#mode;
  }

  @property(Boolean)
  get isConnected(): boolean {
    return this.#connected;
  }

  #connected = false;
  #onConnected: Promise<void>;

  /**
   * Run a command on sway
   */
  public async command(command: string): Promise<CommandResponse[]> {
    await this.#onConnected;
    return JSON.parse(
      (await this.#sendRaw(
        this.#commandSocket,
        SwayMessage.RUN_COMMAND,
        command,
      )) as string,
    );
  }

  /**
   * Get the current sway workspaces
   */
  public async getWorkspaces(): Promise<WorkspaceResponse[]> {
    await this.#onConnected;
    return JSON.parse(
      (await this.#sendRaw(
        this.#commandSocket,
        SwayMessage.GET_WORKSPACES,
      )) as string,
    );
  }

  /**
   * Get the current sway outputs
   */
  public async getOutputs(): Promise<OutputResponse[]> {
    await this.#onConnected;
    return JSON.parse(
      (await this.#sendRaw(
        this.#commandSocket,
        SwayMessage.GET_OUTPUTS,
      )) as string,
    );
  }

  /**
   * Get the current sway tree
   */
  public async getTree(): Promise<TreeNode> {
    await this.#onConnected;
    return JSON.parse(
      (await this.#sendRaw(
        this.#commandSocket,
        SwayMessage.GET_TREE,
      )) as string,
    );
  }

  constructor() {
    super();
    this.#onConnected = new Promise((resolve) => {
      const id = this.connect("connected", () => {
        resolve();
        this.disconnect(id);
      });
    });
    this.#connect();
  }

  static _instance: SwayIPC;

  /**
   * Get the default instance of Sway
   */
  static get instance() {
    if (!SwayIPC._instance) {
      SwayIPC._instance = new SwayIPC();
    }
    return SwayIPC._instance;
  }

  async #createConnection(address: string): Promise<Gio.SocketConnection> {
    console.log(`Connecting to sway socket ${address}`);
    const client = new Gio.SocketClient();
    return new Promise<Gio.SocketConnection>((resolve, reject) => {
      const socket_addr = new Gio.UnixSocketAddress({ path: address });
      client.connect_async(socket_addr, null, (client, result) => {
        try {
          // Type assertion is safe, if it fails we reject the promise in the catch block
          const connection = client!.connect_finish(result);
          resolve(connection);
        } catch (e) {
          console.error(`Failed to connect to socket.`, e);
          reject(e);
        }
      });
    });
  }

  async #getMode(): Promise<string> {
    const result = JSON.parse(
      (await this.#sendRaw(
        this.#commandSocket,
        SwayMessage.GET_BINDING_STATE,
      )) as string,
    );
    return result.name;
  }

  async #connect() {
    console.log("Connecting to sway");
    const address = GLib.getenv("SWAYSOCK");
    if (!address) {
      console.error("SWAYSOCK not set");
      throw new Error("SWAYSOCK not set");
    }

    this.#commandSocket = await this.#createConnection(address);
    this.#subscriptionSocket = await this.#createConnection(address);
    console.log("Connected to sway");

    // Start listening for subscriptions
    this.#startSubscriberListen();

    // Get the current mode
    this.#mode = await this.#getMode();
    this.notify("mode");

    this.connected();
    this.#connected = true;
    this.notify("is-connected");
  }

  static readonly MAGIC = "i3-ipc";
  static readonly HEADER_SIZE = SwayIPC.MAGIC.length + 8;
  static readonly CHUNK_SIZE = 4096;

  async #startSubscriberListen() {
    if (!this.#subscriptionSocket) {
      console.error("Not connected");
      throw new Error("Not connected");
    }
    await this.#subscribe(
      SwayEvent.OUTPUT,
      SwayEvent.WORKSPACE,
      SwayEvent.MODE,
      SwayEvent.SHUTDOWN,
    );
    const read = async () => {
      while (true) {
        try {
          const [event, payload] = await this.#readRaw(
            this.#subscriptionSocket!,
          );
          switch (event) {
            case SwayEvent.OUTPUT:
              this.output(JSON.parse(payload) as OutputEvent);
              break;

            case SwayEvent.WORKSPACE:
              this.workspace(JSON.parse(payload) as WorkspaceEvent);
              break;
            case SwayEvent.MODE:
              this.#mode = JSON.parse(payload).change;
              this.notify("mode");
              break;
            case SwayEvent.SHUTDOWN:
              this.disconnected();
              break;
            default:
              console.warn("Unhandled event", event);
              break;
          }
        } catch (e) {
          console.error("Failed to read event", e);
        }
      }
    };
    read();
  }

  async #readHeader(
    stream: Gio.InputStream,
  ): Promise<{ length: number; type: SwayMessage | SwayEvent }> {
    // We read one byte at a time until we have constructed a full magic string, then we read the rest of the header.
    // This allows us to handle partial reads and corrupted data gracefully.
    let idx = 0;
    while (idx < SwayIPC.MAGIC.length) {
      const buffer = await utils.readFromStreamRaw(stream, 1);
      const byte = buffer[0];
      if (byte !== SwayIPC.MAGIC.charCodeAt(idx)) {
        // Reset if we don't match the magic string
        idx = 0;
      } else {
        // Otherwise, keep going
        idx++;
      }
    }
    const header = new DataView((await utils.readFromStream(stream, 8)).buffer);
    const length = header.getUint32(0, true);
    const type = header.getUint32(4, true);
    return { length, type };
  }

  async #readRaw(
    socket: Gio.SocketConnection,
  ): Promise<[SwayMessage | SwayEvent, string]> {
    const inputStream = socket.input_stream;

    const { length, type } = await this.#readHeader(inputStream);
    const payloadBuf = await utils.readFromStream(inputStream, length);
    const payload = new TextDecoder().decode(payloadBuf);

    return [type as SwayMessage | SwayEvent, payload];
  }

  private messageMutex = new Mutex();

  async #sendRaw(
    socket: Gio.SocketConnection | null,
    type: SwayMessage,
    payload: string = "",
    waitForResponse: boolean = true,
  ): Promise<string | null> {
    return this.messageMutex.with(async () => {
      if (!socket || !socket.is_connected()) {
        console.error("Not connected");
        throw new Error("Not connected");
      }
      if (socket === this.#commandSocket) {
        if (type === SwayMessage.SUBSCRIBE) {
          console.error("Cannot subscribe on command socket");
          throw new Error("Cannot subscribe on command socket");
        }
      }
      // Construct the message
      const buffer = new ArrayBuffer(SwayIPC.HEADER_SIZE + payload.length);
      const magicView = new Uint8Array(buffer, 0, SwayIPC.MAGIC.length);
      const lengthView = new DataView(buffer, SwayIPC.MAGIC.length, 4);
      const typeView = new DataView(buffer, SwayIPC.MAGIC.length + 4, 4);
      const payloadView = new Uint8Array(buffer, SwayIPC.HEADER_SIZE);

      magicView.set(SwayIPC.MAGIC.split("").map((c) => c.charCodeAt(0)));
      lengthView.setUint32(0, payload.length, true);
      typeView.setUint32(0, type, true);
      payloadView.set(payload.split("").map((c) => c.charCodeAt(0)));

      const outputStream = socket.output_stream;

      // Send the message
      try {
        await utils.writeToStream(outputStream, buffer);
      } catch (e) {
        console.error("Failed to write to stream...", e);
        throw e;
      }
      if (!waitForResponse) {
        return null;
      }

      // Read the response
      const [resultType, result] = await this.#readRaw(socket);
      if (resultType !== type) {
        throw new Error(`Unexpected response type: ${resultType}`);
      }
      return result;
    });
  }

  async #subscribe(...events: SwayEvent[]) {
    if (!this.#subscriptionSocket) {
      console.error("Not connected");
      throw new Error("Not connected");
    }
    const eventNames = events.map((e) => SwayEvent[e].toLowerCase());
    const payload = JSON.stringify(eventNames);
    console.log("Subscribing to events:", payload);
    try {
      await this.#sendRaw(
        this.#subscriptionSocket,
        SwayMessage.SUBSCRIBE,
        payload,
      );

      console.log(`Subscribed to events: ${payload}`);
      this.subscribed(events);
    } catch (e) {
      console.error(`Failed to subscribe:`, e);
      throw e;
    }
  }
}