package com.foxx.androidcast.discovery;

import android.content.Context;
import android.net.wifi.WifiManager;

import com.foxx.androidcast.CastConfig;
import com.foxx.androidcast.network.CastProtoHeader;

import org.json.JSONObject;

import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/** Quick UDP cast discovery scan for developer network self-test. */
public final class CastLanPeerProbe {
    private static final int SCAN_MS = 2_500;
    private static final int PROBE_INTERVAL_MS = 700;

    public static final class Result {
        public final int peers;
        public final String summary;
        public final long elapsedMs;

        Result(int peers, String summary, long elapsedMs) {
            this.peers = peers;
            this.summary = summary;
            this.elapsedMs = elapsedMs;
        }
    }

    private CastLanPeerProbe() {}

    public static Result scan(Context context, AtomicBoolean cancel) {
        long t0 = System.currentTimeMillis();
        Map<String, String> peers = new LinkedHashMap<>();
        WifiManager.MulticastLock lock = acquireMulticastLock(context);
        try (DatagramSocket sock = openSocket()) {
            if (sock == null) {
                return new Result(0, "socket bind failed", System.currentTimeMillis() - t0);
            }
            sock.setSoTimeout(250);
            byte[] buf = new byte[2048];
            long deadline = System.currentTimeMillis() + SCAN_MS;
            long nextProbe = System.currentTimeMillis();
            while (System.currentTimeMillis() < deadline) {
                if (cancel != null && cancel.get()) {
                    break;
                }
                long now = System.currentTimeMillis();
                if (now >= nextProbe) {
                    sendDiscoverProbe(sock);
                    nextProbe = now + PROBE_INTERVAL_MS;
                }
                DatagramPacket packet = new DatagramPacket(buf, buf.length);
                try {
                    sock.receive(packet);
                    recordPeer(context, packet, peers);
                } catch (java.net.SocketTimeoutException ignored) {
                } catch (Exception ignored) {
                }
            }
        } catch (Exception e) {
            long ms = System.currentTimeMillis() - t0;
            return new Result(0, e.getClass().getSimpleName(), ms);
        } finally {
            releaseMulticastLock(lock);
        }
        long ms = System.currentTimeMillis() - t0;
        if (peers.isEmpty()) {
            return new Result(0, "UDP :" + CastConfig.DISCOVERY_PORT + " · " + ms + " ms", ms);
        }
        StringBuilder names = new StringBuilder();
        for (Map.Entry<String, String> e : peers.entrySet()) {
            if (names.length() > 0) {
                names.append(", ");
            }
            names.append(e.getValue()).append('@').append(e.getKey());
            if (names.length() > 96) {
                names.append("…");
                break;
            }
        }
        return new Result(peers.size(), names.toString(), ms);
    }

    private static DatagramSocket openSocket() {
        try {
            DatagramSocket sock = new DatagramSocket(null);
            sock.setReuseAddress(true);
            sock.setBroadcast(true);
            try {
                sock.bind(new InetSocketAddress(CastConfig.DISCOVERY_PORT));
            } catch (Exception ignored) {
                sock.bind(null);
            }
            return sock;
        } catch (Exception e) {
            return null;
        }
    }

    private static void sendDiscoverProbe(DatagramSocket sock) {
        try {
            JSONObject json = new JSONObject();
            json.put("magic", CastConfig.DISCOVERY_MAGIC);
            json.put("discover", true);
            byte[] bytes = json.toString().getBytes(StandardCharsets.UTF_8);
            InetAddress broadcast = InetAddress.getByName("255.255.255.255");
            DatagramPacket packet = new DatagramPacket(bytes, bytes.length, broadcast, CastConfig.DISCOVERY_PORT);
            sock.send(packet);
        } catch (Exception ignored) {
        }
    }

    private static void recordPeer(Context context, DatagramPacket packet, Map<String, String> peers) {
        try {
            JSONObject json = parseAnnounceJson(packet);
            if (json == null) {
                return;
            }
            if (!CastConfig.DISCOVERY_MAGIC.equals(json.optString("magic"))) {
                return;
            }
            if (json.optBoolean("discover", false)) {
                return;
            }
            String host = packet.getAddress().getHostAddress();
            String remoteUuid = json.optString("device_uuid", "");
            if (DiscoverySelfFilter.isSelf(context, host, remoteUuid)) {
                return;
            }
            String name = json.optString("name", "Unknown");
            int port = json.optInt("port", CastConfig.CAST_PORT);
            boolean listening = json.optBoolean("listening", true);
            if (!listening) {
                return;
            }
            peers.put(host + ":" + port, name);
        } catch (Exception ignored) {
        }
    }

    private static JSONObject parseAnnounceJson(DatagramPacket packet) {
        try {
            byte[] raw = packet.getData();
            int off = packet.getOffset();
            int len = packet.getLength();
            String text;
            if (len >= CastProtoHeader.HEADER_SIZE_V1 && readMagicLe(raw, off) == CastProtoHeader.MAGIC) {
                CastProtoHeader hdr = CastProtoHeader.decode(raw, off, len);
                int bodyOff = off + hdr.headerSize;
                int bodyLen = len - hdr.headerSize;
                if (bodyLen > 5) {
                    byte type = raw[bodyOff];
                    int plen = readIntBe(raw, bodyOff + 1);
                    if (plen >= 0 && bodyOff + 5 + plen <= off + len) {
                        text = new String(raw, bodyOff + 5, plen, StandardCharsets.UTF_8);
                    } else {
                        text = new String(raw, bodyOff, bodyLen, StandardCharsets.UTF_8);
                    }
                } else {
                    text = new String(raw, bodyOff, bodyLen, StandardCharsets.UTF_8);
                }
            } else {
                text = new String(raw, off, len, StandardCharsets.UTF_8);
            }
            return new JSONObject(text);
        } catch (Exception e) {
            return null;
        }
    }

    private static int readMagicLe(byte[] b, int off) {
        return (b[off] & 0xff) | ((b[off + 1] & 0xff) << 8)
                | ((b[off + 2] & 0xff) << 16) | ((b[off + 3] & 0xff) << 24);
    }

    private static int readIntBe(byte[] b, int off) {
        return ((b[off] & 0xff) << 24) | ((b[off + 1] & 0xff) << 16)
                | ((b[off + 2] & 0xff) << 8) | (b[off + 3] & 0xff);
    }

    private static WifiManager.MulticastLock acquireMulticastLock(Context context) {
        if (context == null) {
            return null;
        }
        WifiManager wifi = context.getApplicationContext().getSystemService(WifiManager.class);
        if (wifi == null) {
            return null;
        }
        WifiManager.MulticastLock lock = wifi.createMulticastLock("androidcast-lan-probe");
        lock.setReferenceCounted(false);
        lock.acquire();
        return lock;
    }

    private static void releaseMulticastLock(WifiManager.MulticastLock lock) {
        if (lock == null) {
            return;
        }
        try {
            if (lock.isHeld()) {
                lock.release();
            }
        } catch (RuntimeException ignored) {
        }
    }
}
