package com.sw.dualscreen.socket;

import android.annotation.SuppressLint;

import org.json.JSONException;
import org.json.JSONObject;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 局域网多客户端通信管理器：支持并发、心跳、认证、广播、点对点发送
 * 可用于 Android 和 JVM 程序。
 */
public class LanCommunicationManager {

    // ============ 监听配置 ============
    private final int port;
    private final long HEARTBEAT_TIMEOUT_MS;
    private final int MAX_CLIENT_THREADS;

    // ============ 状态 ============
    private volatile boolean running = false;
    private ServerSocket serverSocket;

    // ============ 线程池 ============
    private final ExecutorService acceptExecutor = Executors.newSingleThreadExecutor();
    private final ExecutorService clientExecutor;
    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

    // ============ 客户端会话 ============
    private final ConcurrentHashMap<String, ClientSession> clients = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<Socket, ClientSession> unAuthSessions = new ConcurrentHashMap<>();

    // ============ 回调接口 ============
    public interface Listener {
        // 新客户端完成 AUTH / 认证
        void onClientConnected(String clientId);

        // 客户端断开
        void onClientDisconnected(String clientId);

        // 收到业务消息（type != heartbeat/auth)
        void onMessageReceived(String clientId, JSONObject message);
    }

    private Listener listener;

    // ============ 构造 ============
    public LanCommunicationManager(int port, int maxClientThreads, long heartbeatTimeoutMs) {
        this.port = port;
        this.MAX_CLIENT_THREADS = maxClientThreads;
        this.HEARTBEAT_TIMEOUT_MS = heartbeatTimeoutMs;
        this.clientExecutor = Executors.newFixedThreadPool(Math.max(2, maxClientThreads));
    }

    public void setListener(Listener listener) {
        this.listener = listener;
    }

    // ============ 启动服务端 ============
    @SuppressLint("DiscouragedApi")
    public void start() throws IOException {
        if (running) return;
        running = true;

        serverSocket = new ServerSocket(port);
        serverSocket.setSoTimeout(2000);

        acceptExecutor.execute(this::acceptLoop);
        scheduler.scheduleAtFixedRate(this::heartbeatCheck,
                HEARTBEAT_TIMEOUT_MS,
                HEARTBEAT_TIMEOUT_MS,
                TimeUnit.MILLISECONDS);

        System.out.println("LanCommunicationManager started on port " + port);
    }

    private void acceptLoop() {
        while (running) {
            try {
                Socket socket = serverSocket.accept();
                socket.setSoTimeout((int) HEARTBEAT_TIMEOUT_MS * 2);

                ClientSession session = new ClientSession(socket);
                unAuthSessions.put(socket, session);

                clientExecutor.execute(() -> clientReadLoop(session));

            } catch (SocketTimeoutException ignore) {
            } catch (Exception e) {
                if (running) e.printStackTrace();
            }
        }
    }

    // ============ 处理客户端数据读取 ============
    private void clientReadLoop(ClientSession session) {
        Socket socket = session.socket;

        try (DataInputStream in = new DataInputStream(socket.getInputStream())) {
            while (running && !socket.isClosed()) {

                int len;
                try {
                    len = in.readInt();
                } catch (SocketTimeoutException ste) {
                    continue;
                }

                if (len <= 0 || len > 10 * 1024 * 1024) break;

                byte[] buf = new byte[len];
                in.readFully(buf);
                session.updateLastSeen();

                JSONObject msg = new JSONObject(new String(buf));
                handleMessage(session, msg);
            }
        } catch (Exception ignored) {
        } finally {
            closeSession(session);
        }
    }

    private void handleMessage(ClientSession session, JSONObject msg) {
        String type = msg.optString("type", "");

        switch (type) {
            case "auth":
                handleAuth(session, msg);
                break;
            case "heartbeat":
                session.updateLastSeen();
                break;
            default:
                if (listener != null && session.clientId != null) {
                    listener.onMessageReceived(session.clientId, msg);
                }
                break;
        }
    }

    private void handleAuth(ClientSession session, JSONObject msg) {
        String clientId = msg.optString("clientId", null);
        if (clientId == null) return;

        session.clientId = clientId;

        // 移动到已认证 map
        unAuthSessions.remove(session.socket);
        clients.put(clientId, session);

        if (listener != null) listener.onClientConnected(clientId);

        sendToSession(session, ack("auth_ok"));
    }

    // ============ 心跳超时 ============
    private void heartbeatCheck() {
        long now = System.currentTimeMillis();

        for (Map.Entry<String, ClientSession> e : clients.entrySet()) {
            ClientSession s = e.getValue();
            if (now - s.lastSeen > HEARTBEAT_TIMEOUT_MS) {
                closeSession(s);
            }
        }

        for (ClientSession s : unAuthSessions.values()) {
            if (now - s.lastSeen > HEARTBEAT_TIMEOUT_MS * 2) {
                closeSession(s);
            }
        }
    }

    // ============ 发送 ============
    public boolean sendToClient(String clientId, JSONObject json) {
        ClientSession s = clients.get(clientId);
        return s != null && sendToSession(s, json);
    }

    public void broadcast(JSONObject json) {
        for (ClientSession s : clients.values()) {
            sendToSession(s, json);
        }
    }

    private boolean sendToSession(ClientSession s, JSONObject json) {
        try {
            DataOutputStream out = s.out;

            synchronized (out) {
                byte[] data = json.toString().getBytes();
                out.writeInt(data.length);
                out.write(data);
                out.flush();
            }
            return true;

        } catch (Exception e) {
            closeSession(s);
            return false;
        }
    }

    // ============ ACK ============
    private JSONObject ack(String type) {
        JSONObject j = new JSONObject();
        try {
            j.put("type", "ack");
            j.put("ack", type);
        } catch (JSONException e) {
            throw new RuntimeException(e);
        }

        return j;
    }

    // ============ 停止 ============
    public void stop() {
        running = false;

        try {
            serverSocket.close();
        } catch (Exception ignored) {
        }

        for (ClientSession s : clients.values()) closeSession(s);
        for (ClientSession s : unAuthSessions.values()) closeSession(s);

        acceptExecutor.shutdownNow();
        clientExecutor.shutdownNow();
        scheduler.shutdownNow();

        System.out.println("LanCommunicationManager stopped");
    }

    // ============ 会话类 ============
    public static class ClientSession {
        public final Socket socket;
        public final DataOutputStream out;
        public volatile long lastSeen = System.currentTimeMillis();
        public volatile String clientId;

        public ClientSession(Socket socket) throws IOException {
            this.socket = socket;
            this.out = new DataOutputStream(socket.getOutputStream());
        }

        public void updateLastSeen() {
            lastSeen = System.currentTimeMillis();
        }
    }

    public void closeSession(ClientSession session) {
        if (session == null) return;

        try {
            Socket socket = session.socket;

            // 1. 从已认证表移除
            if (session.clientId != null) {
                ClientSession removed = clients.remove(session.clientId);
                if (removed != null && listener != null) {
                    listener.onClientDisconnected(session.clientId);
                }
            }

            // 2. 从未认证表移除
            unAuthSessions.remove(socket);

            // 3. 关闭输出流
            try {
                session.out.close();
            } catch (Exception ignored) {
            }

            // 4. 关闭 socket
            try {
                if (!socket.isClosed()) socket.close();
            } catch (Exception ignored) {
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}