package invalid.lena.scrcpy;

import android.media.MediaCodec;
import android.media.MediaFormat;
import android.media.MediaMuxer;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;

// MP4 muxer that taps the H.264/H.265/AV1 elementary stream coming
// straight from the scrcpy server - no re-encode, the bytes go onto
// disk verbatim. AV1 muxing into MP4 needs Android 11+ (MediaMuxer
// learned the codec there); we are minSdk 31 so that's fine.
//
// State machine: IDLE → (arm) → ARMED → (next keyframe) → RECORDING
//                                                          → (stop) → IDLE
//
// Ordering:
//   onMeta(fcc, w, h)   from VideoStream session-meta packets.
//   onFrame(...)         every frame, in wire order. Config frames
//                        carry SPS/PPS; we cache the most recent one
//                        as csd-0 for MediaMuxer.addTrack.
//
// Resize during recording closes the file (addTrack-after-start is
// illegal). User can re-arm; a fresh output gets the new dimensions.
public final class MuxRecorder implements VideoRecorder {

    private enum State { IDLE, ARMED, RECORDING }

    private final Object lock = new Object();

    private State state = State.IDLE;
    private int   fourcc, w, h;
    private byte[] csd;

    private File         outFile;
    private MediaMuxer   muxer;
    private int          trackIdx = -1;
    private long         bytesWritten;
    private long         firstPtsUs = -1;
    private long         lastPtsUs  = -1;

    // Externally-driven controls.

    public void arm(File out) {
        synchronized (lock) {
            if (state != State.IDLE) {
                Log.w("rec: arm ignored in state %s", state);
                return;
            }
            outFile = out;
            state = State.ARMED;
            firstPtsUs = -1;
            lastPtsUs  = -1;
            bytesWritten = 0;
        }
        Log.i("rec: armed -> %s", out);
    }

    // Returns true if a recording was actually written (a keyframe landed
    // and the muxer ran); false if we were still ARMED, so the caller can
    // tell the user the truth instead of claiming a file was saved.
    public boolean stop() {
        long bytes;
        boolean wrote;
        synchronized (lock) {
            if (state == State.IDLE) return false;
            wrote = state == State.RECORDING && closeMuxerLocked();
            if (!wrote) deleteOutputLocked();
            bytes = bytesWritten;
            state = State.IDLE;
            outFile = null;
        }
        Log.i("rec: stopped (%d bytes)", bytes);
        return wrote;
    }

    public boolean isActive() {
        synchronized (lock) {
            return state == State.ARMED || state == State.RECORDING;
        }
    }

    // VideoRecorder callbacks.

    @Override
    public void onMeta(int fcc, int width, int height) {
        synchronized (lock) {
            // Resize during recording → close the file. addTrack after
            // start is illegal; we don't try to splice tracks together.
            // The user can re-arm; a fresh output gets the new dimensions.
            if (state == State.RECORDING && (fcc != fourcc || width != w || height != h)) {
                Log.w("rec: resize/codec change during recording - closing (%d bytes)",
                        bytesWritten);
                closeMuxerLocked();
                state = State.IDLE;
                outFile = null;
            }
            fourcc = fcc; w = width; h = height;
        }
    }

    @Override
    public void onFrame(byte[] data, long ptsUs, boolean isConfig, boolean isKeyframe) {
        if (isConfig) {
            // Cache the latest CSD (SPS/PPS for h264 - vendor packs as
            // Annex-B NAL units, exactly what MediaMuxer wants in csd-0
            // for AVC/HEVC/AV1).
            synchronized (lock) {
                csd = data.clone();
            }
            return;
        }
        synchronized (lock) {
            switch (state) {
                case IDLE:
                    return;
                case ARMED:
                    if (!isKeyframe || csd == null) return;
                    if (!startMuxerLocked()) {
                        failLocked("start");
                        return;
                    }
                    if (!writeSampleLocked(data, ptsUs, true)) {
                        failLocked("write");
                        return;
                    }
                    state = State.RECORDING;
                    return;
                case RECORDING:
                    if (!writeSampleLocked(data, ptsUs, isKeyframe)) failLocked("write");
            }
        }
    }

    @Override
    public void close() {
        stop();
    }

    // --- internals; all called with `lock` held ---

    private boolean startMuxerLocked() {
        String mime = mimeFor(fourcc);
        if (mime == null) {
            Log.e("rec: unsupported codec %s", Wire.fourccName(fourcc));
            return false;
        }
        try {
            muxer = new MediaMuxer(outFile.getAbsolutePath(),
                    MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
            MediaFormat fmt = MediaFormat.createVideoFormat(mime, w, h);
            fmt.setByteBuffer("csd-0", ByteBuffer.wrap(csd));
            trackIdx = muxer.addTrack(fmt);
            muxer.start();
            Log.i("rec: muxer start %s %dx%d -> %s", mime, w, h, outFile);
            return true;
        } catch (IOException | IllegalStateException e) {
            Log.e(e, "rec: muxer start failed");
            try { if (muxer != null) muxer.release(); } catch (Exception ignored) {}
            muxer = null;
            trackIdx = -1;
            return false;
        }
    }

    private boolean writeSampleLocked(byte[] data, long ptsUs, boolean isKeyframe) {
        if (muxer == null) return false;
        // PTS is rebased to zero so the file is self-contained; some
        // players choke on absolute PTS that doesn't start near zero.
        if (firstPtsUs < 0) firstPtsUs = ptsUs;
        long rebased = ptsUs - firstPtsUs;
        if (rebased < 0) rebased = 0;
        // Monotonic guard - MediaMuxer fails the write if pts goes back.
        if (rebased <= lastPtsUs) rebased = lastPtsUs + 1;
        lastPtsUs = rebased;

        try {
            ByteBuffer buf = ByteBuffer.wrap(data);
            MediaCodec.BufferInfo info = new MediaCodec.BufferInfo();
            info.set(0, data.length, rebased,
                    isKeyframe ? MediaCodec.BUFFER_FLAG_KEY_FRAME : 0);
            muxer.writeSampleData(trackIdx, buf, info);
            bytesWritten += data.length;
            return true;
        } catch (IllegalStateException e) {
            Log.w("rec: writeSampleData: %s", e);
            return false;
        }
    }

    private boolean closeMuxerLocked() {
        MediaMuxer m = muxer;
        muxer = null;
        trackIdx = -1;
        if (m == null) return false;
        boolean ok = true;
        try { m.stop(); } catch (Exception e) { ok = false; Log.w("rec: muxer stop: %s", e); }
        try { m.release(); } catch (Exception ignored) {}
        return ok;
    }

    private void failLocked(String operation) {
        Log.e("rec: %s failed, discarding output", operation);
        closeMuxerLocked();
        deleteOutputLocked();
        state = State.IDLE;
        outFile = null;
    }

    private void deleteOutputLocked() {
        if (outFile != null && outFile.exists() && !outFile.delete()) {
            Log.w("rec: could not delete incomplete output %s", outFile);
        }
    }

    private static String mimeFor(int fourcc) {
        switch (fourcc) {
            case Wire.CODEC_H264: return MediaFormat.MIMETYPE_VIDEO_AVC;
            case Wire.CODEC_H265: return MediaFormat.MIMETYPE_VIDEO_HEVC;
            case Wire.CODEC_AV1:  return MediaFormat.MIMETYPE_VIDEO_AV1;
            default: return null;
        }
    }
}
