package com.android.commands.monkey.ape.jev;

import java.io.*;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.Map;
import java.util.Collections;
import java.util.LinkedHashMap;

/** One bounded HTTP worker. Deadlines include the queue and never wait for socket cleanup. */
public final class JevTransport implements Closeable {
    interface Completion { void completed(Result result); }
    public static final class Result {
        public final String raw,error;
        public final int httpStatus;
        public final long submittedNanos,sentNanos,completedNanos;
        private Result(Call call,String raw,String error,int status) {
            this.raw=raw;this.error=error;httpStatus=status;submittedNanos=call==null?System.nanoTime():call.submitted;
            sentNanos=call==null?submittedNanos:call.sent;completedNanos=System.nanoTime();
        }
        static Result cached(String raw){return new Result(null,raw,null,200);}
        public long waitNanos(){return completedNanos-submittedNanos;}
        public long latencyNanos(){return sentNanos==0?0:completedNanos-sentNanos;}
    }
    public final class Call {
        public final String requestId;
        private final byte[] body;
        private final int responseLimit;
        private final long submitted=System.nanoTime(),deadline;
        private volatile long sent;
        private volatile HttpURLConnection connection;
        private final java.util.concurrent.atomic.AtomicReference<Result> result=new java.util.concurrent.atomic.AtomicReference<Result>();
        private final java.util.concurrent.atomic.AtomicBoolean consumed=new java.util.concurrent.atomic.AtomicBoolean();
        private final java.util.concurrent.atomic.AtomicBoolean notified=new java.util.concurrent.atomic.AtomicBoolean();
        private volatile Completion completion;
        private final java.util.concurrent.CountDownLatch ready=new java.util.concurrent.CountDownLatch(1);
        private TimerTask expiry;
        private Call(String id,String body,int timeout,int limit,boolean retainAfterWait) {
            requestId=id;this.body=body.getBytes(StandardCharsets.UTF_8);responseLimit=limit;
            long requestDeadline=submitted+TimeUnit.MILLISECONDS.toNanos(timeout);
            deadline=retainAfterWait?requestDeadline:Math.min(JevDeadline.absolute(),requestDeadline);
        }
        private boolean finish(String raw,String error,int status) {
            if(!result.compareAndSet(null,new Result(this,raw,error,status)))return false;
            ready.countDown();
            if(expiry!=null)expiry.cancel();notifyCompletion();return true;
        }
        // The listener only enqueues immutable receipt data; it must never wait or touch GUI state.
        void onComplete(Completion listener) {completion=listener;notifyCompletion();}
        private void notifyCompletion() {
            Completion listener=completion;Result value=result.get();
            if(listener!=null && value!=null && notified.compareAndSet(false,true))
                listener.completed(value);
        }
        public Result poll() {
            if(result.get()==null && System.nanoTime()>=deadline)cancel("TIMEOUT");
            return result.get();
        }
        /** Return a completed result once. A completed HTTP call is not a stream. */
        public Result consume() {
            Result value=poll();
            return value!=null && consumed.compareAndSet(false,true)?value:null;
        }
        public Result await() {
            if(poll()==null)try {
                if(!ready.await(Math.max(0,deadline-System.nanoTime()),TimeUnit.NANOSECONDS))cancel("TIMEOUT");
            } catch(InterruptedException interrupted){cancel("INTERRUPTED");Thread.currentThread().interrupt();}
            return result.get();
        }
        /** End the caller's wait without cancelling an already submitted request. */
        public Result awaitUntil(long waitDeadline) {
            if(poll()==null)try {
                ready.await(Math.max(0,Math.min(waitDeadline,deadline)-System.nanoTime()),TimeUnit.NANOSECONDS);
            }catch(InterruptedException e){Thread.currentThread().interrupt();}
            return poll();
        }
        public void cancel(String reason) {
            if(finish(null,reason==null?"CANCELLED":reason,0)) {
                waiting.remove(this);
                HttpURLConnection current=connection;if(current!=null)disconnects.offer(current);
            }
        }
        private int remaining() throws SocketTimeoutException {
            long nanos=deadline-System.nanoTime();
            if(nanos<=0 || poll()!=null)throw new SocketTimeoutException();
            return (int)Math.min(Integer.MAX_VALUE,Math.max(1,TimeUnit.NANOSECONDS.toMillis(nanos)));
        }
        private void execute() {
            if(poll()!=null)return;
            int status=0;
            try {
                connection=(HttpURLConnection)endpoint.openConnection();
                if(poll()!=null)return;sent=System.nanoTime();
                connection.setConnectTimeout(remaining());connection.setReadTimeout(remaining());
                connection.setInstanceFollowRedirects(false);connection.setRequestMethod("POST");
                connection.setRequestProperty("Content-Type","application/json; charset=utf-8");
                connection.setRequestProperty("X-StartLux-Request-Class",requestClass);
                connection.setRequestProperty("X-APE-Request-ID",requestId);
                for(Map.Entry<String,String> field:snapshotIdentity.entrySet())
                    connection.setRequestProperty("X-APE-"+field.getKey().replace('_','-'),field.getValue());
                connection.setRequestProperty("Accept","application/json");connection.setRequestProperty("Connection","close");
                if(!token.isEmpty())connection.setRequestProperty("Authorization","Bearer "+token);
                connection.setDoOutput(true);connection.setFixedLengthStreamingMode(body.length);
                OutputStream stream=connection.getOutputStream();
                try{stream.write(body);}finally{stream.close();}
                connection.setReadTimeout(remaining());status=connection.getResponseCode();
                if(status!=200){finish(null,"HTTP_"+status,status);return;}
                ByteArrayOutputStream bytes=new ByteArrayOutputStream();byte[] buffer=new byte[4096];
                InputStream input=connection.getInputStream();
                try {
                    while(true) {
                        connection.setReadTimeout(remaining());int size=input.read(buffer);
                        if(size<0)break;
                        if(size>responseLimit-bytes.size()){finish(null,"RESPONSE_TOO_LARGE",status);return;}
                        bytes.write(buffer,0,size);
                    }
                }finally{input.close();}
                if(System.nanoTime()>=deadline)cancel("TIMEOUT");
                else finish(new String(bytes.toByteArray(),StandardCharsets.UTF_8),null,status);
            }catch(SocketTimeoutException failure){cancel("TIMEOUT");}
            catch(IOException failure){finish(null,"TRANSPORT_FAILURE",status);}
            catch(RuntimeException failure){finish(null,"TRANSPORT_FAILURE",status);}
            finally{HttpURLConnection current=connection;connection=null;if(current!=null)current.disconnect();}
        }
    }
    private final URL endpoint;
    private final String token,requestClass;
    private final ArrayBlockingQueue<Call> waiting;
    private final ArrayBlockingQueue<HttpURLConnection> disconnects=new ArrayBlockingQueue<HttpURLConnection>(1);
    private final Timer deadlines;
    private final Thread worker,cleanup;
    private volatile boolean closed;
    private volatile Call active;
    private volatile Map<String,String> snapshotIdentity=Collections.emptyMap();
    private boolean used,identitySet;
    private final java.util.concurrent.locks.ReentrantLock submitLock=new java.util.concurrent.locks.ReentrantLock();

    public void setSnapshotIdentity(Map<String,String> identity) {JevDeadline.lock(submitLock);try {
        if(used || identitySet)throw new IllegalStateException("Snapshot identity must precede all requests");
        Map<String,String> copy=new LinkedHashMap<String,String>();
        for(Map.Entry<String,String> field:identity.entrySet()) {
            if(!field.getKey().matches("experiment|device|app|rep|branch|checkpoint_id")
                    || !field.getValue().matches("[A-Za-z0-9_-]{1,160}"))throw new IllegalArgumentException("Invalid Snapshot HTTP identity");
            copy.put(field.getKey(),field.getValue());
        }
        snapshotIdentity=Collections.unmodifiableMap(copy);identitySet=true;
    } finally {submitLock.unlock();}}

    public JevTransport(String endpoint,String token,int capacity,String requestClass)throws IOException {
        if(!"probe".equals(requestClass) && !"region".equals(requestClass))throw new IllegalArgumentException("invalid request class");
        this.requestClass=requestClass;
        if(capacity<1 || capacity>64)throw new IllegalArgumentException("invalid request capacity");
        this.endpoint=new URL(endpoint);
        String host=this.endpoint.getHost(),scheme=this.endpoint.getProtocol();
        if(this.endpoint.getUserInfo()!=null || !("https".equals(scheme) || "http".equals(scheme)
            && ("127.0.0.1".equals(host) || "localhost".equals(host) || "211.71.15.47".equals(host))))
            throw new IllegalArgumentException("invalid decision endpoint");
        this.token=token==null?"":token;waiting=new ArrayBlockingQueue<Call>(capacity);
        deadlines=new Timer("ape-"+requestClass+"-deadline",true);
        worker=new Thread(new Runnable(){public void run(){
            while(!closed)try{Call call=waiting.take();active=call;try{if(closed)call.cancel("CLOSED");else call.execute();}finally{active=null;}}
            catch(InterruptedException interrupted){if(closed)return;}
        }},"ape-"+requestClass+"-http");worker.setDaemon(true);worker.start();
        cleanup=new Thread(new Runnable(){public void run(){
            while(!closed || !disconnects.isEmpty())try {
                HttpURLConnection connection=disconnects.poll(100,TimeUnit.MILLISECONDS);if(connection!=null)connection.disconnect();
            }catch(InterruptedException interrupted){if(closed && disconnects.isEmpty())return;}
            catch(RuntimeException ignored){ }
        }},"ape-"+requestClass+"-cleanup");cleanup.setDaemon(true);cleanup.start();
    }
    public Call submit(String requestId,String body,int timeoutMs,int maxResponseBytes) {
        return submit(requestId,body,timeoutMs,maxResponseBytes,false);
    }
    /** Only the network lifetime outlives the caller's budget; enqueue remains budgeted. */
    Call submitRetained(String requestId,String body,int timeoutMs,int maxResponseBytes) {
        return submit(requestId,body,timeoutMs,maxResponseBytes,true);
    }
    private Call submit(String requestId,String body,int timeoutMs,int maxResponseBytes,boolean retained) {JevDeadline.lock(submitLock);try {
        if(requestId==null || requestId.isEmpty() || body==null || timeoutMs<=0 || maxResponseBytes<=0)
            throw new IllegalArgumentException("invalid request bounds");
        if(closed)return null;
        JevDeadline.check();used=true;
        final Call call=new Call(requestId,body,timeoutMs,maxResponseBytes,retained);
        call.expiry=new TimerTask(){public void run(){call.cancel("TIMEOUT");}};
        // Schedule before publishing: an immediate response may cancel this task.
        deadlines.schedule(call.expiry,Math.max(1,TimeUnit.NANOSECONDS.toMillis(Math.max(0,call.deadline-System.nanoTime()))));
        if(!waiting.offer(call)){call.cancel("QUEUE_FULL");return null;}
        return call;
    } finally {submitLock.unlock();}}
    int outstanding(){Call call=active;return waiting.size()+(call!=null && call.poll()==null?1:0);}
    public void close() {JevDeadline.lock(submitLock);try {
        if(closed)return;closed=true;
        Call call=active;if(call!=null)call.cancel("CLOSED");
        while((call=waiting.poll())!=null)call.cancel("CLOSED");
        deadlines.cancel();worker.interrupt();cleanup.interrupt();
    } finally {submitLock.unlock();}}
}