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);

/**
 * L-system / vine growth generator. A Script SOP runs a turtle that iterates a
 * context-free rewriting grammar from `axiom` for `generations` steps, then walks the
 * resulting string emitting unit-length forward strokes and `branchAngle` turns. The
 * resulting polyline tree is thickened with a Tube SOP, recentred, and rendered.
 *
 * Turtle alphabet:
 *   F  forward (draw)            f  forward (no draw)
 *   + -  yaw  ± branchAngle      & ^  pitch ± branchAngle      \ /  roll ± branchAngle
 *   [  push state                ]  pop state
 *   anything else = no-op (use X/A/B as grammar variables that expand but don't draw).
 *
 * Stochastic rules: when multiple rules share `from`, one is picked per step with
 * weighted-random choice (defaults to weight=1).
 */

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

const ruleSchema = z.object({
  from: z.string().min(1),
  to: z.string(),
  weight: z.coerce.number().positive().optional(),
});

export const createGrowthSystemSchema = z.object({
  name: z.string().default("growth_system").describe("Container baseCOMP name."),
  parent: z
    .string()
    .default("/project1")
    .describe("Parent network where the container is created."),
  rules: z
    .array(ruleSchema)
    .default([{ from: "F", to: "F[+F]F[-F][F]" }])
    .describe(
      "Context-free rewriting rules. Multiple rules sharing the same `from` symbol trigger weighted-random stochastic choice (weight defaults to 1).",
    ),
  generations: z.coerce
    .number()
    .int()
    .min(0)
    .max(7)
    .default(4)
    .describe(
      "Rewrite iterations. Capped at 7 because string length grows ~k^n and freezes the SOP cook.",
    ),
  axiom: z.string().min(1).default("F").describe("Initial string before rewriting."),
  branchAngle: z.coerce
    .number()
    .min(0)
    .max(180)
    .default(22.5)
    .describe("Turtle turn angle (degrees) for + / - / & / ^ / \\ / / symbols."),
  step_length: z.coerce.number().positive().default(0.15).describe("World units per F stroke."),
  thickness: z.coerce
    .number()
    .positive()
    .default(0.01)
    .describe("Tube SOP radius for the rendered branches."),
  color: z
    .tuple([rgb, rgb, rgb])
    .default([0.5, 0.9, 0.4])
    .describe("Constant MAT colour (RGB, 0..1)."),
  seed: z.coerce.number().int().default(1).describe("RNG seed for stochastic rule selection."),
  expose_controls: z
    .boolean()
    .default(true)
    .describe("Expose Generations / BranchAngle / StepLength / Thickness on the container."),
});
type CreateGrowthSystemArgs = z.infer<typeof createGrowthSystemSchema>;

/**
 * Script SOP cook callback. Reads the JSON rules from the docked text DAT, iterates the
 * grammar, then walks the resulting string as a 3D turtle, appending one polyline per
 * contiguous run of F's. A 200000-char safety brake stops runaway expansions.
 */
const GROW_CALLBACK = `# Auto-generated by create_growth_system.
import json, math, random

def cook(scriptOp):
    scriptOp.clear()
    try:
        cfg = json.loads(op('rules').text)
    except Exception:
        return
    # Live-override JSON defaults from parent COMP custom params if present, so
    # exposed controls (Generations/BranchAngle/StepLength/Seed) actually drive
    # the grammar instead of being inert.
    try:
        p = parent()
        for json_key, par_name in (
            ('generations', 'Generations'),
            ('branchAngle', 'BranchAngle'),
            ('step_length', 'StepLength'),
            ('seed', 'Seed'),
        ):
            par = getattr(p.par, par_name, None)
            if par is not None:
                cfg[json_key] = par.eval()
    except Exception:
        pass
    rules_by_from = {}
    for r in cfg.get('rules', []):
        rules_by_from.setdefault(r['from'], []).append(r)
    axiom = cfg.get('axiom', 'F')
    rng = random.Random(cfg.get('seed', 1))

    s = axiom
    for _ in range(int(cfg.get('generations', 0))):
        out = []
        for ch in s:
            choices = rules_by_from.get(ch)
            if not choices:
                out.append(ch)
                continue
            if len(choices) == 1:
                out.append(choices[0]['to'])
            else:
                weights = [c.get('weight', 1.0) for c in choices]
                out.append(rng.choices(choices, weights=weights, k=1)[0]['to'])
        s = ''.join(out)
        if len(s) > 200000:
            break

    step = float(cfg.get('step_length', 0.15))
    ang = math.radians(float(cfg.get('branchAngle', 22.5)))

    H = [0.0, 1.0, 0.0]
    L = [-1.0, 0.0, 0.0]
    U = [0.0, 0.0, 1.0]
    pos = [0.0, 0.0, 0.0]
    stack = []
    cur_poly = [tuple(pos)]

    def flush(cp):
        if len(cp) >= 2:
            try:
                p = scriptOp.appendPoly(len(cp), closed=False, addPoints=True)
                for i, pt in enumerate(cp):
                    p[i].point.x = pt[0]
                    p[i].point.y = pt[1]
                    p[i].point.z = pt[2]
            except Exception:
                pass

    def rot(axis, theta):
        ca, sa = math.cos(theta), math.sin(theta)
        ax, ay, az = axis
        def r(v):
            vx, vy, vz = v
            d = ax*vx + ay*vy + az*vz
            cx = ay*vz - az*vy
            cy = az*vx - ax*vz
            cz = ax*vy - ay*vx
            return [vx*ca + cx*sa + ax*d*(1-ca),
                    vy*ca + cy*sa + ay*d*(1-ca),
                    vz*ca + cz*sa + az*d*(1-ca)]
        return r

    for ch in s:
        if ch == 'F':
            pos = [pos[i] + H[i]*step for i in range(3)]
            cur_poly.append(tuple(pos))
        elif ch == 'f':
            flush(cur_poly)
            pos = [pos[i] + H[i]*step for i in range(3)]
            cur_poly = [tuple(pos)]
        elif ch == '+':
            r = rot(U,  ang); H = r(H); L = r(L); U = r(U)
        elif ch == '-':
            r = rot(U, -ang); H = r(H); L = r(L); U = r(U)
        elif ch == '&':
            r = rot(L,  ang); H = r(H); L = r(L); U = r(U)
        elif ch == '^':
            r = rot(L, -ang); H = r(H); L = r(L); U = r(U)
        elif ch == '\\\\':
            r = rot(H,  ang); H = r(H); L = r(L); U = r(U)
        elif ch == '/':
            r = rot(H, -ang); H = r(H); L = r(L); U = r(U)
        elif ch == '[':
            flush(cur_poly)
            stack.append((pos[:], H[:], L[:], U[:]))
            cur_poly = [tuple(pos)]
        elif ch == ']':
            flush(cur_poly)
            if stack:
                pos, H, L, U = [list(x) for x in stack.pop()]
            cur_poly = [tuple(pos)]
    flush(cur_poly)
    return
`;

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

    // The rules JSON DAT — the Script SOP reads this each cook to drive the grammar.
    const rulesDat = await builder.add("textDAT", "rules");
    const rulesPayload = {
      rules: args.rules,
      axiom: args.axiom,
      generations: args.generations,
      branchAngle: args.branchAngle,
      step_length: args.step_length,
      seed: args.seed,
    };
    await builder.python(`op(${q(rulesDat)}).text = ${q(JSON.stringify(rulesPayload, null, 2))}`);

    // Script SOP + docked callback DAT.
    const grow = await builder.add("scriptSOP", "grow");
    const growCb = await builder.add("textDAT", "grow_cb");
    await builder.python(
      `_cb = op(${q(growCb)})\n_cb.text = ${q(GROW_CALLBACK)}\n_s = op(${q(grow)})\n_s.par.callbacks = _cb.name`,
    );

    // Thicken polylines into tubes for shaded render.
    const thicken = await builder.add("tubeSOP", "thicken", {
      rad1: args.thickness,
      rad2: args.thickness,
      cols: 4,
    });
    await builder.connect(grow, thicken);

    // Recentre bbox around origin so camera framing is stable. A boundSOP
    // computes the centre, and the transform translates the geometry by its
    // negation. Wrapped in try/except in Python in case bound par names differ
    // across TD builds — a missing par falls back to a no-op transform.
    const bounds = await builder.add("boundSOP", "bounds");
    await builder.connect(thicken, bounds);
    const center = await builder.add("transformSOP", "center");
    await builder.connect(thicken, center);
    await builder.python(
      `try:\n` +
        `    _b = op(${q(bounds)})\n` +
        `    _c = op(${q(center)})\n` +
        `    _c.par.tx.expr = "-op('bounds').par.centerx"\n` +
        `    _c.par.ty.expr = "-op('bounds').par.centery"\n` +
        `    _c.par.tz.expr = "-op('bounds').par.centerz"\n` +
        `except Exception as _e:\n` +
        `    pass`,
    );

    // Geometry COMP renders the centred SOP. NetworkBuilder.add clears the default torus.
    const geo = await builder.add("geometryCOMP", "geo");
    // The render SOP lives inside the geo COMP — wire the centred SOP into a render-flagged
    // in SOP nested under geo. Simplest path: make `center` itself the rendered SOP by
    // moving its output through an in-SOP. We instead nest a small render SOP that pulls
    // from the outer chain via a Select SOP (cross-container handle).
    const renderIn = await builder.add("selectSOP", "in1", { sop: center }, 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 camDist = Math.max(2.0, args.generations * args.step_length * 8);
    const cam = await builder.add("cameraCOMP", "cam", { tz: camDist });
    const light = await builder.add("lightCOMP", "light", { tx: 3, ty: 4, tz: 4 });
    const render = await builder.add("renderTOP", "render", {
      camera: cam,
      geometry: geo,
      lights: light,
      bgcolorr: 0.02,
      bgcolorg: 0.02,
      bgcolorb: 0.03,
      bgcolora: 1,
    });
    const out = await builder.add("nullTOP", "out1");
    await builder.connect(render, out);

    builder.warnings.push(
      "L-system string growth is k^n in `generations` — capped at 7 with a 200000-char runtime brake. Script SOP cooks on parameter change (stable under timeline pause). Stochastic rules re-seed each cook from `seed`, so geometry is reproducible.",
    );

    const controls: ControlSpec[] = args.expose_controls
      ? [
          {
            name: "Generations",
            type: "int",
            min: 0,
            max: 7,
            default: args.generations,
          },
          {
            name: "BranchAngle",
            type: "float",
            min: 0,
            max: 180,
            default: args.branchAngle,
          },
          {
            name: "StepLength",
            type: "float",
            min: 0.01,
            max: 1,
            default: args.step_length,
          },
          {
            name: "Thickness",
            type: "float",
            min: 0.001,
            max: 0.1,
            default: args.thickness,
            bind_to: [`${thicken}.rad1`, `${thicken}.rad2`],
          },
        ]
      : [];

    return finalize(ctx, {
      summary: `Built an L-system growth network (axiom=${JSON.stringify(args.axiom)}, generations=${args.generations}, branchAngle=${args.branchAngle}°, ${args.rules.length} rule(s)) → ${out}.`,
      builder,
      outputPath: out,
      controls,
      extra: {
        output_path: out,
        rules_dat: rulesDat,
        grow_sop: grow,
        generations: args.generations,
        axiom: args.axiom,
        branch_angle: args.branchAngle,
        step_length: args.step_length,
        thickness: args.thickness,
      },
    });
  });
}

export const registerCreateGrowthSystem: ToolRegistrar = (server, ctx) => {
  server.registerTool(
    "create_growth_system",
    {
      title: "Create growth system",
      description:
        "Build an L-system / vine-growth generator: a Script SOP iterates a context-free rewriting grammar from `axiom` for `generations` steps, then walks the resulting string as a 3D turtle to draw a polyline tree. Recognised symbols: F (forward draw), f (forward no draw), + - (yaw ± branchAngle), & ^ (pitch), \\ / (roll), [ ] (push/pop state). Other symbols are no-op constants (use X/A/B as grammar variables that expand but don't draw). Multiple rules sharing a `from` symbol trigger weighted-random stochastic selection (`weight` defaults to 1; `seed` controls the RNG). The polyline tree is thickened with a Tube SOP, recentred, and rendered. Complements create_particle_flock (boids) and create_gpu_particle_field (curl-noise) as the deterministic CPU-geometry idiom. Returns a summary plus a JSON block with the container path, output path, rules DAT path, exposed controls, errors, warnings, and an inline preview.",
      inputSchema: createGrowthSystemSchema.shape,
      annotations: { readOnlyHint: false, destructiveHint: false, openWorldHint: true },
    },
    (args) => createGrowthSystemImpl(ctx, args),
  );
};