// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

import { formatInputs } from "./formatters/inputFormatter.js";
import { formatGates } from "./formatters/gateFormatter.js";
import { formatRegisters } from "./formatters/registerFormatter.js";
import { processOperations } from "./process.js";
import {
  Circuit,
  CircuitGroup,
  ComponentGrid,
  Operation,
  SourceLocation,
  Qubit,
} from "./circuit.js";
import { GateRenderData } from "./gateRenderData.js";
import {
  gateHeight,
  minGateWidth,
  minToolboxHeight,
  svgNS,
} from "./constants.js";
import { createDropzones } from "./draggable.js";
import { enableEvents } from "./events.js";
import { createPanel, enableRunButton } from "./panel.js";
import { getMinMaxRegIdx } from "./utils.js";
import type { StateColumn } from "./state-viz/stateViz.js";
import type { PrepareStateVizOptions } from "./state-viz/worker/stateVizPrep.js";

/**
 * Contains render data for visualization.
 */
interface ComposedSqore {
  /** Width of visualization. */
  width: number;
  /** Height of visualization. */
  height: number;
  /** SVG elements the make up the visualization. */
  elements: SVGElement[];
}

/**
 * Defines the mapping of unique location to each operation. Used for enabling
 * interactivity.
 */
type GateRegistry = {
  [location: string]: Operation;
};

export type EditorHandlers = {
  editCallback: (circuitGroup: CircuitGroup) => void;
  // When provided, enables the Run button in the toolbox.
  runCallback?: () => void;
  // Optional callback to offload state visualization computation.
  // When provided (e.g., by the VS Code webview), the state visualizer can
  // compute state in a Web Worker without relying on globals.
  computeStateVizColumnsForCircuitModel?: (
    model: Circuit,
    opts?: PrepareStateVizOptions,
  ) => Promise<StateColumn[]>;
};

export type DrawOptions = {
  renderDepth?: number;
  renderLocations?: (l: SourceLocation[]) => { title: string; href: string };
  /**
   * When provided, enables editing behaviors (dropzones, run button, etc.) and
   * requires the callbacks necessary to support those behaviors.
   */
  editor?: EditorHandlers;
  /**
   * When provided, enables zoom-to-fit behavior. The callback is called with the new zoom level whenever it changes.
   */
  onZoomChange?: (zoomLevel: number) => void;
};

/**
 * Entrypoint class for rendering circuit visualizations.
 */
export class Sqore {
  circuit: Circuit;
  gateRegistry: GateRegistry = {};
  renderDepth: number;
  container: HTMLElement | null = null;
  zoomOnResize: boolean = true;
  zoomLevel: number = 100;
  /**
   * Initializes Sqore object.
   *
   * @param circuitGroup Group of circuits to be visualized.
   * @param options Optional rendering/interaction options.
   */
  constructor(
    public circuitGroup: CircuitGroup,
    private options: DrawOptions = {},
  ) {
    if (
      this.circuitGroup == null ||
      this.circuitGroup.circuits == null ||
      this.circuitGroup.circuits.length === 0
    ) {
      throw new Error(
        `No circuit found in file. Please provide a valid circuit.`,
      );
    }
    this.renderDepth = options.renderDepth ?? 0;
    // For now we only visualize the first circuit in the group
    this.circuit = this.circuitGroup.circuits[0];
  }

  /**
   * Render circuit into `container` at the specified layer depth.
   *
   * @param container HTML element for rendering visualization into.
   */
  draw(container: HTMLElement) {
    // Inject into container
    if (container == null) throw new Error(`Container not provided.`);

    this.container = container;

    this.renderCircuit(container);

    if (this.options.onZoomChange != null) {
      this.zoomToFit();
      window.addEventListener("resize", () => this.onResize());
    }
  }

  /**
   * Window resize handler to recalculate and set the zoom level
   * based on the new window width.
   */
  private onResize() {
    if (!this.zoomOnResize) {
      return;
    }

    // Recalculate the zoom level based on the container width
    this.zoomToFit();
  }

  /**
   * Calculate and set the zoom level to fit the circuit within the container.
   */
  private zoomToFit() {
    if (this.options.onZoomChange == null || this.container == null) {
      return;
    }
    const zoomLevel = this.calculateZoomToFit(this.container);
    this.updateZoomLevel(zoomLevel);
    this.options.onZoomChange?.(zoomLevel);
  }

  /**
   * Update the zoom level setting and apply it to the SVG element.
   */
  updateZoomLevel(zoomLevel: number) {
    this.zoomLevel = zoomLevel;
    const svg = this.container?.querySelector("svg.qviz");
    if (svg) {
      this.updateSvgWidth(svg as SVGElement, zoomLevel);
    }
  }

  /**
   * Update the width of the SVG element based on the zoom level.
   */
  updateSvgWidth(svg: SVGElement, zoomLevel: number) {
    // The width attribute contains the true width.
    // We'll leave this attribute untouched, so we can use it again if the
    // zoom level is ever updated.
    const width = svg.getAttribute("width")!;

    // We'll set the width in the style attribute to (true width * zoom level).
    // This value takes precedence over the true width in the width attribute.
    svg.setAttribute(
      "style",
      `max-width: ${width}; width: ${(parseInt(width) * (zoomLevel || 100)) / 100}; height: auto`,
    );
  }

  /**
   * Calculate the zoom level that will fit the circuit into the current size of the container.
   */
  calculateZoomToFit(container: HTMLElement): number {
    const svg = container.querySelector("svg.qviz") as SVGElement;
    const containerWidth = container.clientWidth;
    // width and height are the true dimensions generated by qviz
    const width = parseInt(svg.getAttribute("width")!);
    const height = svg.getAttribute("height")!;

    svg.setAttribute("viewBox", `0 0 ${width} ${height}`);
    const zoom = Math.min(Math.ceil((containerWidth / width) * 100), 100);
    return zoom;
  }

  /**
   * Render circuit into `container`.
   *
   * @param container HTML element for rendering visualization into.
   * @param circuit Circuit object to be rendered.
   */
  private renderCircuit(container: HTMLElement, circuit?: Circuit): void {
    // Create copy of circuit to prevent mutation
    const _circuit: Circuit =
      circuit ?? JSON.parse(JSON.stringify(this.circuit));

    // Assign unique locations to each operation
    _circuit.componentGrid.forEach((col, colIndex) =>
      col.components.forEach((op, i) =>
        this.fillGateRegistry(op, `${colIndex},${i}`),
      ),
    );

    // Expand operations to the specified render depth
    this.expandOperationsToDepth(_circuit.componentGrid, this.renderDepth);

    // Auto-expand any groups with single children
    this.expandIfSingleOperation(_circuit.componentGrid);

    // Create visualization components
    const composedSqore: ComposedSqore = this.compose(_circuit);
    const svg: SVGElement = this.generateSvg(composedSqore);
    this.setViewBox(svg);
    if (this.options.onZoomChange != null) {
      this.updateSvgWidth(svg, this.zoomLevel);
    }
    const previousSvg = container.querySelector("svg.qviz");
    if (previousSvg == null) {
      container.appendChild(svg);
    } else {
      const wrapper = previousSvg.parentElement;
      if (wrapper) {
        wrapper.replaceChild(svg, previousSvg);
      } else {
        container.replaceChild(svg, previousSvg);
      }
    }
    this.addGateClickHandlers(container, _circuit);

    const editor = this.options.editor;
    const isEditable = editor != null;
    if (isEditable) {
      createDropzones(container, this);
      createPanel(container, editor.computeStateVizColumnsForCircuitModel);
      if (editor.runCallback) {
        enableRunButton(container, editor.runCallback);
      }
      enableEvents(container, this, () => this.renderCircuit(container));
      editor.editCallback(this.minimizeCircuits(this.circuitGroup));
    }
  }

  private expandOperationsToDepth(
    componentGrid: ComponentGrid,
    targetDepth: number,
    currentDepth: number = 0,
  ) {
    for (const col of componentGrid) {
      for (const op of col.components) {
        if (currentDepth < targetDepth && op.children != null) {
          op.dataAttributes = op.dataAttributes || {};
          op.dataAttributes["expanded"] = "true";
          this.expandOperationsToDepth(
            op.children,
            targetDepth,
            currentDepth + 1,
          );
        }
      }
    }
  }

  private expandIfSingleOperation(grid: ComponentGrid) {
    if (grid.length == 1 && grid[0].components.length == 1) {
      const onlyComponent = grid[0].components[0];
      if (
        onlyComponent.dataAttributes != null &&
        Object.prototype.hasOwnProperty.call(
          onlyComponent.dataAttributes,
          "location",
        ) &&
        onlyComponent.dataAttributes["expanded"] !== "false"
      ) {
        const location: string = onlyComponent.dataAttributes["location"];
        this.expandOperation(grid, location);
      }
    }
    // Recursively expand if the only child is also a single operation
    for (const col of grid) {
      for (const op of col.components) {
        this.expandIfSingleOperation(op.children || []);
      }
    }
  }

  /**
   * Sets the viewBox attribute of the SVG element to enable zooming and panning.
   *
   * @param svg The SVG element to set the viewBox for.
   */
  private setViewBox(svg: SVGElement) {
    // width and height are the true dimensions generated by qviz
    const width = parseInt(svg.getAttribute("width")!);
    const height = parseInt(svg.getAttribute("height")!);
    svg.setAttribute("viewBox", `0 0 ${width} ${height}`);
  }

  /**
   * Generates the components required for visualization.
   *
   * @param circuit Circuit to be visualized.
   *
   * @returns `ComposedSqore` object containing render data for visualization.
   */
  private compose(circuit: Circuit): ComposedSqore {
    const add = (
      acc: GateRenderData[],
      gate: GateRenderData | GateRenderData[],
    ): void => {
      if (Array.isArray(gate)) {
        gate.forEach((g) => add(acc, g));
      } else {
        acc.push(gate);
        gate.children?.forEach((col) => col.forEach((g) => add(acc, g)));
      }
    };

    const flatten = (renderData: GateRenderData[][]): GateRenderData[] => {
      const result: GateRenderData[] = [];
      renderData.forEach((col) => col.forEach((g) => add(result, g)));
      return result;
    };

    const { qubits, componentGrid } = circuit;

    // Calculate the row heights, which may vary depending on how many
    // expanded group borders need to fit between qubit wires.
    const rowHeights = getRowHeights(qubits, componentGrid);

    const isEditable = this.options.editor != null;

    // Draw the qubit labels.
    // Also calculate other register render data to be used later in the rendering.
    const { qubitLabels, registers, svgHeight } = formatInputs(
      qubits,
      rowHeights,
      isEditable ? undefined : this.options.renderLocations,
    );

    // Calculate the render data for the operations.
    const topY = qubits[0] ? registers[qubits[0].id].y : -1;
    const bottomY = qubits[qubits.length - 1]
      ? registers[qubits[qubits.length - 1].id].y
      : -1;
    const { renderDataArray, svgWidth } = processOperations(
      componentGrid,
      topY,
      bottomY,
      registers,
      isEditable ? undefined : this.options.renderLocations,
    );

    // Draw the operations.
    const formattedGates: SVGElement = formatGates(renderDataArray);

    // Draw the lines that represent qubit and classical wires.
    const formattedRegs: SVGElement = formatRegisters(
      registers,
      flatten(renderDataArray),
      svgWidth,
    );

    const composedSqore: ComposedSqore = {
      width: svgWidth,
      height: svgHeight,
      elements: [qubitLabels, formattedRegs, formattedGates],
    };
    return composedSqore;
  }

  /**
   * Generates visualization of `composedSqore` as an SVG.
   *
   * @param composedSqore ComposedSqore to be visualized.
   *
   * @returns SVG representation of circuit visualization.
   */
  private generateSvg(composedSqore: ComposedSqore): SVGElement {
    const { width, height, elements } = composedSqore;

    const svg: SVGElement = document.createElementNS(svgNS, "svg");
    svg.setAttribute("class", "qviz");
    svg.setAttribute("width", width.toString());
    svg.setAttribute("height", height.toString());

    // Add styles
    document.documentElement.style.setProperty(
      "--minToolboxHeight",
      `${minToolboxHeight}px`,
    );
    document.documentElement.style.setProperty(
      "--minGateWidth",
      `${minGateWidth}px`,
    );
    document.documentElement.style.setProperty(
      "--gateHeight",
      `${gateHeight}px`,
    );

    // Add body elements
    elements.forEach((element: SVGElement) => svg.appendChild(element));

    return svg;
  }

  /**
   * Depth-first traversal to assign unique location string to `operation`.
   * The operation is assigned the location `location` and its `i`th child
   * in its `colIndex` column is recursively given the location
   * `${location}-${colIndex},${i}`.
   *
   * @param operation Operation to be assigned.
   * @param location: Location to assign to `operation`.
   *
   */
  private fillGateRegistry(operation: Operation, location: string): void {
    if (operation.dataAttributes == null) operation.dataAttributes = {};
    operation.dataAttributes["location"] = location;

    // Note: `dataAttributes["expanded"]` is intentionally not defaulted here.
    // Expansion is controlled by:
    // - `renderDepth` (see `expandOperationsToDepth`),
    // - user interaction (expand/collapse), and
    // - `expandIfSingleOperation`, which auto-expands a single top-level op
    //   unless it has been explicitly collapsed.
    this.gateRegistry[location] = operation;
    operation.children?.forEach((col, colIndex) =>
      col.components.forEach((childOp, i) => {
        this.fillGateRegistry(childOp, `${location}-${colIndex},${i}`);
      }),
    );
  }

  /**
   * Add interactive click handlers to circuit HTML elements.
   *
   * @param container HTML element containing visualized circuit.
   * @param circuit Circuit to be visualized.
   *
   */
  private addGateClickHandlers(container: HTMLElement, circuit: Circuit): void {
    this.addZoomHandlers(container, circuit);
  }

  /**
   * Add interactive click handlers for expand/collapse functionality.
   *
   * @param container HTML element containing visualized circuit.
   * @param circuit Circuit to be visualized.
   */
  private addZoomHandlers(container: HTMLElement, circuit: Circuit): void {
    container.querySelectorAll(".gate .gate-control").forEach((ctrl) => {
      // Zoom in on clicked gate
      ctrl.addEventListener("click", (ev: Event) => {
        const gateId: string | null | undefined =
          ctrl.parentElement?.getAttribute("data-location");
        if (typeof gateId == "string") {
          if (ctrl.classList.contains("gate-collapse")) {
            this.collapseOperation(circuit.componentGrid, gateId);
          } else if (ctrl.classList.contains("gate-expand")) {
            this.expandOperation(circuit.componentGrid, gateId);
          }
          this.zoomOnResize = false;
          this.renderCircuit(container, circuit);

          ev.stopPropagation();
        }
      });
    });
  }

  /**
   * Expand selected composite operation.
   *
   * @param componentGrid Grid of circuit components.
   * @param location Location of operation to expand.
   */
  private expandOperation(
    componentGrid: ComponentGrid,
    location: string,
  ): void {
    componentGrid.forEach((col) =>
      col.components.forEach((op) => {
        if (op.children != null) this.expandOperation(op.children, location);
        if (op.dataAttributes == null) return op;
        const opId: string = op.dataAttributes["location"];
        if (opId === location && op.children != null) {
          op.dataAttributes["expanded"] = "true";
        }
      }),
    );
  }

  /**
   * Collapse selected composite operation.
   *
   * @param componentGrid Grid of circuit components.
   * @param parentLoc Location of operation to collapse.
   */
  private collapseOperation(
    componentGrid: ComponentGrid,
    parentLoc: string,
  ): void {
    componentGrid.forEach((col) =>
      col.components.forEach((op) => {
        if (op.children != null) this.collapseOperation(op.children, parentLoc);
        if (op.dataAttributes == null) return op;
        const opId: string = op.dataAttributes["location"];
        if (opId === parentLoc) {
          // Explicitly collapse the targeted parent operation.
          op.dataAttributes["expanded"] = "false";
        } else if (opId.startsWith(parentLoc)) {
          // For descendants, remove the explicit expanded attribute rather than
          // forcing it to "false". This allows default expansion rules (e.g.
          // classically controlled groups default to expanded) to re-apply
          // correctly when the parent is later re-expanded.
          delete op.dataAttributes["expanded"];
        }
      }),
    );
  }

  // Minimize the circuits in a circuit group to remove dataAttributes
  minimizeCircuits(circuitGroup: CircuitGroup): CircuitGroup {
    // Create a deep copy of the circuit group
    const minimizedCircuits: CircuitGroup = JSON.parse(
      JSON.stringify(circuitGroup),
    );
    minimizedCircuits.circuits.forEach((circuit) => {
      circuit.componentGrid.forEach((col) => {
        col.components.forEach(this.minimizeOperation);
      });
    });
    return minimizedCircuits;
  }

  // Minimize the operation to remove dataAttributes
  minimizeOperation = (operation: Operation): void => {
    if (operation.children !== undefined) {
      operation.children.forEach((col) =>
        col.components.forEach(this.minimizeOperation),
      );
    }
    operation.dataAttributes = undefined;
  };
}

/**
 * Recursively computes vertical space required to render group borders.
 *
 * The resulting `heightAboveWire` and `heightBelowWire` values per qubit are
 * later used by `formatInputs` to leave sufficient space between qubit wires.
 *
 * @param qubits Array of qubits in the circuit.
 * @param componentGrid Grid of circuit components to traverse.
 *
 * @returns Mapping from qubit index to required heights above and below their wires.
 */
function getRowHeights(
  qubits: Qubit[],
  componentGrid: ComponentGrid,
): {
  [qubitIndex: number]: {
    heightAboveWire: number;
    heightBelowWire: number;
  };
} {
  const rowHeights: {
    [qubitIndex: number]: {
      currentGroupBordersAboveWire: number;
      currentGroupBordersBelowWire: number;
      heightAboveWire: number;
      heightBelowWire: number;
    };
  } = {};

  for (const q of qubits) {
    const { id } = q;
    rowHeights[id] = {
      currentGroupBordersBelowWire: 0,
      currentGroupBordersAboveWire: 0,
      heightBelowWire: 0,
      heightAboveWire: 0,
    };
  }

  updateRowHeights(componentGrid, rowHeights, qubits.length);
  return rowHeights;
}

function updateRowHeights(
  componentGrid: ComponentGrid,
  rowHeights: {
    [qubitIndex: number]: {
      currentGroupBordersAboveWire: number;
      currentGroupBordersBelowWire: number;
      heightAboveWire: number;
      heightBelowWire: number;
    };
  },
  numQubits: number,
) {
  for (const col of componentGrid) {
    for (const component of col.components) {
      if (isExpandedGroup(component)) {
        // We're in an expanded group. There is a dashed border above
        // the top qubit, and below the bottom qubit.
        const [topQubit, bottomQubit] = getMinMaxRegIdx(component);

        // Increment the current count of dashed group borders for
        // the top and bottom rows for this operation.
        // If the max height for this row has been exceeded above or below the wire,
        // update it.
        rowHeights[topQubit].currentGroupBordersAboveWire++;
        rowHeights[topQubit].heightAboveWire = Math.max(
          rowHeights[topQubit].heightAboveWire,
          rowHeights[topQubit].currentGroupBordersAboveWire,
        );

        rowHeights[bottomQubit].currentGroupBordersBelowWire++;
        rowHeights[bottomQubit].heightBelowWire = Math.max(
          rowHeights[bottomQubit].heightBelowWire,
          rowHeights[bottomQubit].currentGroupBordersBelowWire,
        );

        // recurse
        updateRowHeights(component.children || [], rowHeights, numQubits);

        // decrement
        rowHeights[topQubit].currentGroupBordersAboveWire--;
        rowHeights[bottomQubit].currentGroupBordersBelowWire--;
      }
    }
  }
}

/**
 * An "expanded group" here is any operation that is to be rendered showing
 * its children, with a dashed box around the children.
 */
function isExpandedGroup(component: Operation) {
  const expandedAttr = component.dataAttributes?.["expanded"];
  if (expandedAttr != null) {
    return expandedAttr === "true";
  }

  const hasChildren =
    component.children != null && component.children.length > 0;
  const hasClassicalControls =
    component.kind === "unitary" &&
    (((component.controls ?? []).some((reg) => reg.result != null) ?? false) ||
      (component.metadata?.controlResultIds?.length ?? 0) > 0);

  // Classically controlled groups default to expanded when not explicitly set.
  return hasChildren && hasClassicalControls;
}
