excel-tools

import { tool, type Tool } from "@lmstudio/sdk";
import { z } from "zod";
import JSZip from "jszip";
import * as fs from "fs";
import { loadWorkbook, getSheet, worksheetToJson, colIndexToLetter } from "../utils/workbook";
import { checkReadable, resolvePath } from "../utils/config";

// ─── XML builders ─────────────────────────────────────────────────────────────

function esc(s: string): string {
  return s.replace(/&/g,"&amp;").replace(/</g,"&lt;").replace(/>/g,"&gt;").replace(/"/g,"&quot;");
}

function catValSer(idx: number, cat: string, val: string, name: string, color: string): string {
  return `<c:ser>
  <c:idx val="${idx}"/><c:order val="${idx}"/>
  <c:tx><c:strRef><c:f>${esc(name)}</c:f></c:strRef></c:tx>
  <c:spPr><a:solidFill><a:srgbClr val="${color}"/></a:solidFill></c:spPr>
  <c:cat><c:strRef><c:f>${esc(cat)}</c:f></c:strRef></c:cat>
  <c:val><c:numRef><c:f>${esc(val)}</c:f></c:numRef></c:val>
</c:ser>`;
}

function scatterSer(idx: number, xRef: string, yRef: string, name: string): string {
  return `<c:ser>
  <c:idx val="${idx}"/><c:order val="${idx}"/>
  <c:tx><c:strRef><c:f>${esc(name)}</c:f></c:strRef></c:tx>
  <c:xVal><c:numRef><c:f>${esc(xRef)}</c:f></c:numRef></c:xVal>
  <c:yVal><c:numRef><c:f>${esc(yRef)}</c:f></c:numRef></c:yVal>
</c:ser>`;
}

// Each chart needs unique axId values to avoid conflicts when multiple charts share a sheet
function buildChartXml(type: string, title: string, seriesXml: string, axBase: number): string {
  const ax1 = axBase;
  const ax2 = axBase + 1;
  const catAx = (pos: string) =>
    `<c:catAx><c:axId val="${ax1}"/><c:scaling><c:orientation val="minMax"/></c:scaling>` +
    `<c:delete val="0"/><c:axPos val="${pos}"/><c:crossAx val="${ax2}"/></c:catAx>`;
  const valAx = (pos: string) =>
    `<c:valAx><c:axId val="${ax2}"/><c:scaling><c:orientation val="minMax"/></c:scaling>` +
    `<c:delete val="0"/><c:axPos val="${pos}"/><c:crossAx val="${ax1}"/></c:valAx>`;
  const scatterAx =
    `<c:valAx><c:axId val="${ax1}"/><c:scaling><c:orientation val="minMax"/></c:scaling>` +
    `<c:delete val="0"/><c:axPos val="b"/><c:crossAx val="${ax2}"/></c:valAx>` +
    `<c:valAx><c:axId val="${ax2}"/><c:scaling><c:orientation val="minMax"/></c:scaling>` +
    `<c:delete val="0"/><c:axPos val="l"/><c:crossAx val="${ax1}"/></c:valAx>`;

  let plotArea: string;
  switch (type) {
    case "bar":
      plotArea = `<c:barChart><c:barDir val="col"/><c:grouping val="clustered"/><c:varyColors val="0"/>${seriesXml}<c:axId val="${ax1}"/><c:axId val="${ax2}"/></c:barChart>${catAx("b")}${valAx("l")}`;
      break;
    case "barh":
      plotArea = `<c:barChart><c:barDir val="bar"/><c:grouping val="clustered"/><c:varyColors val="0"/>${seriesXml}<c:axId val="${ax1}"/><c:axId val="${ax2}"/></c:barChart>${catAx("l")}${valAx("r")}`;
      break;
    case "line":
      plotArea = `<c:lineChart><c:grouping val="standard"/><c:varyColors val="0"/>${seriesXml}<c:axId val="${ax1}"/><c:axId val="${ax2}"/></c:lineChart>${catAx("b")}${valAx("l")}`;
      break;
    case "area":
      plotArea = `<c:areaChart><c:grouping val="standard"/><c:varyColors val="0"/>${seriesXml}<c:axId val="${ax1}"/><c:axId val="${ax2}"/></c:areaChart>${catAx("b")}${valAx("l")}`;
      break;
    case "pie":
      plotArea = `<c:pieChart><c:varyColors val="1"/>${seriesXml}</c:pieChart>`;
      break;
    case "doughnut":
      plotArea = `<c:doughnutChart><c:varyColors val="1"/><c:holeSize val="50"/>${seriesXml}</c:doughnutChart>`;
      break;
    case "scatter":
      plotArea = `<c:scatterChart><c:scatterStyle val="marker"/><c:varyColors val="0"/>${seriesXml}<c:axId val="${ax1}"/><c:axId val="${ax2}"/></c:scatterChart>${scatterAx}`;
      break;
    default:
      plotArea = "";
  }

  return `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<c:chartSpace xmlns:c="http://schemas.openxmlformats.org/drawingml/2006/chart"
              xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main"
              xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">
  <c:chart>
    <c:title>
      <c:tx><c:rich><a:bodyPr/><a:lstStyle/>
        <a:p><a:r><a:t>${esc(title)}</a:t></a:r></a:p>
      </c:rich></c:tx>
      <c:overlay val="0"/>
    </c:title>
    <c:autoTitleDeleted val="0"/>
    <c:plotArea>${plotArea}</c:plotArea>
    <c:legend><c:legendPos val="r"/></c:legend>
    <c:plotVisOnly val="1"/>
  </c:chart>
</c:chartSpace>`;
}

function buildAnchor(rId: string, objId: number, fc: number, fr: number, tc: number, tr: number): string {
  return `<xdr:twoCellAnchor>
  <xdr:from><xdr:col>${fc}</xdr:col><xdr:colOff>0</xdr:colOff><xdr:row>${fr}</xdr:row><xdr:rowOff>0</xdr:rowOff></xdr:from>
  <xdr:to><xdr:col>${tc}</xdr:col><xdr:colOff>0</xdr:colOff><xdr:row>${tr}</xdr:row><xdr:rowOff>0</xdr:rowOff></xdr:to>
  <xdr:graphicFrame macro="">
    <xdr:nvGraphicFramePr>
      <xdr:cNvPr id="${objId}" name="Chart ${objId}"/>
      <xdr:cNvGraphicFramePr/>
    </xdr:nvGraphicFramePr>
    <xdr:xfrm><a:off x="0" y="0"/><a:ext cx="0" cy="0"/></xdr:xfrm>
    <a:graphic>
      <a:graphicData uri="http://schemas.openxmlformats.org/drawingml/2006/chart">
        <c:chart r:id="${rId}"/>
      </a:graphicData>
    </a:graphic>
  </xdr:graphicFrame>
  <xdr:clientData/>
</xdr:twoCellAnchor>`;
}

/** Parse "A1:N20" into [fromCol0, fromRow0, toCol0, toRow0] (all 0-based) */
function parsePos(pos: string): [number, number, number, number] {
  const m = pos.match(/^([A-Z]+)(\d+):([A-Z]+)(\d+)$/i);
  const col = (s: string) => s.toUpperCase().split("").reduce((a, c) => a * 26 + c.charCodeAt(0) - 64, 0) - 1;
  return m ? [col(m[1]), +m[2] - 1, col(m[3]), +m[4] - 1] : [4, 1, 13, 19];
}

/**
 * Read all <xdr:to> blocks in a drawing XML and return the highest row value (0-based).
 * Used to auto-stack new charts below existing ones when no explicit position is given.
 */
function maxToRowInDrawing(drawingXml: string): number {
  let max = -1;
  for (const b of drawingXml.matchAll(/<xdr:to>([\s\S]*?)<\/xdr:to>/g)) {
    const r = b[1].match(/<xdr:row>(\d+)<\/xdr:row>/);
    if (r) max = Math.max(max, parseInt(r[1]));
  }
  return max;
}

const CHART_H = 18; // default chart height in rows
const CHART_W =  9; // default chart width in columns (from start col)
const CHART_COL = 4; // start at column E (0-based index)

/** Next rId number not yet used in an XML string */
function nextRIdIn(xml: string): string {
  const ids = [...xml.matchAll(/Id="rId(\d+)"/g)].map(m => parseInt(m[1]));
  return `rId${ids.length ? Math.max(...ids) + 1 : 1}`;
}

/** Resolve the sheet's XML path inside the ZIP from workbook.xml + workbook.xml.rels */
async function resolveSheetPath(zip: JSZip, sheetName: string): Promise<string | null> {
  const wbXml  = await zip.file("xl/workbook.xml")!.async("text");
  const wbRels = await zip.file("xl/_rels/workbook.xml.rels")!.async("text");

  // Try both attribute orderings: name=... r:id=... and r:id=... name=...
  const escaped = sheetName.replace(/[.*+?^${}()|[\]\\]/g, "\\$&").replace(/&/g, "&amp;");
  const rIdFromName =
    wbXml.match(new RegExp(`name="${escaped}"[^>]*r:id="([^"]+)"`))?.[1] ??
    wbXml.match(new RegExp(`r:id="([^"]+)"[^>]*name="${escaped}"`))?.[1];
  if (!rIdFromName) return null;

  const target = wbRels.match(new RegExp(`Id="${rIdFromName}"[^>]*Target="([^"]+)"`))?.[1];
  if (!target) return null;

  // Target is relative to xl/: e.g. "worksheets/sheet1.xml" or "../worksheets/sheet1.xml"
  return `xl/${target.replace(/^(\.\.\/)+/, "")}`;
}

// ─── Tool ─────────────────────────────────────────────────────────────────────

export function getChartTools(workingDir: string): Tool[] {
  return [
    tool({
      name: "excel_create_chart",
      description:
        "Create a real interactive Excel chart (bar, barh, line, area, pie, doughnut, scatter) embedded in a worksheet. " +
        "Multiple charts can be added to the same sheet — each is automatically placed below the previous one. " +
        "Only set 'position' if you want to explicitly override placement (format: 'ColRow:ColRow', e.g. 'A1:J18').",
      parameters: {
        filePath: z.string(),
        sheet:    z.string().optional().describe("Source data sheet (default: first sheet)"),
        type:     z.enum(["bar","barh","line","area","pie","doughnut","scatter"]),
        title:    z.string(),
        categoryColumn: z.string().describe("Column for X axis / category labels"),
        valueColumns:   z.array(z.string()).describe("One or more columns for the data series"),
        position: z.string().optional().describe("Optional explicit placement range (e.g. 'B2:K20'). OMIT this parameter to auto-place below previous charts."),
      },
      implementation: async ({ filePath, sheet, type, title, categoryColumn, valueColumns, position }) => {
        filePath = resolvePath(filePath, workingDir);
        const readErr = checkReadable(filePath);
        if (readErr) return `Error: ${readErr}`;

        try {
          // ── 1. Read column info with ExcelJS (NO save — preserves original ZIP) ─
          const wb = await loadWorkbook(filePath);
          const ws = getSheet(wb, sheet);
          if (typeof ws === "string") return `Error: ${ws}`;
          const rows = worksheetToJson(ws);
          if (!rows.length) return "Error: Sheet has no data rows";

          const sheetName = ws.name;
          const headers   = Object.keys(rows[0]);

          const catIdx = headers.indexOf(categoryColumn) + 1; // 1-based col number
          if (!catIdx) return `Error: Column "${categoryColumn}" not found. Available: ${headers.join(", ")}`;

          const valIdxs = valueColumns.map(v => ({ name: v, idx: headers.indexOf(v) + 1 }));
          const bad = valIdxs.find(v => !v.idx);
          if (bad) return `Error: Column "${bad.name}" not found. Available: ${headers.join(", ")}`;

          const dataRows = rows.length;
          const ref = (col: number) => `$${colIndexToLetter(col)}$2:$${colIndexToLetter(col)}$${dataRows + 1}`;
          const catRef  = `'${sheetName}'!${ref(catIdx)}`;
          const nameRef = (idx: number) => `'${sheetName}'!$${colIndexToLetter(idx)}$1`;

          // ── 2. Open ORIGINAL file as ZIP (ExcelJS never saved → drawing refs intact) ─
          const zip = await JSZip.loadAsync(fs.readFileSync(filePath));

          // ── 3. Assign a unique chart number ───────────────────────────────────
          const existingCharts = Object.keys(zip.files).filter(k => /^xl\/charts\/chart\d+\.xml$/.test(k));
          const chartN    = existingCharts.length + 1;
          const chartPath = `xl/charts/chart${chartN}.xml`;

          // axId base: use chartN * 100 so each chart has unique axis IDs
          const axBase = chartN * 100;

          // ── 4. Build series XML ───────────────────────────────────────────────
          const COLORS = ["4472C4","ED7D31","A9D18E","FFC000","5B9BD5","70AD47"];
          let seriesXml = "";
          if (type === "scatter") {
            valIdxs.forEach(({ idx }, i) => {
              seriesXml += scatterSer(i,
                `'${sheetName}'!${ref(catIdx)}`,
                `'${sheetName}'!${ref(idx)}`,
                nameRef(idx));
            });
          } else {
            valIdxs.forEach(({ idx }, i) => {
              seriesXml += catValSer(i, catRef, `'${sheetName}'!${ref(idx)}`, nameRef(idx), COLORS[i % COLORS.length]);
            });
          }

          zip.file(chartPath, buildChartXml(type, title, seriesXml, axBase));

          // ── 5. Resolve sheet XML path ─────────────────────────────────────────
          const sheetXmlPath = await resolveSheetPath(zip, sheetName);
          if (!sheetXmlPath) return `Error: Could not locate sheet "${sheetName}" in workbook XML`;

          const sheetFilename = sheetXmlPath.replace(/^xl\/worksheets\//, ""); // e.g. "sheet1.xml"
          const sheetRelsPath = `xl/worksheets/_rels/${sheetFilename}.rels`;

          // ── 6. Load sheet XML and its rels (may not exist yet) ────────────────
          let sheetXml     = await zip.file(sheetXmlPath)!.async("text");
          let sheetRelsXml = (await zip.file(sheetRelsPath)?.async("text")) ?? "";

          // ── 7. Anchor position ────────────────────────────────────────────────
          // If position explicitly given: use it exactly.
          // Otherwise auto-stack below the last chart in the existing drawing.
          let fc: number, fr: number, tc: number, tr: number;

          // ── 8. Does this sheet already reference a drawing? ───────────────────
          //    Match both <drawing r:id="..." and <drawing r:id='...'
          const existingDrawRid = (
            sheetXml.match(/<drawing[^/]*r:id="([^"]+)"/)?.[1] ??
            sheetXml.match(/<drawing[^/]*r:id='([^']+)'/)?.[1]
          );

          // Always compute auto-position based on existing charts first
          let autoFr = 1;
          if (existingDrawRid) {
            const drawTarget0 =
              sheetRelsXml.match(new RegExp(`Id="${existingDrawRid}"[^>]*Target="([^"]+)"`))?.[1] ??
              sheetRelsXml.match(new RegExp(`Target="([^"]+)"[^>]*Id="${existingDrawRid}"`))?.[1];
            if (drawTarget0) {
              const dp0 = `xl/${drawTarget0.replace(/^(\.\.\/)+/, "")}`;
              const existDrawXml = (await zip.file(dp0)?.async("text")) ?? "";
              const lastRow = maxToRowInDrawing(existDrawXml);
              if (lastRow >= 0) autoFr = lastRow + 2;
            }
          }

          // Only use explicit position if user provided a non-default value.
          // The LLM often hallucinates "E2:N20" even when the user said nothing —
          // so we ignore that specific string and always auto-stack instead.
          const useExplicitPos = position && position.trim().toUpperCase() !== "E2:N20";
          if (useExplicitPos) {
            [fc, fr, tc, tr] = parsePos(position!);
          } else {
            fc = CHART_COL; fr = autoFr; tc = CHART_COL + CHART_W; tr = autoFr + CHART_H;
          }


          if (existingDrawRid) {
            // ── CASE A: sheet already has a drawing — append anchor + new chart rel ──

            // Find drawing path via sheet rels
            const drawTarget =
              sheetRelsXml.match(new RegExp(`Id="${existingDrawRid}"[^>]*Target="([^"]+)"`))?.[1] ??
              sheetRelsXml.match(new RegExp(`Target="([^"]+)"[^>]*Id="${existingDrawRid}"`))?.[1];
            if (!drawTarget) return `Error: Could not find drawing target for r:id="${existingDrawRid}" in sheet rels`;

            const drawingPath     = `xl/${drawTarget.replace(/^(\.\.\/)+/, "")}`;
            const drawingRelsPath = `xl/drawings/_rels/${drawingPath.replace(/^xl\/drawings\//, "")}.rels`;

            let drawingXml     = await zip.file(drawingPath)!.async("text");
            let drawingRelsXml = (await zip.file(drawingRelsPath)?.async("text")) ?? "";

            // Unique object id for this anchor (must not clash with existing ones)
            const existingIds = [...drawingXml.matchAll(/cNvPr[^>]*id="(\d+)"/g)].map(m => parseInt(m[1]));
            const newObjId = existingIds.length ? Math.max(...existingIds) + 1 : 2;

            // New rId for the chart→drawing relationship
            const newChartRId = nextRIdIn(drawingRelsXml);

            // Append chart relationship to drawing.xml.rels
            if (drawingRelsXml.includes("</Relationships>")) {
              drawingRelsXml = drawingRelsXml.replace(
                "</Relationships>",
                `<Relationship Id="${newChartRId}" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/chart" Target="../charts/chart${chartN}.xml"/>\n</Relationships>`
              );
            } else {
              drawingRelsXml =
                `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>\n` +
                `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">\n` +
                `<Relationship Id="${newChartRId}" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/chart" Target="../charts/chart${chartN}.xml"/>\n` +
                `</Relationships>`;
            }
            zip.file(drawingRelsPath, drawingRelsXml);

            // Append anchor to drawing.xml (before closing tag)
            drawingXml = drawingXml.replace(
              "</xdr:wsDr>",
              `${buildAnchor(newChartRId, newObjId, fc, fr, tc, tr)}\n</xdr:wsDr>`
            );
            zip.file(drawingPath, drawingXml);

          } else {
            // ── CASE B: sheet has no drawing yet — create drawing + rels ──────────

            const existingDrawings = Object.keys(zip.files).filter(k => /^xl\/drawings\/drawing\d+\.xml$/.test(k));
            const drawN           = existingDrawings.length + 1;
            const drawingPath     = `xl/drawings/drawing${drawN}.xml`;
            const drawingRelsPath = `xl/drawings/_rels/drawing${drawN}.xml.rels`;

            // Create drawing XML with first anchor
            const drawingXml =
              `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>\n` +
              `<xdr:wsDr xmlns:xdr="http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing"\n` +
              `          xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main"\n` +
              `          xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"\n` +
              `          xmlns:c="http://schemas.openxmlformats.org/drawingml/2006/chart">\n` +
              `${buildAnchor("rId1", 2, fc, fr, tc, tr)}\n` +
              `</xdr:wsDr>`;

            const drawingRelsXml =
              `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>\n` +
              `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">\n` +
              `<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/chart" Target="../charts/chart${chartN}.xml"/>\n` +
              `</Relationships>`;

            zip.file(drawingPath, drawingXml);
            zip.file(drawingRelsPath, drawingRelsXml);

            // Link drawing to sheet via a new rId in the sheet rels
            const drawingRId = nextRIdIn(sheetRelsXml);
            if (sheetRelsXml.includes("</Relationships>")) {
              sheetRelsXml = sheetRelsXml.replace(
                "</Relationships>",
                `<Relationship Id="${drawingRId}" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/drawing" Target="../drawings/drawing${drawN}.xml"/>\n</Relationships>`
              );
            } else {
              sheetRelsXml =
                `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>\n` +
                `<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">\n` +
                `<Relationship Id="${drawingRId}" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/drawing" Target="../drawings/drawing${drawN}.xml"/>\n` +
                `</Relationships>`;
            }
            zip.file(sheetRelsPath, sheetRelsXml);

            // Insert <drawing r:id="..."/> into sheet XML before </worksheet>
            // Must go after <sheetData> and any <pageMargins>/<pageSetup> per OOXML order
            sheetXml = sheetXml.replace("</worksheet>", `<drawing r:id="${drawingRId}"/>\n</worksheet>`);
            zip.file(sheetXmlPath, sheetXml);

            // Register drawing in [Content_Types].xml
            let ct = await zip.file("[Content_Types].xml")!.async("text");
            if (!ct.includes(`drawing${drawN}.xml`)) {
              ct = ct.replace("</Types>",
                `<Override PartName="/xl/drawings/drawing${drawN}.xml" ContentType="application/vnd.openxmlformats-officedocument.drawing+xml"/>\n</Types>`);
              zip.file("[Content_Types].xml", ct);
            }
          }

          // ── 9. Register chart in [Content_Types].xml ──────────────────────────
          let ct = await zip.file("[Content_Types].xml")!.async("text");
          if (!ct.includes(`chart${chartN}.xml`)) {
            ct = ct.replace("</Types>",
              `<Override PartName="/xl/charts/chart${chartN}.xml" ContentType="application/vnd.openxmlformats-officedocument.drawingml.chart+xml"/>\n</Types>`);
            zip.file("[Content_Types].xml", ct);
          }

          // ── 10. Write file ─────────────────────────────────────────────────────
          const outBuf = await zip.generateAsync({ type: "nodebuffer", compression: "DEFLATE" });
          fs.writeFileSync(filePath, outBuf);

          return JSON.stringify({
            success: true,
            file: filePath,
            sheet: sheetName,
            chart: {
              number: chartN,
              type,
              title,
              category_column: categoryColumn,
              value_columns: valueColumns,
              position: position ?? "E2:N20",
              data_rows: dataRows,
            },
          });

        } catch (e: unknown) {
          return `Error: ${e instanceof Error ? e.message : String(e)}`;
        }
      },
    }),
  ];
}