import { z } from "zod";
import type { ControlSpec } from "../layer2/createControlPanel.js";
import { createSystemContainer, finalize, runBuild } from "../layer2/orchestration.js";
import type { ToolContext, ToolRegistrar } from "../types.js";

const q = (value: string): string => JSON.stringify(value);

/**
 * Strange-attractor deferred generator. A Script CHOP integrates a chosen ODE system
 * (Lorenz / Aizawa / Halvorsen) using sub-stepped Euler integration, maintaining a rolling
 * ring buffer of `trail_length` points. A Script SOP converts those channels into a single
 * open polyline. An optional Tube SOP thickens it for shaded render.
 *
 * Closing Roadmap Milestone 4 "strange-attractor deferred generator".
 *
 * Paused-timeline note: with TD's timeline paused the Script CHOP integration freezes —
 * this is expected (time-dependent) and is flagged in warnings.
 */

const rgb = z.coerce.number().min(0).max(1);

const ATTRACTOR_DEFAULTS: Record<string, Record<string, number>> = {
  lorenz: { sigma: 10, rho: 28, beta: 8 / 3 },
  aizawa: { a: 0.95, b: 0.7, c: 0.6, d: 3.5, e: 0.25, f: 0.1 },
  halvorsen: { a: 1.4 },
};

// Rough bounding radius per attractor for auto-frame (deterministic; no live bound query).
const ATTRACTOR_RADIUS: Record<string, number> = {
  lorenz: 30,
  aizawa: 2,
  halvorsen: 6,
};

export const createStrangeAttractorSchema = z.object({
  name: z.string().default("strange_attractor").describe("Container baseCOMP name."),
  parent: z
    .string()
    .default("/project1")
    .describe("Parent network where the container is created."),
  attractor: z
    .enum(["lorenz", "aizawa", "halvorsen"])
    .default("lorenz")
    .describe("ODE system to integrate: lorenz (classic butterfly), aizawa, or halvorsen."),
  trail_length: z
    .number()
    .int()
    .min(100)
    .max(20000)
    .default(4000)
    .describe(
      "Points retained in the rolling ring buffer. Higher = longer ribbon, costlier SOP cook.",
    ),
  steps_per_frame: z
    .number()
    .int()
    .min(1)
    .max(50)
    .default(6)
    .describe("RK-style integration sub-steps per cook frame (controls speed and smoothness)."),
  dt: z
    .number()
    .min(0.0001)
    .max(0.05)
    .default(0.005)
    .describe("Integrator time step. Smaller = smoother but slower trajectory."),
  seed: z
    .tuple([z.number(), z.number(), z.number()])
    .default([0.1, 0.0, 0.0])
    .describe(
      "Initial state [x, y, z]. A tiny non-zero offset avoids the Lorenz fixed-point stall at the origin.",
    ),
  params: z
    .record(z.string(), z.number())
    .optional()
    .describe(
      "Override ODE constants. Lorenz: sigma, rho, beta. Aizawa: a, b, c, d, e, f. Halvorsen: a. Unknown keys are ignored.",
    ),
  thickness: z
    .number()
    .min(0)
    .default(0.015)
    .describe(
      "Tube SOP radius. Set to 0 to render the raw polyline (no Tube SOP — lighter on GPU).",
    ),
  color: z
    .tuple([rgb, rgb, rgb])
    .default([0.6, 0.9, 1.0])
    .describe("Constant MAT colour (RGB, 0..1)."),
  bg_color: z
    .tuple([rgb, rgb, rgb])
    .default([0.02, 0.02, 0.04])
    .describe("Render TOP background colour (RGB, 0..1)."),
  auto_frame: z
    .boolean()
    .default(true)
    .describe(
      "Auto-position camera based on attractor bounding radius (deterministic; no live bound query).",
    ),
  expose_controls: z
    .boolean()
    .default(true)
    .describe(
      "Expose StepsPerFrame / Dt / TrailLength / Thickness as custom parameters on the container.",
    ),
});
type CreateStrangeAttractorArgs = z.infer<typeof createStrangeAttractorSchema>;

/**
 * Script CHOP cook callback — integrates the ODE and fills tx/ty/tz ring-buffer channels.
 */
const INTEGRATE_CALLBACK = `# Auto-generated by create_strange_attractor.
import json

def cook(scriptOp):
    try:
        cfg = json.loads(op('config').text)
    except Exception:
        return
    attractor = cfg.get('attractor', 'lorenz')
    trail_length = cfg.get('trail_length', 4000)
    steps_per_frame = cfg.get('steps_per_frame', 6)
    dt = cfg.get('dt', 0.005)
    p = cfg.get('params', {})
    seed_xyz = cfg.get('seed', [0.1, 0.0, 0.0])

    # Live-override from parent custom pars if they exist.
    try:
        par = parent().par
        if hasattr(par, 'Stepsperframe'):
            steps_per_frame = int(par.Stepsperframe.eval())
        if hasattr(par, 'Dt'):
            dt = float(par.Dt.eval())
        if hasattr(par, 'Traillength'):
            trail_length = int(par.Traillength.eval())
    except Exception:
        pass

    storage = scriptOp.storage
    # Re-init when missing or trail_length changed.
    if 'state' not in storage or storage.get('trail_length') != trail_length:
        storage['state'] = list(seed_xyz)
        storage['buf'] = [[seed_xyz[0]], [seed_xyz[1]], [seed_xyz[2]]]
        storage['idx'] = 0
        storage['trail_length'] = trail_length

    state = storage['state']
    buf = storage['buf']
    idx = int(storage['idx'])

    def step_lorenz(x, y, z, dt2, p2):
        sigma = float(p2.get('sigma', 10))
        rho = float(p2.get('rho', 28))
        beta = float(p2.get('beta', 2.6666666666666665))
        dx = sigma * (y - x)
        dy = x * (rho - z) - y
        dz = x * y - beta * z
        return x + dx * dt2, y + dy * dt2, z + dz * dt2

    def step_aizawa(x, y, z, dt2, p2):
        a = float(p2.get('a', 0.95))
        b = float(p2.get('b', 0.7))
        c = float(p2.get('c', 0.6))
        d = float(p2.get('d', 3.5))
        e = float(p2.get('e', 0.25))
        f = float(p2.get('f', 0.1))
        dx = (z - b) * x - d * y
        dy = d * x + (z - b) * y
        dz = c + a * z - (z ** 3) / 3.0 - (x ** 2 + y ** 2) * (1 + e * z) + f * z * (x ** 3)
        return x + dx * dt2, y + dy * dt2, z + dz * dt2

    def step_halvorsen(x, y, z, dt2, p2):
        a2 = float(p2.get('a', 1.4))
        dx = -a2 * x - 4 * y - 4 * z - y * y
        dy = -a2 * y - 4 * z - 4 * x - z * z
        dz = -a2 * z - 4 * x - 4 * y - x * x
        return x + dx * dt2, y + dy * dt2, z + dz * dt2

    step_fn = {'lorenz': step_lorenz, 'aizawa': step_aizawa, 'halvorsen': step_halvorsen}.get(attractor, step_lorenz)

    x, y, z = float(state[0]), float(state[1]), float(state[2])
    for _ in range(steps_per_frame):
        x, y, z = step_fn(x, y, z, dt, p)
        # Ensure buf lists are long enough.
        while len(buf[0]) < trail_length:
            buf[0].append(0.0)
            buf[1].append(0.0)
            buf[2].append(0.0)
        buf[0][idx] = x
        buf[1][idx] = y
        buf[2][idx] = z
        idx = (idx + 1) % trail_length

    storage['state'] = [x, y, z]
    storage['idx'] = idx
    storage['buf'] = buf

    scriptOp.clear()
    scriptOp.numSamples = trail_length
    cx = scriptOp.appendChan('tx')
    cy = scriptOp.appendChan('ty')
    cz = scriptOp.appendChan('tz')
    # Rotate so oldest sample is index 0 and newest is last.
    for i in range(trail_length):
        ri = (idx + i) % trail_length
        cx[i] = buf[0][ri]
        cy[i] = buf[1][ri]
        cz[i] = buf[2][ri]
`;

/**
 * Script SOP cook callback — reads integrate CHOP and builds one open polyline.
 */
const TRAIL_CALLBACK = `# Auto-generated by create_strange_attractor.
def cook(scriptOp):
    integrate = op('integrate')
    if integrate is None:
        return
    trail_length = integrate.numSamples
    scriptOp.clear()
    if trail_length < 2:
        return
    p = scriptOp.appendPoly(trail_length, closed=False, addPoints=True)
    tx = integrate['tx']
    ty = integrate['ty']
    tz = integrate['tz']
    for i in range(trail_length):
        p[i].point.x = tx[i]
        p[i].point.y = ty[i]
        p[i].point.z = tz[i]
`;

export async function createStrangeAttractorImpl(
  ctx: ToolContext,
  args: CreateStrangeAttractorArgs,
) {
  return runBuild(async () => {
    const builder = await createSystemContainer(ctx, args.parent, args.name);

    // Merge per-attractor defaults with user overrides.
    const baseParams = { ...(ATTRACTOR_DEFAULTS[args.attractor] ?? {}) };
    if (args.params) {
      for (const [k, v] of Object.entries(args.params)) {
        if (k in baseParams) {
          baseParams[k] = v;
        }
      }
    }

    // Config DAT — the Script CHOP reads this each cook.
    const configDat = await builder.add("textDAT", "config");
    const configPayload = {
      attractor: args.attractor,
      params: baseParams,
      seed: args.seed,
      trail_length: args.trail_length,
      dt: args.dt,
      steps_per_frame: args.steps_per_frame,
    };
    await builder.python(`op(${q(configDat)}).text = ${q(JSON.stringify(configPayload, null, 2))}`);

    // Script CHOP + docked callback DAT.
    const integrate = await builder.add("scriptCHOP", "integrate");
    const integrateCb = await builder.add("textDAT", "integrate_cb");
    await builder.python(
      `_cb = op(${q(integrateCb)})\n_cb.text = ${q(INTEGRATE_CALLBACK)}\n_s = op(${q(integrate)})\n_s.par.callbacks = _cb.name`,
    );

    // Script SOP + docked callback DAT.
    const trail = await builder.add("scriptSOP", "trail");
    const trailCb = await builder.add("textDAT", "trail_cb");
    await builder.python(
      `_cb2 = op(${q(trailCb)})\n_cb2.text = ${q(TRAIL_CALLBACK)}\n_s2 = op(${q(trail)})\n_s2.par.callbacks = _cb2.name`,
    );

    // Optional Tube SOP to thicken the polyline.
    let geomSource = trail;
    let thicken: string | undefined;
    if (args.thickness > 0) {
      thicken = await builder.add("tubeSOP", "thicken", {
        rad1: args.thickness,
        rad2: args.thickness,
        cols: 4,
      });
      await builder.connect(trail, thicken);
      geomSource = thicken;
    }

    // Bound SOP for potential auto-frame reference (build only; camera uses precomputed radius).
    await builder.add("boundSOP", "bounds");

    // Geometry COMP with a Select SOP pointing at the final SOP output.
    const geo = await builder.add("geometryCOMP", "geo");
    const renderIn = await builder.add("selectSOP", "in1", { sop: geomSource }, geo);
    await builder.python(`_n = op(${q(renderIn)})\n_n.render = True\n_n.display = True`);

    const [cr, cg, cb] = args.color;
    const mat = await builder.add("constantMAT", "mat", { colorr: cr, colorg: cg, colorb: cb });
    await builder.setParams(geo, { material: mat });

    const radius = ATTRACTOR_RADIUS[args.attractor] ?? 30;
    const camTz = args.auto_frame ? radius * 3.0 : 30;
    const cam = await builder.add("cameraCOMP", "cam", { tz: camTz });
    const light = await builder.add("lightCOMP", "light", { tx: 3, ty: 4, tz: 4 });
    const [bgr, bgg, bgb] = args.bg_color;
    const render = await builder.add("renderTOP", "render", {
      camera: cam,
      geometry: geo,
      lights: light,
      bgcolorr: bgr,
      bgcolorg: bgg,
      bgcolorb: bgb,
      bgcolora: 1,
    });
    const out = await builder.add("nullTOP", "out1");
    await builder.connect(render, out);

    builder.warnings.push(
      "Script CHOP integration pauses when TD timeline is paused (time-dependent cook) — this is expected, not a bug. Resume playback to continue the trajectory.",
    );
    if (args.thickness > 0 && args.trail_length > 8000) {
      builder.warnings.push(
        `tubeSOP on ${args.trail_length} points produces ${args.trail_length * 4} quads — consider thickness=0 (raw polyline) for low-end machines.`,
      );
    }

    const controls: ControlSpec[] = args.expose_controls
      ? [
          {
            name: "StepsPerFrame",
            type: "int",
            min: 1,
            max: 50,
            default: args.steps_per_frame,
          },
          {
            name: "Dt",
            type: "float",
            min: 0.0001,
            max: 0.05,
            default: args.dt,
          },
          {
            name: "TrailLength",
            type: "int",
            min: 100,
            max: 20000,
            default: args.trail_length,
          },
          ...(thicken !== undefined
            ? [
                {
                  name: "Thickness",
                  type: "float" as const,
                  min: 0.001,
                  max: 0.2,
                  default: args.thickness,
                  bind_to: [`${thicken}.rad1`, `${thicken}.rad2`],
                },
              ]
            : []),
        ]
      : [];

    return finalize(ctx, {
      summary: `Built a ${args.attractor} strange-attractor network (trail_length=${args.trail_length}, steps_per_frame=${args.steps_per_frame}, dt=${args.dt}) → ${out}.`,
      builder,
      outputPath: out,
      controls,
      extra: {
        output_path: out,
        config_dat: configDat,
        integrate_chop: integrate,
        trail_sop: trail,
        thicken_sop: thicken ?? null,
        attractor: args.attractor,
        params: baseParams,
        trail_length: args.trail_length,
        thickness: args.thickness,
      },
    });
  });
}

export const registerCreateStrangeAttractor: ToolRegistrar = (server, ctx) => {
  server.registerTool(
    "create_strange_attractor",
    {
      title: "Create strange attractor",
      description:
        "Build a strange-attractor deferred geometry generator: a Script CHOP integrates a chosen ODE system (Lorenz / Aizawa / Halvorsen) with configurable sub-steps and maintains a rolling ring buffer of trail_length points. A Script SOP converts the channels into one open polyline; an optional Tube SOP thickens it for shaded render inside a Geometry COMP + Camera + Light + Render TOP pipeline. Closing Roadmap Milestone 4. Complements create_growth_system (L-systems) and create_particle_flock (boids) as the deterministic CPU-geometry idiom. With TD timeline paused the integrator pauses too (time-dependent) — resume playback to continue. Returns a summary plus a JSON block with the container path, output path, exposed controls, errors, warnings, and an inline preview.",
      inputSchema: createStrangeAttractorSchema.shape,
      annotations: { readOnlyHint: false, destructiveHint: false, openWorldHint: true },
    },
    (args) => createStrangeAttractorImpl(ctx, args),
  );
};