package invalid.lena.scrcpy;

import android.media.AudioAttributes;
import android.media.AudioFormat;
import android.media.AudioManager;
import android.media.AudioTrack;
import android.media.MediaCodec;
import android.media.MediaFormat;
import android.os.Handler;
import android.os.HandlerThread;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayDeque;
import java.util.Deque;

// Audio output: AudioTrack writing 48 kHz stereo 16-bit PCM. The
// upstream feed is either raw PCM (passthrough) or Opus packets
// (decoded through MediaCodec into PCM first).
public final class AudioSink implements AudioFrames {

    private static final int SAMPLE_RATE = 48_000;
    private static final int CHANNEL_OUT = AudioFormat.CHANNEL_OUT_STEREO;
    private static final int ENCODING    = AudioFormat.ENCODING_PCM_16BIT;
    private static final int MAX_PENDING_OPUS = 16;

    // Defaults documented at <https://developer.android.com/reference/android/media/MediaCodec#CSD>.
    private static final long DEFAULT_PRE_ROLL_NS = 80_000_000L;

    private AudioTrack    track;
    private volatile MediaCodec opusCodec;
    private HandlerThread opusThread;
    private Handler       opusHandler;
    private boolean       opusConfigured;
    private int           fourcc;

    private final Object lock = new Object();
    private final Runnable onFatalError;
    private final Deque<Integer> freeOpusInputs = new ArrayDeque<>();
    private final Deque<byte[]> pendingOpus = new ArrayDeque<>(MAX_PENDING_OPUS);

    private volatile boolean released;
    public  volatile long    frames; // public read for the status overlay
    private long droppedBytes;
    private long lastDropLogMs;

    public AudioSink(Runnable onFatalError) {
        this.onFatalError = onFatalError;
    }

    @Override
    public void start(int fourcc) {
        this.fourcc = fourcc;

        int minBuf = AudioTrack.getMinBufferSize(SAMPLE_RATE, CHANNEL_OUT, ENCODING);
        if (minBuf <= 0) throw new IllegalStateException("AudioTrack.getMinBufferSize=" + minBuf);
        int bufSize = Math.max(minBuf, 32 * 1024);

        AudioAttributes attrs = new AudioAttributes.Builder()
                .setUsage(AudioAttributes.USAGE_MEDIA)
                .setContentType(AudioAttributes.CONTENT_TYPE_MUSIC)
                .build();
        AudioFormat fmt = new AudioFormat.Builder()
                .setSampleRate(SAMPLE_RATE)
                .setChannelMask(CHANNEL_OUT)
                .setEncoding(ENCODING)
                .build();

        track = new AudioTrack(
                attrs, fmt, bufSize, AudioTrack.MODE_STREAM, AudioManager.AUDIO_SESSION_ID_GENERATE);
        track.play();
        Log.i("audio sink: AudioTrack started sr=%d ch=2 buf=%d", SAMPLE_RATE, bufSize);

        if (fourcc == Wire.CODEC_OPUS) {
            opusThread = new HandlerThread("audio-mc");
            opusThread.start();
            opusHandler = new Handler(opusThread.getLooper());
            // Codec is created here, but only configured once the first
            // FLAG_CONFIG frame arrives with the OpusHead bytes.
            try {
                opusCodec = MediaCodec.createDecoderByType(MediaFormat.MIMETYPE_AUDIO_OPUS);
            } catch (IOException e) {
                throw new IllegalStateException("audio sink: no opus decoder", e);
            }
             opusCodec.setCallback(new MediaCodec.Callback() {
                 @Override public void onInputBufferAvailable(MediaCodec mc, int idx) {
                    synchronized (lock) {
                        if (released || mc != opusCodec) return;
                        byte[] packet = pendingOpus.pollFirst();
                        if (packet == null) freeOpusInputs.offerLast(idx);
                        else submitOpus(mc, idx, packet);
                    }
                }
                @Override public void onOutputBufferAvailable(MediaCodec mc, int idx,
                                                              MediaCodec.BufferInfo info) {
                    try {
                        ByteBuffer out = mc.getOutputBuffer(idx);
                        if (out != null && info.size > 0 && !released) {
                            byte[] pcm = new byte[info.size];
                            out.position(info.offset);
                            out.limit(info.offset + info.size);
                            out.get(pcm);
                            writePcm(pcm, 0, pcm.length);
                        }
                    } catch (IllegalStateException ignored) {
                    } finally {
                        try { mc.releaseOutputBuffer(idx, false); }
                        catch (IllegalStateException ignored) {}
                    }
                }
                @Override public void onError(MediaCodec mc, MediaCodec.CodecException e) {
                    Log.e(e, "audio sink: opus codec error");
                }
                @Override public void onOutputFormatChanged(MediaCodec mc, MediaFormat f) {
                    Log.i("audio sink: opus output format %s", f);
                }
            }, opusHandler);
        }
    }

    @Override
    public void feed(byte[] data, int off, int len, boolean isConfig) {
        if (released || len <= 0) return;
        if (fourcc == Wire.CODEC_OPUS) {
            feedOpus(data, off, len, isConfig);
        } else {
            // Raw PCM: passthrough.
            frames++;
            writePcm(data, off, len);
        }
    }

    private void feedOpus(byte[] data, int off, int len, boolean isConfig) {
        if (isConfig) {
            if (opusConfigured) return; // already configured
            // OpusHead is 19 bytes; pre_skip lives at bytes [10..11]. Reject
            // anything shorter before indexing into it.
            if (len < 19) {
                Log.e("audio sink: opus head too short (%d bytes)", len);
                return;
            }
            try {
                byte[] head = new byte[len];
                System.arraycopy(data, off, head, 0, len);
                int preSkipSamples = ((head[10] & 0xff) | ((head[11] & 0xff) << 8));
                long preSkipNs = preSkipSamples * 1_000_000_000L / SAMPLE_RATE;

                MediaFormat fmt = MediaFormat.createAudioFormat(
                        MediaFormat.MIMETYPE_AUDIO_OPUS, SAMPLE_RATE, 2);
                fmt.setByteBuffer("csd-0", ByteBuffer.wrap(head));
                fmt.setByteBuffer("csd-1", longLeBytes(preSkipNs));
                fmt.setByteBuffer("csd-2", longLeBytes(DEFAULT_PRE_ROLL_NS));
                opusCodec.configure(fmt, null, null, 0);
                opusCodec.start();
                opusConfigured = true;
                Log.i("audio sink: opus configured, pre_skip=%d ns", preSkipNs);
            } catch (Exception e) {
                Log.e(e, "audio sink: opus configure");
            }
            return;
        }
        if (!opusConfigured) return;
        byte[] packet = new byte[len];
        System.arraycopy(data, off, packet, 0, len);
        synchronized (lock) {
            if (released || opusCodec == null) return;
            Integer idx = freeOpusInputs.pollFirst();
            if (idx != null) {
                submitOpus(opusCodec, idx, packet);
                return;
            }
            if (pendingOpus.size() == MAX_PENDING_OPUS) pendingOpus.pollFirst();
            pendingOpus.offerLast(packet);
        }
    }

    // Must be called with lock held. Async MediaCodec input indices belong
    // to the codec instance that delivered them.
    private void submitOpus(MediaCodec codec, int idx, byte[] packet) {
        try {
            ByteBuffer in = codec.getInputBuffer(idx);
            if (in == null || packet.length > in.capacity()) {
                Log.e("audio sink: opus packet exceeds codec input (%d bytes)", packet.length);
                if (onFatalError != null) onFatalError.run();
                return;
            }
            in.clear();
            in.put(packet);
            codec.queueInputBuffer(idx, 0, packet.length, 0, 0);
            frames++;
        } catch (IllegalStateException e) {
            Log.w("audio sink: opus queueInputBuffer: %s", e);
        }
    }

    // Encode a long little-endian for MediaFormat csd-1 / csd-2.
    private static ByteBuffer longLeBytes(long v) {
        ByteBuffer b = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(v);
        b.flip();
        return b;
    }

    private void writePcm(byte[] data, int off, int len) {
        AudioTrack t = track;
        if (released || t == null || len <= 0) return;
        int written = t.write(data, off, len, AudioTrack.WRITE_NON_BLOCKING);
        if (written < 0) {
            Log.w("audio sink: write rc=%d", written);
            return;
        }
        if (written < len) {
            droppedBytes += (len - written);
            long now = System.currentTimeMillis();
            if (now - lastDropLogMs > 1000L) {
                lastDropLogMs = now;
                Log.w("audio sink: ring full, dropped %d bytes so far", droppedBytes);
            }
        }
    }

    @Override
    public void release() {
        MediaCodec c;
        synchronized (lock) {
            released = true;
            c = opusCodec;
            opusCodec = null;
            freeOpusInputs.clear();
            pendingOpus.clear();
        }
        HandlerThread ht = opusThread;
        opusThread = null;
        opusHandler = null;
        if (c != null) {
            try { c.stop(); } catch (Exception ignored) {}
            try { c.release(); } catch (Exception ignored) {}
        }
        if (ht != null) ht.quitSafely();

        AudioTrack t = track;
        track = null;
        if (t == null) return;
        try { t.pause(); t.flush(); t.stop(); } catch (Exception ignored) {}
        try { t.release(); } catch (Exception ignored) {}
    }
}