import com.sun.net.httpserver.Headers;
import com.sun.net.httpserver.HttpServer;

import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;

public class KshitijAIProxyServer {

    private static final int PORT = 7860;

    // Secrets (set in HF Space settings)
    private static final String BACKEND_URL = getenvOrThrow("BACKEND_URL"); // e.g. https://user-kshitijai-backend.hf.space
    private static final String HF_TOKEN = getenvOrThrow("HF_TOKEN");       // token that can access private space

    public static void main(String[] args) throws Exception {
        final Path baseDir = Paths.get(System.getProperty("user.dir")).toRealPath();

        HttpServer server = HttpServer.create(new InetSocketAddress(PORT), 0);

        server.createContext("/health", exchange -> {
            addCors(exchange.getResponseHeaders());
            if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
                exchange.sendResponseHeaders(405, -1);
                return;
            }
            byte[] ok = "ok".getBytes(StandardCharsets.UTF_8);
            exchange.getResponseHeaders().set("Content-Type", "text/plain; charset=utf-8");
            exchange.sendResponseHeaders(200, ok.length);
            try (OutputStream os = exchange.getResponseBody()) { os.write(ok); }
        });

        // Proxy: browser -> public space -> private backend space
        server.createContext("/api/chat", exchange -> {
            Headers h = exchange.getResponseHeaders();
            addCors(h);

            // Handle preflight
            if (exchange.getRequestMethod().equalsIgnoreCase("OPTIONS")) {
                exchange.sendResponseHeaders(204, -1);
                return;
            }

            if (!exchange.getRequestMethod().equalsIgnoreCase("POST")) {
                exchange.sendResponseHeaders(405, -1);
                return;
            }

            byte[] requestBytes = exchange.getRequestBody().readAllBytes();

            HttpClient client = HttpClient.newBuilder()
                    .connectTimeout(Duration.ofSeconds(10))
                    .build();

            String target = BACKEND_URL.endsWith("/")
                    ? BACKEND_URL + "api/chat"
                    : BACKEND_URL + "/api/chat";

            HttpRequest req = HttpRequest.newBuilder()
                    .uri(URI.create(target))
                    .timeout(Duration.ofSeconds(300))
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + HF_TOKEN)
                    .POST(HttpRequest.BodyPublishers.ofByteArray(requestBytes))
                    .build();

            HttpResponse<byte[]> resp;
            try {
                resp = client.send(req, HttpResponse.BodyHandlers.ofByteArray());
            } catch (Exception e) {
                byte[] err = ("Failed to reach backend at " + target + "\n" + e)
                        .getBytes(StandardCharsets.UTF_8);
                h.set("Content-Type", "text/plain; charset=utf-8");
                exchange.sendResponseHeaders(502, err.length);
                try (OutputStream os = exchange.getResponseBody()) { os.write(err); }
                return;
            }

            h.set("Content-Type", "application/json; charset=utf-8");
            exchange.sendResponseHeaders(resp.statusCode(), resp.body().length);
            try (OutputStream os = exchange.getResponseBody()) { os.write(resp.body()); }
        });

        // Static files: / -> index.html, /style.css, /app.js
        server.createContext("/", exchange -> {
            addCors(exchange.getResponseHeaders());

            if (!exchange.getRequestMethod().equalsIgnoreCase("GET")) {
                exchange.sendResponseHeaders(405, -1);
                return;
            }

            String urlPath = exchange.getRequestURI().getPath();
            if (urlPath == null || urlPath.isBlank() || urlPath.equals("/")) urlPath = "/index.html";

            String rel = urlPath.startsWith("/") ? urlPath.substring(1) : urlPath;
            Path requested = baseDir.resolve(rel).normalize();

            if (!requested.startsWith(baseDir)) { exchange.sendResponseHeaders(403, -1); return; }
            if (!Files.exists(requested) || Files.isDirectory(requested)) { exchange.sendResponseHeaders(404, -1); return; }

            String ct = guessContentType(requested);
            byte[] bytes = Files.readAllBytes(requested);

            exchange.getResponseHeaders().set("Content-Type", ct);
            exchange.sendResponseHeaders(200, bytes.length);
            try (OutputStream os = exchange.getResponseBody()) { os.write(bytes); }
        });

        server.start();
        System.out.println("Public UI + Proxy running on port " + PORT);
    }

    private static void addCors(Headers h) {
        h.set("Access-Control-Allow-Origin", "*");
        h.set("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
        h.set("Access-Control-Allow-Headers", "Content-Type, Authorization");
    }

    private static String guessContentType(Path file) throws IOException {
        String name = file.getFileName().toString().toLowerCase();
        if (name.endsWith(".html")) return "text/html; charset=utf-8";
        if (name.endsWith(".css")) return "text/css; charset=utf-8";
        if (name.endsWith(".js")) return "application/javascript; charset=utf-8";
        String probe = Files.probeContentType(file);
        return probe != null ? probe : "application/octet-stream";
    }

    private static String getenvOrThrow(String key) {
        String v = System.getenv(key);
        if (v == null || v.isBlank()) throw new IllegalStateException("Missing env var: " + key);
        return v;
    }
}