/* Copyright 2026 Marimo. All rights reserved. */
import { describe, expect, it } from "vitest";
import type { Outline } from "@/core/cells/outline";
import {
  canCollapseOutline,
  findCollapseRange,
  mergeOutlines,
  parseOutline,
} from "../outline";

describe("parseOutline", () => {
  it("can parse html outline", () => {
    const html = `
    <span class="markdown">
      <h1 id="welcome-to-marimo">Welcome to marimo! 🌊🍃</h1>
      <h2 id="what-is-marimo">What is <b>marimo</b>?</h2>
      <span class="paragraph">marimo is a Python library for creating reactive and interactive notebooks</span>
      <h2 id="how-do-i-use-marimo">How do I use marimo?</h2>
      <span class="paragraph">pip install marimo</span>
    </span>
    `;
    const outline = parseOutline({
      mimetype: "text/html",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "id": "welcome-to-marimo",
            },
            "level": 1,
            "name": "Welcome to marimo! 🌊🍃",
          },
          {
            "by": {
              "id": "what-is-marimo",
            },
            "html": "What is <b>marimo</b>?",
            "level": 2,
            "name": "What is marimo?",
          },
          {
            "by": {
              "id": "how-do-i-use-marimo",
            },
            "level": 2,
            "name": "How do I use marimo?",
          },
        ],
      }
    `);
  });

  it("can parse html outline with duplicate nested headings", () => {
    const html = `
    <span class="markdown">
      <h1 id="experiment-1">Experiment 1</h1>
      <h2 id="setup">Setup</h2>
      <h2 id="instructions">Instructions</h2>
      <h1 id="experiment-2">Experiment 2</h1>
      <h2 id="setup">Setup</h2>
      <h2 id="instructions">Instructions</h2>
      <h1 id="ack">Acknowledgements</h1>
      <h3 id="marimo">marimo</h1>
    </span>
    `;
    const outline = parseOutline({
      mimetype: "text/html",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "id": "experiment-1",
            },
            "level": 1,
            "name": "Experiment 1",
          },
          {
            "by": {
              "id": "setup",
            },
            "level": 2,
            "name": "Setup",
          },
          {
            "by": {
              "id": "instructions",
            },
            "level": 2,
            "name": "Instructions",
          },
          {
            "by": {
              "id": "experiment-2",
            },
            "level": 1,
            "name": "Experiment 2",
          },
          {
            "by": {
              "id": "setup",
            },
            "level": 2,
            "name": "Setup",
          },
          {
            "by": {
              "id": "instructions",
            },
            "level": 2,
            "name": "Instructions",
          },
          {
            "by": {
              "id": "ack",
            },
            "level": 1,
            "name": "Acknowledgements",
          },
          {
            "by": {
              "id": "marimo",
            },
            "level": 3,
            "name": "marimo",
          },
        ],
      }
    `);
  });

  it("can parse markdown outline", () => {
    const markdown = `
    <span class="markdown">
      <h1 id="introduction">Introduction</h1>
      <h2 id="getting-started">Getting Started</h2>
      <h3 id="installation">Installation</h3>
    </span>
    `;
    const outline = parseOutline({
      mimetype: "text/markdown",
      timestamp: 0,
      channel: "output",
      data: markdown,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "id": "introduction",
            },
            "level": 1,
            "name": "Introduction",
          },
          {
            "by": {
              "id": "getting-started",
            },
            "level": 2,
            "name": "Getting Started",
          },
          {
            "by": {
              "id": "installation",
            },
            "level": 3,
            "name": "Installation",
          },
        ],
      }
    `);
  });

  it("can handle non-html outline", () => {
    const html = "foo";
    const outline = parseOutline({
      mimetype: "text/plain",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toEqual(null);
  });

  it("can handle empty/null outline", () => {
    expect(parseOutline(null)).toEqual(null);
    expect(parseOutline(undefined!)).toEqual(null);

    const html = "";
    expect(
      parseOutline({
        mimetype: "text/html",
        timestamp: 0,
        channel: "output",
        data: html,
      }),
    ).toEqual({ items: [] });
  });

  it("can handle invalid outline", () => {
    const html = "<h1>foo</h1><h2>bar</h2><h3>baz</h3>";
    const outline = parseOutline({
      mimetype: "text/html",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "path": "//H1[contains(., "foo")]",
            },
            "level": 1,
            "name": "foo",
          },
          {
            "by": {
              "path": "//H2[contains(., "bar")]",
            },
            "level": 2,
            "name": "bar",
          },
          {
            "by": {
              "path": "//H3[contains(., "baz")]",
            },
            "level": 3,
            "name": "baz",
          },
        ],
      }
    `);
  });

  it("can parse h4, h5, and h6 headings", () => {
    const html = `
    <div>
      <h1 id="chapter">Chapter</h1>
      <h2 id="section">Section</h2>
      <h3 id="subsection">Subsection</h3>
      <h4 id="paragraph">Paragraph</h4>
      <h5 id="subparagraph">Subparagraph</h5>
      <h6 id="detail">Detail</h6>
    </div>
    `;
    const outline = parseOutline({
      mimetype: "text/html",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "id": "chapter",
            },
            "level": 1,
            "name": "Chapter",
          },
          {
            "by": {
              "id": "section",
            },
            "level": 2,
            "name": "Section",
          },
          {
            "by": {
              "id": "subsection",
            },
            "level": 3,
            "name": "Subsection",
          },
          {
            "by": {
              "id": "paragraph",
            },
            "level": 4,
            "name": "Paragraph",
          },
          {
            "by": {
              "id": "subparagraph",
            },
            "level": 5,
            "name": "Subparagraph",
          },
          {
            "by": {
              "id": "detail",
            },
            "level": 6,
            "name": "Detail",
          },
        ],
      }
    `);
  });

  it("preserves LaTeX HTML in headings", () => {
    const html = `
    <span class="markdown">
      <h1 id="intro">Introduction</h1>
      <h2 id="the-equation">The equation <marimo-tex class="arithmatex">||(E = mc^2||)</marimo-tex></h2>
      <h2 id="plain">Plain heading</h2>
    </span>
    `;
    const outline = parseOutline({
      mimetype: "text/html",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "id": "intro",
            },
            "level": 1,
            "name": "Introduction",
          },
          {
            "by": {
              "id": "the-equation",
            },
            "html": "The equation <marimo-tex class="arithmatex">||(E = mc^2||)</marimo-tex>",
            "level": 2,
            "name": "The equation ||(E = mc^2||)",
          },
          {
            "by": {
              "id": "plain",
            },
            "level": 2,
            "name": "Plain heading",
          },
        ],
      }
    `);
  });

  it("excludes headings within excluded tags", () => {
    const html = `
    <div>
      <h1 id="included-heading">Included Heading</h1>
      <marimo-carousel>
        <h2 id="excluded-heading-carousel">Excluded Heading in Carousel</h2>
      </marimo-carousel>
      <marimo-tabs>
        <h2 id="excluded-heading-tabs">Excluded Heading in Tabs</h2>
      </marimo-tabs>
      <marimo-accordion>
        <h2 id="excluded-heading-accordion">Excluded Heading in Accordion</h2>
      </marimo-accordion>
      <marimo-sidebar>
        <h2 id="excluded-heading-sidebar">Excluded Heading in Sidebar</h2>
      </marimo-sidebar>
      <h2 id="another-included-heading">Another Included Heading</h2>
    </div>
    `;
    const outline = parseOutline({
      mimetype: "text/html",
      timestamp: 0,
      channel: "output",
      data: html,
    });
    expect(outline).toMatchInlineSnapshot(`
      {
        "items": [
          {
            "by": {
              "id": "included-heading",
            },
            "level": 1,
            "name": "Included Heading",
          },
          {
            "by": {
              "id": "another-included-heading",
            },
            "level": 2,
            "name": "Another Included Heading",
          },
        ],
      }
    `);
  });
});

const OUTLINE_1: Outline = {
  items: [
    {
      name: "h1",
      level: 1,
      by: { id: "h1" },
    },
    {
      name: "h2",
      level: 2,
      by: { id: "h2" },
    },
    {
      name: "h3",
      level: 3,
      by: { id: "h3" },
    },
  ],
};

const OUTLINE_2: Outline = {
  items: [
    {
      name: "other-h1",
      level: 1,
      by: { path: "other-h1" },
    },
    {
      name: "other-h2",
      level: 2,
      by: { path: "other-h2" },
    },
  ],
};

it("mergeOutlines", () => {
  expect(mergeOutlines([OUTLINE_1, null, OUTLINE_2, null]))
    .toMatchInlineSnapshot(`
    {
      "items": [
        {
          "by": {
            "id": "h1",
          },
          "level": 1,
          "name": "h1",
        },
        {
          "by": {
            "id": "h2",
          },
          "level": 2,
          "name": "h2",
        },
        {
          "by": {
            "id": "h3",
          },
          "level": 3,
          "name": "h3",
        },
        {
          "by": {
            "path": "other-h1",
          },
          "level": 1,
          "name": "other-h1",
        },
        {
          "by": {
            "path": "other-h2",
          },
          "level": 2,
          "name": "other-h2",
        },
      ],
    }
  `);
});

const makeOutline = (levels: number[]) => {
  return {
    items: levels.map((level) => ({
      name: `h${level}`,
      level,
      by: { id: `h${level}` },
    })),
  };
};

it("canCollapseOutline", () => {
  expect(canCollapseOutline(null)).toBe(false);
  expect(canCollapseOutline(makeOutline([1]))).toBe(true);
  expect(canCollapseOutline(makeOutline([1, 2]))).toBe(true);
  expect(canCollapseOutline(makeOutline([2]))).toBe(true);
  expect(canCollapseOutline(makeOutline([3]))).toBe(true);
  expect(canCollapseOutline(makeOutline([1, 4]))).toBe(true);
  expect(canCollapseOutline(makeOutline([4]))).toBe(false);
});

describe("findCollapseRange", () => {
  it("can collapse range", () => {
    expect(findCollapseRange(0, [makeOutline([1, 2, 3, 4])])).toEqual([0, 0]);
  });

  it("can collapse range with gaps", () => {
    const outlines = [
      makeOutline([1, 2, 3, 4]),
      makeOutline([2, 3, 4]),
      null,
      makeOutline([2]),
      makeOutline([1]),
      makeOutline([2]),
    ];
    expect(findCollapseRange(0, outlines)).toEqual([0, 3]);
    expect(findCollapseRange(1, outlines)).toEqual([1, 2]);
    expect(findCollapseRange(4, outlines)).toEqual([4, 5]);
    expect(findCollapseRange(5, outlines)).toEqual([5, 5]);
    // bad ranges
    expect(findCollapseRange(10, outlines)).toEqual(null);
    expect(findCollapseRange(2, outlines)).toEqual(null);
  });
});