package com.elphel.imagej.cuas.rt;

import java.awt.Rectangle;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.file.Files;

/**
 * TCP client for the DGX remote DNN inference server (attic/imagej-elphel-internal/c5p_dnn/infer_server.py).
 * By Claude on 06/20/2026.
 *
 * Stateful, big-endian protocol (Java DataInput/OutputStream is big-endian; the Python side uses
 * struct '>' / numpy '>f4', so they match without byte-swapping):
 *   UPLOAD  : cmd=1, int T,H,W, then T*H*W float32 (the SUBAVG+LoG-conditioned stack)
 *             -> reply: int n_levels, n_levels x int frames_per_level, int N, double build_ms
 *             The DGX builds the temporal pyramid (0.5*(now+prev), == temporalAverageLReLU).
 *   INFER   : cmd=2, int level, newest, roi_x, roi_y, roi_w, roi_h
 *             -> reply: double gpu_ms, int H,W, 5*H*W float32 (offset5 dx,dy,s,Vx,Vy plane-major),
 *                       int rh,rw,nvel, rh*rw*nvel float32 (ROI softmax*s, pixel-major)
 *   READBACK: cmd=3, int level, frame -> reply: int H,W, H*W float32 (one conditioned pyramid frame)
 *   BYE     : cmd=0
 */
public class CuasDnnRemote implements AutoCloseable {
	private static final int CMD_BYE = 0, CMD_UPLOAD = 1, CMD_INFER = 2, CMD_READBACK = 3, CMD_STATUS = 4;
	private final Socket sock;
	private final DataInputStream  in;
	private final DataOutputStream out;
	private int nFrames = 0;   // temporal depth N, learned from the UPLOAD reply

	public CuasDnnRemote(String hostPort) throws Exception {
		String [] hp = hostPort.split(":");
		int port = (hp.length > 1) ? Integer.parseInt(hp[1].trim()) : 5577;
		sock = new Socket(hp[0].trim(), port);
		sock.setTcpNoDelay(true);
		in  = new DataInputStream (new BufferedInputStream (sock.getInputStream(),  1 << 20));
		out = new DataOutputStream(new BufferedOutputStream(sock.getOutputStream(), 1 << 20));
	}

	public int getNFrames() { return nFrames; }

	/** Upload the conditioned stack frames[T][H*W] (row-major) once; the DGX builds the pyramid.
	 *  Returns the per-level frame counts (and sets getNFrames()). */
	public int [] upload(float [][] frames, int H, int W) throws Exception {
		int T = frames.length;
		out.writeInt(CMD_UPLOAD); out.writeInt(T); out.writeInt(H); out.writeInt(W);
		byte [] buf = new byte [H * W * 4];                 // one frame, reused
		ByteBuffer bb = ByteBuffer.wrap(buf);               // big-endian by default
		for (int t = 0; t < T; t++) {
			bb.clear();
			float [] fr = frames[t];
			for (int p = 0; p < fr.length; p++) bb.putFloat(fr[p]);
			out.write(buf);
		}
		out.flush();
		int nl = in.readInt();
		int [] counts = new int [nl];
		for (int i = 0; i < nl; i++) counts[i] = in.readInt();
		nFrames = in.readInt();
		in.readDouble();                                    // build_ms (logged server-side)
		return counts;
	}

	/** Batched INFER result for `count` scenes (newest = start + s*stride): full-frame offset5
	 *  {dx,dy,s,Vx,Vy}[count][5][H*W] + ROI field [count][rh*rw][nvel] (softmax*s, ghostbusted on
	 *  the GPU) + the total pure-GPU compute ms (continuous = production throughput). */
	public static class BatchResult {
		public double gpuMs;
		public int H, W, count, nvel, rh, rw, nch;   // nch = offset channels: 5 {dx,dy,s,Vx,Vy} or 6 (+L2 age). By Claude 06/24/2026
		public float [][][] offset5;    // [count][nch][H*W]
		public float [][][] roiField;   // [count][rh*rw][nvel]
	}

	/** Infer `count` scenes of a level in one round-trip (newest_s = start + s*stride). rmaxCells>0
	 *  enables the on-GPU ghostbuster (== CuasDetectRT.dnnGhostbust). Keep `count` modest so the
	 *  reply byte[] stays < 2GB (count*5*H*W*4): count<=64 is ~419MB at 640x512. */
	public BatchResult inferBatch(int level, int start, int count, int stride, Rectangle roi, double rmaxCells,
			boolean l2Enable, boolean l2Reset, double noiseScale) throws Exception {
		out.writeInt(CMD_INFER); out.writeInt(level); out.writeInt(start); out.writeInt(count); out.writeInt(stride);
		out.writeInt(roi.x); out.writeInt(roi.y); out.writeInt(roi.width); out.writeInt(roi.height);
		out.writeDouble(rmaxCells);
		out.writeInt(l2Enable ? 1 : 0); out.writeInt(l2Reset ? 1 : 0);   // run Layer-2 on the DGX; reset hidden state at a level's first chunk. By Claude 06/22/2026
		out.writeDouble(noiseScale);   // per-level L1-input noise scale (Java is the source of truth; <=0 -> server sqrt fallback). By Claude 06/24/2026
		out.flush();
		BatchResult r = new BatchResult();
		r.gpuMs = in.readDouble(); r.H = in.readInt(); r.W = in.readInt();
		r.count = in.readInt(); r.nch = in.readInt(); r.nvel = in.readInt(); r.rh = in.readInt(); r.rw = in.readInt();   // +nch. By Claude 06/24/2026
		int hw = r.H * r.W, rn = r.rh * r.rw;
		byte [] ob = new byte [r.count * r.nch * hw * 4]; in.readFully(ob);
		ByteBuffer obb = ByteBuffer.wrap(ob);
		r.offset5 = new float [r.count][r.nch][hw];
		for (int s = 0; s < r.count; s++) for (int c = 0; c < r.nch; c++) for (int p = 0; p < hw; p++) r.offset5[s][c][p] = obb.getFloat();
		byte [] rb = new byte [r.count * rn * r.nvel * 4]; in.readFully(rb);
		ByteBuffer rbb = ByteBuffer.wrap(rb);
		r.roiField = new float [r.count][rn][r.nvel];
		for (int s = 0; s < r.count; s++) for (int p = 0; p < rn; p++) for (int v = 0; v < r.nvel; v++) r.roiField[s][p][v] = rbb.getFloat();
		return r;
	}

	/** Debug: read back one conditioned pyramid frame [H*W] (sets hw[0]=H, hw[1]=W if provided). */
	public float [] readback(int level, int frame, int [] hw) throws Exception {
		out.writeInt(CMD_READBACK); out.writeInt(level); out.writeInt(frame); out.flush();
		int H = in.readInt(), W = in.readInt();
		if ((hw != null) && (hw.length >= 2)) { hw[0] = H; hw[1] = W; }
		byte [] b = new byte [H * W * 4]; in.readFully(b);
		ByteBuffer bb = ByteBuffer.wrap(b);
		float [] f = new float [H * W];
		for (int i = 0; i < f.length; i++) f[i] = bb.getFloat();
		return f;
	}

	@Override
	public void close() {
		try { out.writeInt(CMD_BYE); out.flush(); } catch (Exception e) { /* ignore */ }
		try { sock.close(); } catch (Exception e) { /* ignore */ }
	}

	// ---- Auto-launch (no manual interaction): ensure the DGX server is running, else deploy + start it ----
	// Mirrors the GPU-kernel default-vs-override: server scripts come from the bundled jar resource
	// (/cuas_dnn/...) unless srcdir is set, then from that local dir. By Claude on 06/20/2026

	private static boolean canConnect(String host, int port, int timeoutMs) {
		try (Socket s = new Socket()) { s.connect(new InetSocketAddress(host, port), timeoutMs); return true; }
		catch (Exception e) { return false; }
	}

	/** Read a server script: bundled jar resource (/cuas_dnn/<name>) by default, or <srcdir>/<name> if set. */
	private static byte [] readServerScript(String name, String srcdir) throws Exception {
		if ((srcdir != null) && !srcdir.isEmpty()) {
			return Files.readAllBytes(new File(srcdir, name).toPath());           // local-repo override (config)
		}
		try (InputStream is = CuasDnnRemote.class.getResourceAsStream("/cuas_dnn/" + name)) {
			if (is == null) throw new Exception("bundled resource /cuas_dnn/" + name + " not found");
			return is.readAllBytes();                                             // working version in jar
		}
	}

	private static void deployScript(String name, String srcdir, String sshTarget, String dest) throws Exception {
		byte [] content = readServerScript(name, srcdir);
		ProcessBuilder pb = new ProcessBuilder("ssh", sshTarget, "cat > " + dest);
		pb.redirectErrorStream(true);
		Process p = pb.start();
		try (OutputStream os = p.getOutputStream()) { os.write(content); }
		if (p.waitFor() != 0) throw new Exception("deploy " + name + " to " + sshTarget + ":" + dest + " failed");
	}

	/** Query the running server for its loaded {L1 model, L2 model} paths (L2="" if L1-only), or null
	 *  if unreachable or the server is too old to answer CMD_STATUS. A short read timeout keeps an old
	 *  server (which ignores the opcode and never replies) from hanging the client. By Claude 06/24/2026 */
	private static String [] queryModels(String host, int port) {
		try (Socket s = new Socket()) {
			s.connect(new InetSocketAddress(host, port), 1500);
			s.setSoTimeout(2500);
			DataOutputStream o = new DataOutputStream(s.getOutputStream());
			DataInputStream  i = new DataInputStream (s.getInputStream());
			o.writeInt(CMD_STATUS); o.flush();
			int n1 = i.readInt(); byte [] b1 = new byte [n1]; i.readFully(b1);
			int n2 = i.readInt(); byte [] b2 = new byte [n2]; i.readFully(b2);
			try { o.writeInt(CMD_BYE); o.flush(); } catch (Exception e) { /* ignore */ }
			return new String [] { new String(b1, java.nio.charset.StandardCharsets.UTF_8),
			                       new String(b2, java.nio.charset.StandardCharsets.UTF_8) };
		} catch (Exception e) {
			return null;                                                          // unreachable, or old server (status read timed out)
		}
	}

	/** ssh-stop the DGX server (docker rm -f) so the next launch can load a different model. By Claude 06/24/2026 */
	private static void stopServer(String sshTarget, String code, int port) throws Exception {
		ProcessBuilder pb = new ProcessBuilder("ssh", sshTarget,
				"cd " + code + " && PORT=" + port + " ./run_infer_server.sh stop");
		pb.inheritIO();
		pb.start().waitFor();
	}

	/** Ensure the DGX server is up at hostPort WITH the requested L1+L2 models; if a different model is
	 *  loaded (or the server is an old build that can't report), tear it down and relaunch. Then deploy the
	 *  (bundled/override) scripts and ssh-launch run_infer_server.sh with RUN=model, polling until it accepts
	 *  connections. No manual step. All decisions logged to System.out (-> ImageJ log file). By Claude 06/24/2026 */
	public static void ensureServer(String hostPort, String model, String l2model, String srcdir) throws Exception {
		String [] hp = hostPort.split(":");
		String host = hp[0].trim();
		int port = (hp.length > 1) ? Integer.parseInt(hp[1].trim()) : 5577;
		String sshTarget = "elphel@" + host;                                      // DGX login user
		String code = "/home/elphel/c5p_dnn";                                     // DGX dir (model.py + runs/ live here)
		String wantL2 = (l2model == null) ? "" : l2model;                         // normalize (server reports "" for L1-only)
		if (canConnect(host, port, 1500)) {                                       // a server is listening - is it the right one?
			String [] cur = queryModels(host, port);
			if ((cur != null) && cur[0].equals(model) && cur[1].equals(wantL2)) {
				System.out.println("CuasDnnRemote.ensureServer(): server at " + host + ":" + port
						+ " already loaded with matching models (L1=" + model + ", L2=" + (wantL2.isEmpty()?"off":wantL2) + ") - reusing");
				return;
			}
			if (cur == null) {
				System.out.println("CuasDnnRemote.ensureServer(): server at " + host + ":" + port
						+ " up but did not report models (old build) - restarting for L1=" + model + ", L2=" + (wantL2.isEmpty()?"off":wantL2));
			} else {
				System.out.println("CuasDnnRemote.ensureServer(): server model MISMATCH at " + host + ":" + port
						+ " (has L1=" + cur[0] + ", L2=" + (cur[1].isEmpty()?"off":cur[1])
						+ "; want L1=" + model + ", L2=" + (wantL2.isEmpty()?"off":wantL2) + ") - restarting");
			}
			stopServer(sshTarget, code, port);                                    // free the port so the relaunch loads the new model
		}
		String run2 = !wantL2.isEmpty() ? (" RUN2=" + wantL2) : "";               // optional Layer-2. By Claude 06/22/2026
		System.out.println("CuasDnnRemote.ensureServer(): launching server on " + sshTarget
				+ " (model=" + model + ", l2=" + (run2.isEmpty()?"off":wantL2) + ")");
		deployScript("infer_server.py",     srcdir, sshTarget, code + "/infer_server.py");
		deployScript("layer2.py",            srcdir, sshTarget, code + "/layer2.py");   // Layer-2 model module (Layer2Net). By Claude 06/22/2026
		deployScript("run_infer_server.sh",  srcdir, sshTarget, code + "/run_infer_server.sh");
		ProcessBuilder pb = new ProcessBuilder("ssh", sshTarget,
				"cd " + code + " && chmod +x run_infer_server.sh && RUN=" + model + run2 + " PORT=" + port + " ./run_infer_server.sh start");
		pb.inheritIO();
		pb.start().waitFor();
		long deadline = System.currentTimeMillis() + 90000;                       // model load + warm-up can take a bit
		while (System.currentTimeMillis() < deadline) {
			if (canConnect(host, port, 1500)) { System.out.println("CuasDnnRemote.ensureServer(): server up at " + host + ":" + port); return; }
			Thread.sleep(2000);
		}
		throw new Exception("CuasDnnRemote.ensureServer(): server did not come up at " + host + ":" + port);
	}
}
