import java.io.*;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.*;
import java.util.concurrent.*;
import java.util.regex.*;

/**
 * P2P聊天程序 - Web版本服务器
 * 支持 WebSocket 连接和 HTTP 静态文件服务
 */
public class P2PChatWebServer {
    private static final int DEFAULT_PORT = 8888;
    private static final int HTTP_PORT = 8080;
    private ServerSocket serverSocket;
    private ServerSocket httpSocket;
    private Map<String, WebSocketClient> connectedPeers = new ConcurrentHashMap<>();
    private ExecutorService threadPool = Executors.newCachedThreadPool();
    private volatile boolean running = true;
    private String username;
    
    public P2PChatWebServer(String username, int port, int httpPort) {
        this.username = username;
        try {
            serverSocket = new ServerSocket(port);
            httpSocket = new ServerSocket(httpPort);
            System.out.println("=== P2P聊天程序 Web 服务器启动成功 ===");
            System.out.println("用户名: " + username);
            System.out.println("WebSocket 端口: " + port);
            System.out.println("HTTP 端口: " + httpPort);
            System.out.println("本机IP: " + InetAddress.getLocalHost().getHostAddress());
            System.out.println("访问地址: http://localhost:" + httpPort);
            System.out.println("=========================================");
        } catch (IOException e) {
            System.err.println("启动失败: " + e.getMessage());
            System.exit(1);
        }
    }
    
    // 启动 WebSocket 服务器
    public void startWebSocketServer() {
        threadPool.execute(() -> {
            while (running) {
                try {
                    Socket clientSocket = serverSocket.accept();
                    WebSocketClient client = new WebSocketClient(clientSocket);
                    threadPool.execute(client);
                } catch (IOException e) {
                    if (running) {
                        System.err.println("接受连接失败: " + e.getMessage());
                    }
                }
            }
        });
    }
    
    // 启动 HTTP 服务器
    public void startHttpServer() {
        threadPool.execute(() -> {
            while (running) {
                try {
                    Socket clientSocket = httpSocket.accept();
                    threadPool.execute(() -> handleHttpRequest(clientSocket));
                } catch (IOException e) {
                    if (running) {
                        System.err.println("HTTP 请求处理失败: " + e.getMessage());
                    }
                }
            }
        });
    }
    
    // 处理 HTTP 请求
    private void handleHttpRequest(Socket socket) {
        try {
            BufferedReader reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
            String line = reader.readLine();
            
            if (line != null && line.startsWith("GET")) {
                String[] parts = line.split(" ");
                String path = parts[1];
                
                if (path.equals("/")) {
                    path = "/index.html";
                }
                
                sendHttpResponse(socket, path);
            }
            
            socket.close();
        } catch (IOException e) {
            System.err.println("HTTP 请求处理错误: " + e.getMessage());
        }
    }
    
    // 发送 HTTP 响应
    private void sendHttpResponse(Socket socket, String path) throws IOException {
        OutputStream out = socket.getOutputStream();
        
        String content = getFileContent(path);
        String contentType = getContentType(path);
        
        if (content != null) {
            String response = "HTTP/1.1 200 OK\r\n" +
                            "Content-Type: " + contentType + "; charset=UTF-8\r\n" +
                            "Content-Length: " + content.getBytes(StandardCharsets.UTF_8).length + "\r\n" +
                            "Connection: close\r\n\r\n" +
                            content;
            out.write(response.getBytes(StandardCharsets.UTF_8));
        } else {
            String response = "HTTP/1.1 404 Not Found\r\n" +
                            "Content-Type: text/html; charset=UTF-8\r\n" +
                            "Connection: close\r\n\r\n" +
                            "<h1>404 Not Found</h1>";
            out.write(response.getBytes(StandardCharsets.UTF_8));
        }
        
        out.flush();
    }
    
    // 获取文件内容
    private String getFileContent(String path) {
        try {
            File file = new File("web" + path);
            if (file.exists() && file.isFile()) {
                return new String(java.nio.file.Files.readAllBytes(file.toPath()), StandardCharsets.UTF_8);
            }
        } catch (IOException e) {
            System.err.println("读取文件失败: " + e.getMessage());
        }
        return null;
    }
    
    // 获取内容类型
    private String getContentType(String path) {
        if (path.endsWith(".html")) return "text/html";
        if (path.endsWith(".css")) return "text/css";
        if (path.endsWith(".js")) return "application/javascript";
        if (path.endsWith(".json")) return "application/json";
        return "text/plain";
    }
    
    // 广播消息
    public void broadcastMessage(String message, WebSocketClient sender) {
        String fullMessage = "[" + sender.peerId + "]: " + message;
        for (WebSocketClient client : connectedPeers.values()) {
            if (client != sender) {
                client.sendMessage(fullMessage);
            }
        }
    }
    
    // 列出所有对等端
    public String listPeers() {
        StringBuilder sb = new StringBuilder();
        for (String peerId : connectedPeers.keySet()) {
            sb.append(peerId).append(",");
        }
        return sb.toString();
    }
    
    // WebSocket 客户端处理器
    class WebSocketClient implements Runnable {
        private Socket socket;
        private InputStream in;
        private OutputStream out;
        private String peerId;
        private boolean handshakeDone = false;
        
        public WebSocketClient(Socket socket) {
            this.socket = socket;
            try {
                in = socket.getInputStream();
                out = socket.getOutputStream();
            } catch (IOException e) {
                System.err.println("初始化连接失败: " + e.getMessage());
            }
        }
        
        @Override
        public void run() {
            try {
                if (!performHandshake()) {
                    socket.close();
                    return;
                }
                
                while (true) {
                    String message = receiveMessage();
                    if (message == null) break;
                    
                    handleMessage(message);
                }
            } catch (IOException e) {
                System.out.println("连接断开: " + (peerId != null ? peerId : "未知"));
            } finally {
                close();
            }
        }
        
        // WebSocket 握手
        private boolean performHandshake() throws IOException {
            BufferedReader reader = new BufferedReader(new InputStreamReader(in));
            String line;
            String key = null;
            
            while ((line = reader.readLine()) != null && !line.isEmpty()) {
                if (line.startsWith("Sec-WebSocket-Key:")) {
                    key = line.substring(19).trim();
                }
            }
            
            if (key == null) return false;
            
            String accept = generateAcceptKey(key);
            String response = "HTTP/1.1 101 Switching Protocols\r\n" +
                            "Upgrade: websocket\r\n" +
                            "Connection: Upgrade\r\n" +
                            "Sec-WebSocket-Accept: " + accept + "\r\n\r\n";
            
            out.write(response.getBytes(StandardCharsets.UTF_8));
            out.flush();
            
            handshakeDone = true;
            return true;
        }
        
        // 生成 WebSocket Accept Key
        private String generateAcceptKey(String key) {
            try {
                String magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                MessageDigest md = MessageDigest.getInstance("SHA-1");
                byte[] hash = md.digest((key + magic).getBytes(StandardCharsets.UTF_8));
                return Base64.getEncoder().encodeToString(hash);
            } catch (Exception e) {
                return "";
            }
        }
        
        // 接收 WebSocket 消息
        private String receiveMessage() throws IOException {
            int b = in.read();
            if (b == -1) return null;
            
            boolean fin = (b & 0x80) != 0;
            int opcode = b & 0x0F;
            
            if (opcode == 8) return null; // Close frame
            
            b = in.read();
            if (b == -1) return null;
            
            boolean masked = (b & 0x80) != 0;
            long payloadLen = b & 0x7F;
            
            if (payloadLen == 126) {
                payloadLen = (in.read() << 8) | in.read();
            } else if (payloadLen == 127) {
                payloadLen = 0;
                for (int i = 0; i < 8; i++) {
                    payloadLen = (payloadLen << 8) | in.read();
                }
            }
            
            byte[] maskingKey = new byte[4];
            if (masked) {
                in.read(maskingKey);
            }
            
            byte[] payload = new byte[(int) payloadLen];
            in.read(payload);
            
            if (masked) {
                for (int i = 0; i < payload.length; i++) {
                    payload[i] ^= maskingKey[i % 4];
                }
            }
            
            return new String(payload, StandardCharsets.UTF_8);
        }
        
        // 发送 WebSocket 消息
        public void sendMessage(String message) {
            try {
                byte[] payload = message.getBytes(StandardCharsets.UTF_8);
                
                out.write(0x81); // FIN + Text frame
                
                if (payload.length <= 125) {
                    out.write(payload.length);
                } else if (payload.length <= 65535) {
                    out.write(126);
                    out.write(payload.length >> 8);
                    out.write(payload.length & 0xFF);
                } else {
                    out.write(127);
                    for (int i = 7; i >= 0; i--) {
                        out.write((int) (payload.length >> (i * 8)) & 0xFF);
                    }
                }
                
                out.write(payload);
                out.flush();
            } catch (IOException e) {
                System.err.println("发送消息失败: " + e.getMessage());
            }
        }
        
        // 处理消息
        private void handleMessage(String message) {
            if (message.startsWith("USERNAME:")) {
                peerId = message.substring(9);
                connectedPeers.put(peerId, this);
                System.out.println(peerId + " 已连接");
                sendMessage("SYSTEM:欢迎 " + peerId + "！");
                sendMessage("PEERS:" + listPeers());
            } else if (message.startsWith("MESSAGE:")) {
                String msg = message.substring(8);
                System.out.println("[" + peerId + "]: " + msg);
                broadcastMessage(msg, this);
            }
        }
        
        public void close() {
            try {
                if (peerId != null) {
                    connectedPeers.remove(peerId);
                }
                socket.close();
            } catch (IOException e) {
                // 忽略
            }
        }
    }
    
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        
        System.out.print("请输入用户名: ");
        String username = scanner.nextLine();
        
        System.out.print("请输入 WebSocket 端口 (默认8888): ");
        String portInput = scanner.nextLine();
        int port = portInput.isEmpty() ? DEFAULT_PORT : Integer.parseInt(portInput);
        
        System.out.print("请输入 HTTP 端口 (默认8080): ");
        String httpPortInput = scanner.nextLine();
        int httpPort = httpPortInput.isEmpty() ? HTTP_PORT : Integer.parseInt(httpPortInput);
        
        P2PChatWebServer server = new P2PChatWebServer(username, port, httpPort);
        server.startWebSocketServer();
        server.startHttpServer();
        
        System.out.println("\n服务器运行中... 输入 'quit' 退出");
        
        while (true) {
            String input = scanner.nextLine();
            if (input.equalsIgnoreCase("quit")) {
                scanner.close();
                System.exit(0);
            }
        }
    }
}