import { z } from "zod";
import { buildPayloadScript, parsePythonReport } from "../pythonReport.js";
import { errorResult, guardTd, jsonResult } from "../result.js";
import type { ToolContext, ToolRegistrar } from "../types.js";

export const createEnergyStructureSchema = z.object({
  name: z.string().min(1).describe("Parent COMP name to create under parent."),
  parent: z.string().default("/").describe("Parent path (default project root)."),
  audioSource: z
    .string()
    .optional()
    .describe(
      "Optional existing CHOP path producing audio (e.g. an Audio Device In or Audio File In). If omitted, an Audio Device In is created inside the COMP as 'audioin'.",
    ),
  windowSec: z
    .number()
    .min(2)
    .max(120)
    .default(20)
    .describe("Length of the rolling energy buffer (sec) used to compute adaptive mean/std."),
  buildThreshold: z
    .number()
    .min(0)
    .max(5)
    .default(0.7)
    .describe("k_build: state becomes BUILD when energy > mu + k_build*sigma."),
  dropThreshold: z
    .number()
    .min(0)
    .max(5)
    .default(0.85)
    .describe(
      "k_drop: state becomes DROP when energy > mu + k_drop*sigma (must be > buildThreshold).",
    ),
});
type CreateEnergyStructureArgs = z.infer<typeof createEnergyStructureSchema>;

interface EnergyStructureReport {
  comp: string;
  ops: string[];
  warnings: string[];
  fatal?: string;
}

// Builds a song-structure (build/drop/breakdown) edge detector COMP.
//
// Topology:
//   {parent}/{name}/
//     audioin    audiodeviceinCHOP  (only when audioSource not provided)
//     env        filterCHOP         long window — slow follower (tcomp=windowSec) → "energy"
//     script     scriptCHOP         rolling buffer + adaptive thresholds (storage)
//     out        nullCHOP           5-channel binding surface
//
// Adaptive thresholds: each cook appends the envelope sample to a rolling buffer
// (last windowSec*rate samples) stored in parent().fetch/store. mu = mean(buf),
// sigma = std(buf). build_level = mu + kB*sigma; drop_level = mu + kD*sigma.
// Hysteresis: 4 cooks above to step up, 30 cooks below to fall back to breakdown.
// Outputs energy (e/max), state (0/1/2), and three 1-sample edge channels.
const ENERGY_SCRIPT = `
import json, base64, traceback
_p = json.loads(base64.b64decode("__PAYLOAD_B64__").decode("utf-8"))
report = {"comp": "", "ops": [], "warnings": []}

def _try(label, fn):
    try:
        return fn()
    except Exception as _e:
        report["warnings"].append(label + ": " + str(_e))
        return None

try:
    _parent = op(_p["parent"])
    if _parent is None:
        report["fatal"] = "Parent not found: " + str(_p["parent"])
    else:
        _comp = _parent.create(baseCOMP, _p["name"])
        report["comp"] = _comp.path

        # --- Audio source ---
        _src = None
        if _p.get("audioSource"):
            _sel = _try("Select CHOP create", lambda: _comp.create(selectCHOP, "audioin"))
            if _sel is not None:
                _try("select chop par", lambda: setattr(_sel.par, "chop", _p["audioSource"]))
                report["ops"].append(_sel.path)
                _src = _sel
        else:
            _ain = _try("Audio Device In create", lambda: _comp.create(audiodeviceinCHOP, "audioin"))
            if _ain is not None:
                report["ops"].append(_ain.path)
                _src = _ain

        # --- Envelope follower (long window → "energy") ---
        # Use a Filter CHOP with time-constant smoothing; tcomp = windowSec.
        # envelope CHOP (TD 099) only models a short widthunit window and has no
        # attack/release params, so it can't model a 20s adaptive baseline cleanly.
        _env = _try("Filter CHOP create", lambda: _comp.create(filterCHOP, "env"))
        if _env is not None:
            report["ops"].append(_env.path)
            _try("env type", lambda: setattr(_env.par, "type", "gauss"))
            _try("env tcomp", lambda: setattr(_env.par, "tcomp", float(_p["windowSec"])))
            if _src is not None:
                _try("env connect", lambda: _env.inputConnectors[0].connect(_src))
            # Rename channel(s) to "energy" via a Rename CHOP would add nodes; instead
            # use a Reorder/Rename pattern: the Script CHOP below reads inputs[0][0]
            # by position, so the channel name on env is not load-bearing.

        # --- Script CHOP: rolling buffer + adaptive thresholds ---
        _script = _try("Script CHOP create", lambda: _comp.create(scriptCHOP, "script"))
        if _script is not None:
            report["ops"].append(_script.path)
            if _env is not None:
                _try("script connect", lambda: _script.inputConnectors[0].connect(_env))

            _cb_code = (
                "# Auto-generated by create_energy_structure. Rolling-buffer adaptive\\n"
                "# threshold detector (build/drop/breakdown) with hysteresis.\\n"
                "def onCook(scriptOp):\\n"
                "\\tparent_ = scriptOp.parent()\\n"
                "\\ttry:\\n"
                "\\t\\twin = float(parent_.par.Windowsec.eval())\\n"
                "\\t\\tkB = float(parent_.par.Buildthreshold.eval())\\n"
                "\\t\\tkD = float(parent_.par.Dropthreshold.eval())\\n"
                "\\texcept Exception:\\n"
                "\\t\\twin = 20.0; kB = 0.7; kD = 0.85\\n"
                "\\trate = max(1.0, float(scriptOp.rate))\\n"
                "\\tcap = int(win * rate)\\n"
                "\\tst = parent_.fetch('energy_state', {'buf': [], 'state': 0, 'above': 0, 'below': 0})\\n"
                "\\tbuf = st['buf']\\n"
                "\\tsrc = scriptOp.inputs[0] if scriptOp.inputs else None\\n"
                "\\te = 0.0\\n"
                "\\tif src is not None and src.numChans:\\n"
                "\\t\\ttry:\\n"
                "\\t\\t\\te = float(src[0][0])\\n"
                "\\t\\texcept Exception:\\n"
                "\\t\\t\\te = 0.0\\n"
                "\\tbuf.append(e)\\n"
                "\\tif len(buf) > cap:\\n"
                "\\t\\tdel buf[: len(buf) - cap]\\n"
                "\\tn = len(buf)\\n"
                "\\tif n < max(8, int(rate * 0.25)):\\n"
                "\\t\\tmu = e; sigma = 0.0; mx = max(1e-6, e)\\n"
                "\\telse:\\n"
                "\\t\\tmu = sum(buf) / n\\n"
                "\\t\\tvar = sum((x - mu) ** 2 for x in buf) / n\\n"
                "\\t\\tsigma = var ** 0.5\\n"
                "\\t\\tmx = max(buf) or 1e-6\\n"
                "\\tbuild_level = mu + kB * sigma\\n"
                "\\tdrop_level = mu + kD * sigma\\n"
                "\\tif drop_level <= build_level:\\n"
                "\\t\\tdrop_level = build_level + 1e-6\\n"
                "\\tprev = st['state']\\n"
                "\\tcandidate = 2 if e > drop_level else (1 if e > build_level else 0)\\n"
                "\\tif candidate > prev:\\n"
                "\\t\\tst['above'] += 1; st['below'] = 0\\n"
                "\\t\\tnew_state = candidate if st['above'] >= 4 else prev\\n"
                "\\telif candidate < prev:\\n"
                "\\t\\tst['below'] += 1; st['above'] = 0\\n"
                "\\t\\tnew_state = candidate if st['below'] >= 30 else prev\\n"
                "\\telse:\\n"
                "\\t\\tnew_state = prev\\n"
                "\\t\\tst['above'] = 0; st['below'] = 0\\n"
                "\\tbuild_edge = 1 if (new_state == 1 and prev != 1) else 0\\n"
                "\\tdrop_edge = 1 if (new_state == 2 and prev != 2) else 0\\n"
                "\\tbreak_edge = 1 if (new_state == 0 and prev != 0) else 0\\n"
                "\\tst['state'] = new_state\\n"
                "\\tparent_.store('energy_state', st)\\n"
                "\\tscriptOp.clear()\\n"
                "\\tscriptOp.numSamples = 1\\n"
                "\\tscriptOp.rate = rate\\n"
                "\\tscriptOp.appendChan('energy').vals = [max(0.0, min(1.0, e / mx))]\\n"
                "\\tscriptOp.appendChan('state').vals = [float(new_state)]\\n"
                "\\tscriptOp.appendChan('build_edge').vals = [float(build_edge)]\\n"
                "\\tscriptOp.appendChan('drop_edge').vals = [float(drop_edge)]\\n"
                "\\tscriptOp.appendChan('breakdown_edge').vals = [float(break_edge)]\\n"
            )
            # Try the callbacks DAT first; fall back to writing the op's own text.
            _ok = _try(
                "script callbacks DAT",
                lambda: (setattr(_script.par.callbacks.eval(), "text", _cb_code), True)[1],
            )
            if not _ok:
                _try("script inline text", lambda: setattr(_script, "text", _cb_code))

        # --- Null CHOP "out": binding handle ---
        _null = _try("Null CHOP create", lambda: _comp.create(nullCHOP, "out"))
        if _null is not None:
            report["ops"].append(_null.path)
            if _script is not None:
                _try("out connect", lambda: _null.inputConnectors[0].connect(_script))

        # --- Custom params on parent COMP (artist tweakables) ---
        def _add_par(label, default_val):
            try:
                page = _comp.appendCustomPage("Energy") if not any(
                    p.name == "Energy" for p in _comp.customPages
                ) else next(p for p in _comp.customPages if p.name == "Energy")
                par = page.appendFloat(label)
                par[0].default = default_val
                par[0].val = default_val
                return True
            except Exception as _e:
                report["warnings"].append("custom par " + label + ": " + str(_e))
                return False

        _add_par("Windowsec", float(_p["windowSec"]))
        _add_par("Buildthreshold", float(_p["buildThreshold"]))
        _add_par("Dropthreshold", float(_p["dropThreshold"]))

        # Bind Windowsec → env (filterCHOP) tcomp via expression so live tweaks
        # to the custom par flow into the filter time-constant too. Best-effort:
        # the rolling buffer cap is already re-read from parent().par.Windowsec
        # every onCook, so the detector remains live-tweakable even if this fails.
        if _env is not None:
            def _bind_tcomp():
                _env.par.tcomp.expr = "parent().par.Windowsec"
                _env.par.tcomp.mode = _env.par.tcomp.mode.EXPRESSION
                return True
            _try("env tcomp expr bind", _bind_tcomp)
except Exception:
    report["fatal"] = traceback.format_exc().splitlines()[-1]
print(json.dumps(report))
`;

export function buildEnergyStructureScript(payload: object): string {
  return buildPayloadScript(ENERGY_SCRIPT, payload);
}

export async function createEnergyStructureImpl(ctx: ToolContext, args: CreateEnergyStructureArgs) {
  if (args.dropThreshold <= args.buildThreshold) {
    return errorResult(
      `dropThreshold (${args.dropThreshold}) must be greater than buildThreshold (${args.buildThreshold}). The drop state is meant to fire on bigger excursions than the build state.`,
    );
  }
  return guardTd(
    async () => {
      const script = buildEnergyStructureScript({
        parent: args.parent,
        name: args.name,
        audioSource: args.audioSource ?? "",
        windowSec: args.windowSec,
        buildThreshold: args.buildThreshold,
        dropThreshold: args.dropThreshold,
      });
      const exec = await ctx.client.executePythonScript(script, true);
      return parsePythonReport<EnergyStructureReport>(exec.stdout);
    },
    (report) => {
      if (report.fatal) {
        return errorResult(`Energy-structure build failed: ${report.fatal}`, report);
      }
      const warnNote = report.warnings.length > 0 ? `, ${report.warnings.length} warning(s)` : "";
      const summary = `Built energy-structure detector at ${report.comp} → out (energy/state/build_edge/drop_edge/breakdown_edge), window ${args.windowSec}s, kB=${args.buildThreshold}, kD=${args.dropThreshold}${warnNote}. Bind via bind_audio_reactive to op('${report.comp}/out')['energy'].`;
      return jsonResult(summary, report);
    },
  );
}

export const registerCreateEnergyStructure: ToolRegistrar = (server, ctx) => {
  server.registerTool(
    "create_energy_structure",
    {
      title: "Create energy structure",
      description:
        "[experimental] Build a song-structure (build / drop / breakdown) edge detector COMP with adaptive thresholds. Listens to an existing audio CHOP (audioSource) or a freshly created Audio Device In, follows a long-window envelope, and runs a Script CHOP that maintains a rolling buffer (last windowSec seconds) to derive an adaptive mean (mu) and std (sigma). Emits a 5-channel Null CHOP `out` with: energy (smoothed RMS 0..1), state (0=breakdown, 1=build, 2=drop), and three 1-sample edge pulses build_edge / drop_edge / breakdown_edge. buildThreshold and dropThreshold are k-multipliers of sigma above mu (NOT absolute amplitudes), so the detector self-calibrates to the current mix loudness. Hysteresis (4 cooks above to step up, 30 below to fall back) stops chattering at thresholds. windowSec/Buildthreshold/Dropthreshold are exposed as custom params on the parent COMP so artists can tweak live. Default audio source builds an Audio Device In CHOP (may pop the macOS mic-permission dialog once — click Allow); pass audioSource to skip the device.",
      inputSchema: createEnergyStructureSchema.shape,
      annotations: { readOnlyHint: false, destructiveHint: false, openWorldHint: true },
    },
    (args) => createEnergyStructureImpl(ctx, args),
  );
};