package com.lumaserv.bgp;

import com.lumaserv.bgp.protocol.message.BGPNotification;
import com.lumaserv.bgp.protocol.message.BGPOpen;
import com.lumaserv.bgp.protocol.message.BGPUpdate;

import lombok.Getter;
import lombok.Setter;

import java.io.IOException;
import java.util.Optional;
import java.util.Timer;
import java.util.TimerTask;

public class BGPFsm implements Callback {
    public static final int CONFIG_HOLD_TIME = 80;  // FIXME use from configuration
    public static final byte CONFIG_VERSION = 4;

    public final State IDLE;
    public final State CONNECT;
    public final State ACTIVE;
    public final State OPEN_SENT;
    public final State OPEN_CONFIRM;
    public final State ESTABLISHED;

    @Getter
    @Setter
    State currentState;
    boolean delayOpen;
    int holdTime;
    int connectRetryCounter;
    int connectRetryTime = 5; //120;
    @Setter
    BGPSession session;
    Timer keepAliveTimer;
    TimerTask keepAliveTimerTask;
    Timer holdTimer;
    TimerTask holdTimerTask;
    Timer connectRetryTimer;
    TimerTask connectRetryTimerTask;

    class State {
        void setState(State next) {
            System.out.println("State transition: " + BGPFsm.this + ":" + this + " -> " + next);
            // if (next == IDLE)
            //     throw new RuntimeException();
            currentState = next;
        }
        void manualStart() { throw new RuntimeException("Not implemented: " + this); }
        void automaticStart() { throw new RuntimeException("Not implemented: " + this); }
        void manualStartPassive() { throw new RuntimeException("Not implemented: " + this); }
        void automaticStartPassive() { throw new RuntimeException("Not implemented: " + this); }
        void manualStop() { throw new RuntimeException("Not implemented: " + this); }
        void automaticStop() { throw new RuntimeException("Not implemented: " + this); }
        void holdTimerExpires() { throw new RuntimeException("Not implemented: " + this); }
        void keepAliveTimerExpires() { throw new RuntimeException("Not implemented: " + this); }
        void connectRetryTimerExpires() { throw new RuntimeException("Not implemented: " + this); }
        void tcpCRAcked() { throw new RuntimeException("Not implemented: " + this); }
        void tcpConnectionConfirmed() { throw new RuntimeException("Not implemented: " + this); }
        void tcpConnectionFails() { throw new RuntimeException("Not implemented: " + this); }
        void openMsg(BGPOpen packet)  { throw new RuntimeException("Not implemented: " + this); }
        void keepAliveMsg() { throw new RuntimeException("Not implemented: " + this); }
        void updateMsg(BGPUpdate packet) { throw new RuntimeException("Not implemented: " + this); }
        void notifMsgVerErr(BGPNotification msg) { throw new RuntimeException("Not implemented: " + this); }
        void notifMsg(BGPNotification msg) { throw new RuntimeException("Not implemented: " + this); }
    }

    class Idle extends State {
        protected void start() {
            // - initializes all BGP resources for the peer connection,
            setConnectRetryTimer(connectRetryTime);
            setState(CONNECT);
            // - listens for a connection that may be initiated by the remote
            //   BGP peer, and
            connectRetryCounter = 0;
            session.startOutgoing();
        }

        @Override
        void manualStart() {
            start();
        }

        @Override
        void automaticStart() {
            start();
        }

        @Override
        void manualStop() {
            // ignore
        }

        @Override
        void automaticStop() {
            // ignore
        }

        protected void startPassive() {
            // - initializes all BGP resources,
            // - sets the ConnectRetryCounter to zero,
            setConnectRetryTimer(connectRetryTime);
            // - listens for a connection that may be initiated by the remote
            //   peer, and
            setState(ACTIVE);
            session.startIncoming();
        }

        @Override
        void manualStartPassive() {
            startPassive();
        }

        @Override
        void automaticStartPassive() {
            startPassive();
        }

        // TODO implement DampPeerOscillations related events?

        // TODO
        // Any other event (Events 9-12, 15-28) received in the Idle state
        // does not cause change in the state of the local system.
        @Override
        void connectRetryTimerExpires() {} // Ignore event 9
        @Override
        void holdTimerExpires() {} // Ignore event 10
        @Override
        void keepAliveTimerExpires() {} // Ignore event 11
        // Event 12 not implemented
        // Event 15 not implemented
        @Override
        void tcpCRAcked() {} // Ignore event 16
        @Override
        void tcpConnectionConfirmed() {} // Ignore event 17
        @Override
        void tcpConnectionFails() {} // Ignore event 18
        @Override
        void openMsg(BGPOpen msg) {} // Ignore event 19
        // Event 20 not implemented
        //@Override
        // TODO void bgpHeaderErr() {} // Ignore event 21
        @Override
        public String toString() { return "Idle"; }
    }

    class Connect extends State {
        // TODO ignore start events 1, 3-7

        protected void stop(boolean isAutomatic) {
            session.dropConnection();
            // - releases all BGP resources,
            if (isAutomatic) {
                connectRetryCounter++;
            } else {
                connectRetryCounter = 0;
            }
            setConnectRetryTimer(0);
            setState(IDLE);
        }

        @Override
        void manualStop() {
            stop(false);
        }

        // AutomaticStop isn't mentioned in the Connect state in RFC 4271,
        // but it's needed at least in WireGuard.
        @Override
        void automaticStop() {
            stop(true);
        }

        @Override
        void connectRetryTimerExpires() {
            System.out.println("connectRetryTimerExpires");
            session.dropConnection();
            setConnectRetryTimer(connectRetryTime);
            // - stops the DelayOpenTimer and resets the timer to zero,
            session.retryConnection();
            // - continues to listen for a connection that may be initiated by
            //   the remote BGP peer, and
        }

        // TODO DelayOpenTimer_Expires event (Event 12)
        // TODO TcpConnection_Valid event (Event 14)
        // TODO Tcp_CR_Invalid event (Event 15)

        protected void tcpCRAckedOrTcpConnectionConfirmed() {
            if (delayOpen) {
                setConnectRetryTimer(0);
                // - sets the DelayOpenTimer to the initial value, and
            } else {
                setConnectRetryTimer(0);
                // - completes BGP initialization

                try {
                    session.sendOpen();
                    setHoldTimer(4 * 60);
                    setState(OPEN_SENT);
                } catch (IOException ex) {
                    tcpConnectionFails();
                }
            }
        }

        @Override
        void tcpCRAcked() {
            tcpCRAckedOrTcpConnectionConfirmed();
        }

        @Override
        void tcpConnectionConfirmed() {
            tcpCRAckedOrTcpConnectionConfirmed();
        }

        // FIXME
        @Override
        void tcpConnectionFails() {
            setConnectRetryTimer(0);
            // - releases all BGP resources,
            // - (optionally) performs peer oscillation damping if the
            //   DampPeerOscillations attribute is set to TRUE, and
            session.dropConnection();
            connectRetryCounter++;
            setState(IDLE);
        }

        // TODO
        // State notifMsgVerErr() {
        //     // - sets the ConnectRetryTimer to zero,
        //     // - releases all BGP resources,
        //     // - drops the TCP connection, and
        //     return IDLE;
        // }

        // State bgpOpen() {
        //     return
        // }

        @Override
        public String toString() { return "Connect"; }
    }

    class Active extends State {
        @Override
        void manualStart() {}
        @Override
        void automaticStart() {}
        @Override
        void manualStartPassive() {}
        @Override
        void automaticStartPassive() {}

        protected void stop() {
            // - If the DelayOpenTimer is running and the
            try {
                BGPNotification notification = new BGPNotification()
                    .setMajorErrorCode(BGPNotification.Error.CEASE.getCode());
                session.sendNotification(notification);
            } catch (IOException ex) {
            }
            // - releases all BGP resources including stopping the
            //   DelayOpenTimer
            session.dropConnection();
            connectRetryCounter = 0;
            setConnectRetryTimer(0);
            setState(IDLE);
        }

        @Override
        void manualStop() {
            stop();
        }

        @Override
        void automaticStop() {
            stop();
        }

        @Override
        void connectRetryTimerExpires() {
            setConnectRetryTimer(connectRetryTime);
            setState(CONNECT);
            session.startOutgoing();
        // - continues to listen for a TCP connection that may be initiated
        //   by a remote BGP peer, and
        }

        @Override
        void tcpConnectionConfirmed() {
            if (delayOpen) {
                setConnectRetryTimer(0);
                // - sets the DelayOpenTimer to the initial value
                //   (DelayOpenTime), and
                // - stays in the Active state.
            } else {
                setConnectRetryTimer(0);
                // - completes the BGP initialization,

                try {
                    session.sendOpen();
                    setHoldTimer(4 * 60);
                    setState(OPEN_SENT);
                } catch (IOException ex) {
                    tcpConnectionFails();
                }
            }
        }

        @Override
        public String toString() { return "Active"; }
    }

    class OpenSent extends State {
        protected void stop(boolean isAutomatic) {
            try {
                BGPNotification notification = new BGPNotification()
                    .setMajorErrorCode(BGPNotification.Error.CEASE.getCode());
                session.sendNotification(notification);
            } catch (IOException ex) {
            }
            setConnectRetryTimer(0);
            // - releases all the BGP resources,
            session.dropConnection();
            if (isAutomatic) {
                connectRetryCounter++;
                // - (optionally) performs peer oscillation damping if the
                //   DampPeerOscillations attribute is set to TRUE, and
            } else {
                connectRetryCounter = 0;
            }
            setState(IDLE);
        }

        @Override
        void manualStop() {
            stop(false);
        }

        @Override
        void automaticStop() {
            stop(true);
        }

        @Override
        void tcpConnectionFails() {
            session.dropConnection();
            setConnectRetryTimer(connectRetryTime);
            // - continues to listen for a connection that may be initiated by
            //   the remote BGP peer, and
            setState(ACTIVE);
        }

        @Override
        void openMsg(BGPOpen msg) {
            System.out.println("openMsg:" + msg);

            Optional<BGPOpen.AS4Capability> as4Opt =
                msg.getAS4Capability();
            int remoteAsn = session.getConfiguration().getRemoteAs();
            int openAsn = as4Opt.isPresent()
                ? as4Opt.get().getAsn()
                : msg.getAsn();

            if(remoteAsn == openAsn) {
                // - resets the DelayOpenTimer to zero,
                // - sets the BGP ConnectRetryTimer to zero,
                session.keepAlive();
                if (msg.getHoldTime() < CONFIG_HOLD_TIME)
                    holdTime = msg.getHoldTime();
                else
                    holdTime = CONFIG_HOLD_TIME;
                setKeepAliveTimer();
                setHoldTimer(holdTime);
                setState(OPEN_CONFIRM);
            } else {
                System.out.println("Bad asn:" + (openAsn & 0xFFFFFFFFL));
                try {
                    BGPNotification notification = new BGPNotification()
                        .setMajorErrorCode(BGPNotification.Error.OPEN_MESSAGE_ERROR.getCode())
                        .setMinorErrorCode(BGPNotification.OpenMessageError.BAD_PEER_AS.getCode());
                    session.sendNotification(notification);
                } catch (IOException ex) {
                }
                setConnectRetryTimer(0);
                // - releases all BGP resources,
                session.dropConnection();
                connectRetryCounter++;
                // - (optionally) performs peer oscillation damping if the
                //   DampPeerOscillations attribute is TRUE, and

                setState(IDLE);
            }
        }

        @Override
        public String toString() { return "OpenSent"; }
    }

    class OpenConfirm extends State {
        @Override
        void keepAliveMsg() {
            // - restarts the HoldTimer and
            setState(ESTABLISHED);
            session.getConfiguration().getListener().onOpen(session);
        }

        // TODO implement keepAliveTimerExpires and holdTimerExpires

        @Override
        void tcpConnectionFails() {
            tcpConnectionFailsOrNotifMsg();
        }

        @Override
        void notifMsg(BGPNotification msg) {
            tcpConnectionFailsOrNotifMsg();
        }

        protected void tcpConnectionFailsOrNotifMsg() {
            setConnectRetryTimer(0);
            // - releases all BGP resources,
            session.dropConnection();
            connectRetryCounter++;
            // - (optionally) performs peer oscillation damping if the
            //   DampPeerOscillations attribute is set to TRUE, and
            setState(IDLE);
        }

        @Override
        public String toString() { return "OpenConfirm"; }
    }

    class Established extends State {
        protected void stop(boolean isAutomatic) {
            try {
                BGPNotification notification = new BGPNotification()
                    .setMajorErrorCode(BGPNotification.Error.CEASE.getCode());
                session.sendNotification(notification);
            } catch (IOException ex) {
            }
            setConnectRetryTimer(0);
            // - deletes all routes associated with this connection,
            // - releases all BGP resources,
            session.dropConnection();
            if (isAutomatic) {
                connectRetryCounter++;
                // - (optionally in automaticStop) performs peer oscillation damping if the
                // DampPeerOscillations attribute is set to TRUE, and
            } else {
                connectRetryCounter = 0;
            }
            setState(IDLE);
        }

        @Override
        void manualStop() {
            stop(false);
        }

        @Override
        void automaticStop() {
            stop(true);
        }

        @Override
        void keepAliveMsg() {
            setHoldTimer(holdTime);
        }

        @Override
        void updateMsg(BGPUpdate packet) {
            // - processes the message,
            setHoldTimer(holdTime);
        }

        @Override
        void holdTimerExpires() {
            // - sends a NOTIFICATION message with the Error Code Hold Timer
            //   Expired,
            setConnectRetryTimer(0);
            // - releases all BGP resources,
            session.dropConnection();
            connectRetryCounter++;
            // - (optionally) performs peer oscillation damping if the
            //   DampPeerOscillations attribute is set to TRUE, and
            setState(IDLE);
        }

        @Override
        void keepAliveTimerExpires() {
            session.keepAlive();
            setKeepAliveTimer();
        }

        @Override
        void tcpConnectionFails() {
            tcpConnectionFailsOrNotifMsg();
        }

        @Override
        void notifMsg(BGPNotification msg) {
            tcpConnectionFailsOrNotifMsg();
        }

        protected void tcpConnectionFailsOrNotifMsg() {
            setConnectRetryTimer(0);
            // - deletes all routes associated with this connection,
            // - releases all BGP resources,
            session.dropConnection();
            connectRetryCounter++;
            setState(IDLE);
        }

        @Override
        public String toString() { return "Established"; }
    }

    BGPFsm(BGPSession session) {
        this.session = session;
        IDLE = new Idle();
        CONNECT = new Connect();
        ACTIVE = new Active();
        OPEN_SENT = new OpenSent();
        OPEN_CONFIRM = new OpenConfirm();
        ESTABLISHED = new Established();
        currentState = IDLE;
        keepAliveTimer = new Timer();
        holdTimer = new Timer();
        connectRetryTimer = new Timer();
    }

    public void cleanup() {
        System.out.println("BGPFsm cleanup");
        if (keepAliveTimerTask != null) {
            keepAliveTimerTask.cancel();
            keepAliveTimerTask = null;
        }
        keepAliveTimer = null;

        if (holdTimerTask != null) {
            holdTimerTask.cancel();
            holdTimerTask = null;
        }
        holdTimer = null;

        if (connectRetryTimerTask != null) {
            connectRetryTimerTask.cancel();
            connectRetryTimerTask = null;
        }
        connectRetryTimer = null;
    }

    void setKeepAliveTimer() {
        //        System.out.println("setKeepAliveTimer: " + holdTime / 3);

        if (keepAliveTimerTask != null) {
            keepAliveTimerTask.cancel();
            keepAliveTimerTask = null;
        }

        if (holdTime <= 0)
            return;

        keepAliveTimerTask = new FsmTimerTask(this);
        keepAliveTimer.schedule(keepAliveTimerTask, 1000 * holdTime / 3);
    }

    private void setHoldTimer(int delay) {
        //        System.out.println("setHoldTimer: " + delay);
        if (holdTimerTask != null) {
            holdTimerTask.cancel();
            holdTimerTask = null;
        }

        if (delay <= 0)
            return;

        holdTimerTask = new FsmTimerTask(this);
        holdTimer.schedule(holdTimerTask, 1000 * delay);
    }

    private void setConnectRetryTimer(int delay) {
        System.out.println("setHoldTimer(" + this + "):" + delay);
        if (connectRetryTimerTask != null) {
            connectRetryTimerTask.cancel();
            connectRetryTimerTask = null;
        }

        if (delay <= 0)
            return;

        connectRetryTimerTask = new FsmTimerTask(this);
        connectRetryTimer.schedule(connectRetryTimerTask, 1000 * delay);
    }

    public void callback(Object task) {
        if (task == connectRetryTimerTask) {
            currentState.connectRetryTimerExpires();
        } else if (task == holdTimerTask) {
            currentState.holdTimerExpires();
        } else if (task == keepAliveTimerTask) {
            currentState.keepAliveTimerExpires();
        } else {
            throw new RuntimeException("Unknown timer task: " + task);
        }
    }

    static class FsmTimerTask extends TimerTask {
        Callback owner;

        FsmTimerTask(Callback owner) {
            this.owner = owner;
        }

        public void run() {
            if (owner != null) {
                owner.callback(this);
            }
        }

        public boolean cancel() {
            boolean res = super.cancel();
            // Break cyclic references to allow GC
            this.owner = null;
            return res;
        }
    }
}

interface Callback {
    void callback(Object cb);
}