package com.elphel.imagej.cuas.rt;

import java.awt.Rectangle;
import java.io.File;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import com.sun.jna.Native;
import com.sun.jna.Pointer;

/**
 * In-process native DNN backend: runs the SAME L1+L2 inference as {@link CuasDnnRemote}, but via
 * LibTorch through JNA (libtpdnn.so, tile_processor_gpu/jna/tp_dnn.cpp) instead of TCP to a Python
 * server. No DGX / no server process. Each model's params (N/P/vr/out_ch, L2 ch_hidden) come from its
 * bundled {@code <model>.ts.pt.meta.json} sidecar (single source of truth). By Claude on 06/27/2026.
 *
 * Model resolution mirrors CuasDnnRemote's bundled-vs-override convention:
 *   - localDir set  -> &lt;localDir&gt;/&lt;model&gt;/model.ts.pt (+ .meta.json)   [dev/override]
 *   - localDir empty -> bundled jar resource /cuas_dnn/&lt;basename&gt;/model.ts.pt extracted to temp
 *
 * Requires {@code -Djna.library.path=<dir with libtpdnn.so>}; libtpdnn.so finds libtorch via its rpath
 * (set by build_dnn.sh) or LD_LIBRARY_PATH.
 */
public class CuasDnnLocal implements CuasDnnBackend {
	private static TpDnnJna LIB = null;
	private static synchronized TpDnnJna lib() {
		if (LIB == null) LIB = Native.load("tpdnn", TpDnnJna.class);
		return LIB;
	}

	private final TpDnnJna   lib;
	private final Pointer    ctx;
	private final int        N, P, vr, nvel;
	private int              nFrames = 0;
	private int              H = 0, W = 0;
	private File             tmpDir = null;     // holds extracted bundled models (deleted on close)

	/** @param localDir  override dir holding &lt;model&gt;/model.ts.pt, or empty/null to use bundled resources.
	 *  @param l1Model   L1 run name (e.g. "runs/weighted9_pm_s").
	 *  @param l2Model   L2 run name (e.g. "runs/mexhat_gaps_boost40"), or empty/null for L1-only. */
	public CuasDnnLocal(String localDir, String l1Model, String l2Model) throws Exception {
		this.lib = lib();
		String [] l1 = resolveModel(localDir, l1Model, ".ts.pt");      // {tsPath, metaJson}
		this.N  = metaInt(l1[1], "N", 9);
		this.P  = metaInt(l1[1], "P", 24);
		this.vr = metaInt(l1[1], "vr", 5);
		this.nvel = (2 * vr + 1) * (2 * vr + 1);
		String l2ts = "";
		int l2ch = 24;
		if ((l2Model != null) && !l2Model.isEmpty()) {
			String [] l2 = resolveModel(localDir, l2Model, ".l2.ts.pt");
			l2ts = l2[0];
			l2ch = metaInt(l2[1], "ch_hidden", 24);
		}
		this.ctx = lib.tpdnn_init(l1[0], l2ts, N, P, vr, l2ch);
		if (ctx == null) throw new Exception("tpdnn_init failed (l1=" + l1[0] + ", l2=" + l2ts + ")");
		System.out.println("CuasDnnLocal: native LibTorch backend up (L1=" + l1Model
				+ ", L2=" + ((l2Model == null || l2Model.isEmpty()) ? "off" : l2Model)
				+ "; N=" + N + " P=" + P + " vr=" + vr + " l2_ch=" + l2ch + ")");
	}

	@Override
	public int [] upload(float [][] frames, int H, int W) throws Exception {
		this.H = H; this.W = W;
		int T = frames.length, hw = H * W;
		float [] flat = new float [T * hw];                            // [T,H,W] row-major
		for (int t = 0; t < T; t++) System.arraycopy(frames[t], 0, flat, t * hw, hw);
		lib.tpdnn_upload(ctx, flat, T, H, W);
		this.nFrames = N;
		int nl = lib.tpdnn_num_levels(ctx);
		int [] counts = new int [nl];
		for (int i = 0; i < nl; i++) counts[i] = lib.tpdnn_level_frames(ctx, i);
		return counts;
	}

	@Override
	public int getNFrames() { return nFrames; }

	@Override
	public CuasDnnRemote.BatchResult inferBatch(int level, int start, int count, int stride, Rectangle roi,
			double rmaxCells, boolean l2Enable, boolean l2Reset, double noiseScale) throws Exception {
		int hw = H * W, rn = roi.width * roi.height;
		int nchAlloc = l2Enable ? 6 : 5;                              // native returns 6 with L2 (+age), else 5
		float [] o5  = new float [count * nchAlloc * hw];
		float [] rfb = new float [count * rn * nvel];
		long t0 = System.nanoTime();
		int nch = lib.tpdnn_infer(ctx, level, start, count, stride,
				roi.x, roi.y, roi.width, roi.height, rmaxCells,
				l2Enable ? 1 : 0, l2Reset ? 1 : 0, noiseScale, o5, rfb);
		double ms = (System.nanoTime() - t0) / 1.0e6;

		CuasDnnRemote.BatchResult r = new CuasDnnRemote.BatchResult();
		r.gpuMs = ms; r.H = H; r.W = W; r.count = count; r.nch = nch; r.nvel = nvel;
		r.rh = roi.height; r.rw = roi.width;
		r.offset5 = new float [count][nch][hw];                       // o5 layout [count][nch][hw]
		for (int s = 0; s < count; s++)
			for (int c = 0; c < nch; c++)
				System.arraycopy(o5, (s * nch + c) * hw, r.offset5[s][c], 0, hw);
		r.roiField = new float [count][rn][nvel];                     // rfb layout [count][rn][nvel]
		for (int s = 0; s < count; s++)
			for (int p = 0; p < rn; p++)
				System.arraycopy(rfb, (s * rn + p) * nvel, r.roiField[s][p], 0, nvel);
		return r;
	}

	@Override
	public void close() {
		try { if (ctx != null) lib.tpdnn_free(ctx); } catch (Exception e) { /* ignore */ }
		if (tmpDir != null) { try { deleteRec(tmpDir); } catch (Exception e) { /* ignore */ } }
	}

	// ---- model resolution + tiny meta parser ----

	/** Returns {tsPath, metaJsonContent} for the model. localDir override, else bundled resource->temp. */
	private String [] resolveModel(String localDir, String model, String suffix) throws Exception {
		String fileName = "model" + suffix;                          // model.ts.pt / model.l2.ts.pt
		if ((localDir != null) && !localDir.isEmpty()) {
			File ts   = new File(new File(localDir, model), fileName);
			File meta = new File(ts.getPath() + ".meta.json");
			if (!ts.exists()) throw new Exception("local model not found: " + ts);
			String mj = meta.exists() ? new String(Files.readAllBytes(meta.toPath())) : "";
			return new String [] { ts.getPath(), mj };
		}
		// bundled: /cuas_dnn/<basename(model)>/model.ts.pt (+ .meta.json) -> extract to a temp dir
		String base = model.substring(model.lastIndexOf('/') + 1);
		if (tmpDir == null) tmpDir = Files.createTempDirectory("cuas_dnn_").toFile();
		String res = "/cuas_dnn/" + base + "/" + fileName;
		Path tsTmp = extractResource(res, base + suffix);
		String mj = "";
		try (InputStream is = CuasDnnLocal.class.getResourceAsStream(res + ".meta.json")) {
			if (is != null) mj = new String(is.readAllBytes());
		}
		return new String [] { tsTmp.toString(), mj };
	}

	private Path extractResource(String res, String tmpName) throws Exception {
		try (InputStream is = CuasDnnLocal.class.getResourceAsStream(res)) {
			if (is == null) throw new Exception("bundled resource not found: " + res);
			Path out = new File(tmpDir, tmpName).toPath();
			Files.copy(is, out, StandardCopyOption.REPLACE_EXISTING);
			return out;
		}
	}

	private static int metaInt(String json, String key, int dflt) {
		if (json == null) return dflt;
		Matcher m = Pattern.compile("\"" + key + "\"\\s*:\\s*(-?\\d+)").matcher(json);
		return m.find() ? Integer.parseInt(m.group(1)) : dflt;
	}

	private static void deleteRec(File f) {
		File [] kids = f.listFiles();
		if (kids != null) for (File k : kids) deleteRec(k);
		f.delete();
	}
}
