from pprint import pprint
from locust import FastHttpUser, between, task

password = "fooobarasdfasdf"
username = "realaravinth"


class Unprotected(FastHttpUser):
    #    wait_time = between(5, 15)
    peers = [
            "http://localhost:9001",
            "http://localhost:9002",
            "http://localhost:9003",
            ]
    leader = "http://localhost:9001"
    host = leader
    captcha_id="locust"

    pipeline_vote = []
    for _ in range(0,1000):
        pipeline_vote.append({"AddVisitor": captcha_id})

#    def on_start(self):
#        resp  = self.client.get(f"{self.leader}/metrics")
#        data = resp.json()
#        leader = data["Ok"]["membership_config"]["log_id"]["leader_id"]["node_id"]
#        self.leader = self.peers[leader - 1]
#        self.host = self.leader
#        print(f"Leader: {self.host}")
#        self.add_captcha(captcha_id="locust")



    def write(self, data):
        resp = self.client.post(f"{self.host}/write", json=data)
        print(f"RPC Status: {resp.status_code}")
        resp = resp.json()
        if "Err" in resp:
            leader = resp["Err"]["APIError"]["ForwardToLeader"]["leader_node"]["addr"]
            print(f"Forwarding write to leader {leader}")
            return write(leader, data)
        return resp["Ok"]["data"]

    def pipeline_write(self, data):
        resp = self.client.post(f"{self.host}/pipeline/write", json=data)
#        print(f"RPC Status: {resp.status_code}")
        resp = resp.json()
        if "Err" in resp:
            leader = resp["Err"]["APIError"]["ForwardToLeader"]["leader_node"]["addr"]
            print(f"Forwarding write to leader {leader}")
            return write(leader, data)
        return resp




    def add_vote(self, captcha_id: str):
        resp = self.write(data={"AddVisitor": captcha_id})
        pprint(resp)

    def add_vote_pipeline(self, captcha_id: str):
        resp = self.pipeline_write(data=self.pipeline_vote)
#        pprint(resp)

    def add_captcha(self, captcha_id: str):
        params = {
            "AddCaptcha": {
                "id": captcha_id,
                "mcaptcha": {
                    "visitor_threshold": 0,
                    "defense": {
                        "levels": [
                            {"visitor_threshold": 50, "difficulty_factor": 500},
                            {"visitor_threshold": 5000, "difficulty_factor": 50000},
                        ],
                        "current_visitor_threshold": 0,
                    },
                    "duration": 30,
                },
            }
        }
        resp = self.write(data=params)
        pprint(f"Captcha added {captcha_id}: {resp}")


    @task
    def unprotected(self):
        self.add_vote_pipeline(captcha_id=self.captcha_id)
        ##self.add_vote(captcha_id="locust")
#        data = {
#            "username": username,
#            "password": username,
#            "confirm_password": username,
#        }
#        self.client.post("/unprotected", data=data)