package org.example;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;

public class LeaderElection implements Watcher {
    private static final String ADDRESS = "192.168.181.136:2181";
    private static final int SESSION_TIMEOUT = 3000;
    private static final String ELECTION_NAMESPACE = "/election";
    private static final String LEADER_NODE_PREFIX = "leader-";

    private String currentNodePath;
    private String previousNodePath;
    private ZooKeeper zooKeeper;
    private Stat previousNodeStat;

    public static void main(String[] args) throws IOException, InterruptedException, KeeperException {
        LeaderElection leaderElection = new LeaderElection();
        leaderElection.connectToZookeeper();
        leaderElection.volunteerForLeadership();
        leaderElection.electLeader();
        leaderElection.run();
    }

    public void connectToZookeeper() throws IOException {
        this.zooKeeper = new ZooKeeper(ADDRESS, SESSION_TIMEOUT, this);
    }

    public void run() {
        synchronized (zooKeeper) {
            try {
                zooKeeper.wait(); // Wait for events
            } catch (InterruptedException e) {
                System.out.println(e.getMessage());
            }
        }
        try {
            zooKeeper.close();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    public void volunteerForLeadership() throws InterruptedException, KeeperException {
        String znodeFullPath = zooKeeper.create(
                ELECTION_NAMESPACE + "/" + LEADER_NODE_PREFIX,
                new byte[]{},
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL_SEQUENTIAL
        );

        System.out.println("Created znode: " + znodeFullPath);
        this.currentNodePath = znodeFullPath.replace(ELECTION_NAMESPACE + "/", "");
    }

    public void electLeader() throws InterruptedException, KeeperException {
        List<String> children = zooKeeper.getChildren(ELECTION_NAMESPACE, false);

        Collections.sort(children);
        int currentNodeIndex = children.indexOf(currentNodePath);

        if (currentNodeIndex == 0) {
            becomeLeader();
        } else {
            String previousNodeName = children.get(currentNodeIndex - 1);
            previousNodePath = ELECTION_NAMESPACE + "/" + previousNodeName;
            System.out.println("I'm not the leader, the leader is: " + children.get(0));
            watchPreviousNode();
        }
    }

    private void becomeLeader() throws KeeperException, InterruptedException {
        System.out.println("I'm the leader");
        // Notify other nodes that you are the leader
        String leaderData = "Leader: " + currentNodePath;
        zooKeeper.setData(ELECTION_NAMESPACE+"/"+currentNodePath, leaderData.getBytes(), -1);
    }

    private void watchPreviousNode() throws KeeperException, InterruptedException {
        previousNodeStat = zooKeeper.exists(previousNodePath, this);
        if (previousNodeStat != null) {
            System.out.println("Watching previous node: " + previousNodePath);
        }
    }
    private void setData() throws InterruptedException, KeeperException {
        // Retrieve leader data from the previous leader's znode, because it contains our new leader
        byte[] data = zooKeeper.getData(previousNodePath, this, previousNodeStat);
        String leaderData = new String(data);
        zooKeeper.setData(ELECTION_NAMESPACE+"/"+currentNodePath,data,-1);
        System.out.println("Current leader: " + leaderData);
    }
    @Override
    public void process(WatchedEvent watchedEvent) {
        switch (watchedEvent.getType()) {
            case None:
                if (watchedEvent.getState() == Event.KeeperState.SyncConnected) {
                    System.out.println("Successfully connected to ZooKeeper");
                } else if (watchedEvent.getState() == Event.KeeperState.Disconnected) {
                    System.out.println("Disconnected from ZooKeeper. Shutting down the program.");
                    synchronized (zooKeeper) {
                        zooKeeper.notifyAll();
                    }
                }
                break;
            case NodeDeleted:
                if (watchedEvent.getPath().equals(previousNodePath)) {
                    System.out.println("Previous leader node deleted. Attempting to become the new leader...");
                    try {
                        electLeader();
                    } catch (InterruptedException | KeeperException e) {
                        throw new RuntimeException(e);
                    }
                }
                break;
            case NodeDataChanged:
                if (watchedEvent.getPath().equals(previousNodePath)) {
                    System.out.println("Previous leader data changed. maybe we have a new leader .");
                    try {
                        setData();
                    } catch (InterruptedException | KeeperException e) {
                        throw new RuntimeException(e);
                    }
                }
                break;
            default:
                break;
        }
    }
}