package org.expeditee.warp;

import org.eclipse.jetty.server.*;
import org.eclipse.jetty.server.handler.ContextHandler;
import org.eclipse.jetty.util.ssl.SslContextFactory;
import org.eclipse.jetty.websocket.server.WebSocketUpgradeHandler;
import org.eclipse.jetty.websocket.api.Session;
    
import org.eclipse.jetty.websocket.api.exceptions.WebSocketException;
import org.eclipse.jetty.websocket.api.Callback;
    
    
import java.awt.*;
import java.io.IOException;
import java.net.URI;
    
import java.security.SecureRandom;
import java.time.Duration;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicBoolean;

import java.nio.ByteBuffer;
import java.nio.file.Paths;

import org.eclipse.jetty.http.HttpHeader;
import org.eclipse.jetty.server.Request;
import org.eclipse.jetty.server.Response;
import org.eclipse.jetty.server.handler.AbstractHandler;

import org.eclipse.jetty.server.handler.ResourceHandler;
import org.eclipse.jetty.util.resource.ResourceFactory;


public class LocalWarpHelper
{

    static String randomToken(int bytes) {
        byte[] b = new byte[bytes];
        new SecureRandom().nextBytes(b);
        return Base64.getUrlEncoder().withoutPadding().encodeToString(b);
    }


    public static class KeystoreMaterial {
        final java.nio.file.Path pkcs12Path;
        final String password;

        KeystoreMaterial(java.nio.file.Path pkcs12Path, String password) {
            this.pkcs12Path = pkcs12Path;
            this.password = password;
        }
    }

    /**
     * First-run keystore generation. Produces a self-signed cert suitable for "localhost" with SANs.
     */
    public static class KeystoreUtil {

        static KeystoreMaterial ensureLocalhostPkcs12() throws Exception {
            java.nio.file.Path dir = java.nio.file.Paths.get(System.getProperty("user.home"), ".expeditee-warp-helper");
            java.nio.file.Files.createDirectories(dir);

            java.nio.file.Path p12 = dir.resolve("localhost.p12");
            java.nio.file.Path passFile = dir.resolve("keystore.pass");

            if (java.nio.file.Files.exists(p12) && java.nio.file.Files.exists(passFile)) {
                String password = java.nio.file.Files.readString(passFile, java.nio.charset.StandardCharsets.UTF_8).trim();
                return new KeystoreMaterial(p12, password);
            }

            String password = randomPassword(24);
            java.nio.file.Files.writeString(passFile, password + "\n",
                java.nio.charset.StandardCharsets.UTF_8,
                java.nio.file.StandardOpenOption.CREATE,
                java.nio.file.StandardOpenOption.TRUNCATE_EXISTING);

            // Generate PKCS12 with self-signed cert that includes SAN localhost + 127.0.0.1
            SelfSignedPkcs12Generator.generate(p12, password);

            return new KeystoreMaterial(p12, password);
        }

        static String randomPassword(int bytes) {
            byte[] b = new byte[bytes];
            new SecureRandom().nextBytes(b);
            return Base64.getUrlEncoder().withoutPadding().encodeToString(b);
        }
    }



    /**
     * Resolve a port from an environment variable.
     * If unset or invalid, returns the provided default.
     */
    static int resolvePortFromEnv(String envVar, int defaultPort) {
        String raw = System.getenv(envVar);
        if (raw == null || raw.isBlank()) {
            return defaultPort;
        }
        try {
            int p = Integer.parseInt(raw.trim());
            if (p < 1 || p > 65535) {
                System.err.println(envVar + " out of range (1..65535): " + raw + "  Using default " + defaultPort);
                return defaultPort;
            }
            return p;
        } catch (NumberFormatException nfe) {
            System.err.println(envVar + " not an integer: " + raw + "  Using default " + defaultPort);
            return defaultPort;
        }
    }


    public static void main(String[] args) throws Exception {
        int port = resolvePortFromEnv("EXPEDITEE_PORT", 8433);
        String token = randomToken(32);

        KeystoreMaterial ks = KeystoreUtil.ensureLocalhostPkcs12();

        Server server = new Server();

        // ---- TLS connector bound to 127.0.0.1 ----
        SslContextFactory.Server ssl = new SslContextFactory.Server();
        ssl.setKeyStorePath(ks.pkcs12Path.toAbsolutePath().toString());
        ssl.setKeyStorePassword(ks.password);
        ssl.setKeyStoreType("PKCS12");
        ssl.setIncludeProtocols("TLSv1.3", "TLSv1.2");

        HttpConfiguration https = new HttpConfiguration();
        https.addCustomizer(new SecureRequestCustomizer());

        ServerConnector connector = new ServerConnector(
            server,
            new SslConnectionFactory(ssl, "http/1.1"),
            new HttpConnectionFactory(https)
        );
        connector.setHost("127.0.0.1");
        connector.setPort(port);
        server.addConnector(connector);
		
	ContextHandler ctx = new ContextHandler("/");
	
	WebSocketUpgradeHandler wsHandler = WebSocketUpgradeHandler.from(server, ctx, container -> {
		container.setMaxTextMessageSize(64 * 1024);
		container.setIdleTimeout(Duration.ofMinutes(10));
		
		container.addMapping("/ws", (rq, rs, cb) -> new WarpSocket(token));
	    });
	
	RootHandler root = new RootHandler();


	
	ResourceHandler resources = new ResourceHandler();
	resources.setDirAllowed(false);
	
	resources.setBaseResource(
	   ResourceFactory.of(resources).newResource(Paths.get("webroot"))
	);
	
	Handler.Sequence handlers = new Handler.Sequence();
	handlers.addHandler(wsHandler);
	handlers.addHandler(root);
	handlers.addHandler(resources);
	
	ctx.setHandler(handlers);
	server.setHandler(ctx);
	
        server.start();

        // ---- Launch browser (token in fragment so it isn't sent to the remote server) ----
        // Replace with your real WebSwing URL when ready.
        //String webSwingUrl = "https://your-server.example/expeditee/";
	String webSwingUrl = "https://localhost:" + port + "/";
        String launchUrl = webSwingUrl + "?warpPort=" + port + "&warpToken=" + token;

        System.out.println("Local Warp Helper running.");
        System.out.println("  Local WS:  wss://127.0.0.1:" + port + "/ws");
        System.out.println("  Token:     " + token);
        System.out.println("Open in browser:  " + launchUrl);

        if (Desktop.isDesktopSupported()) {
            Desktop.getDesktop().browse(new URI(launchUrl));
        }

        server.join();
    }   
}