import requests
import xmltodict
import json
import os


class OkrsokResult:
    def __init__(self, key, yes_count, no_count):
        self.key = key
        self.yes_count = yes_count
        self.no_count = no_count

    def to_dict(self):
        return {
            "key": self.key,
            "yes_count": self.yes_count,
            "no_count": self.no_count
        }


class OkrsokCZ2024Result:
    def __init__(self, key, party_14_percent, party_23_percent, party_17_percent, other_percent, count):
        self.key = key
        self.party_14_percent = party_14_percent
        self.party_23_percent = party_23_percent
        self.party_17_percent = party_17_percent
        self.other_percent = other_percent
        self.count = count

    def to_dict(self):
        return {
            "key": self.key,
            "party_14_percent": self.party_14_percent,
            "party_23_percent": self.party_23_percent,
            "party_17_percent": self.party_17_percent,
            "other_percent": self.other_percent,
            "count": self.count
        }


def prez_2023():
    results = []

    for i in range(1, 38):
        url = f"https://www.volby.cz/pls/prez2023/vysledky_okrsky?kolo=2&davka={i}"
        response = requests.get(url)
        doc = xmltodict.parse(response.content)
        json_text = json.loads(json.dumps(doc))
        vys = json_text["VYSLEDKY_OKRSKY"]
        davka = vys["DAVKA"]
        obec_key = "@CIS_OBEC"
        okrsek_key = "@CIS_OKRSEK"

        if isinstance(vys["OKRSEK"], dict):
            yes_count = 0
            no_count = 0
            val = vys["OKRSEK"]["HLASY_OKRSEK"]
            hlas_key = "@HLASY"
            if not isinstance(val, list):
                if val["@PORADOVE_CISLO"] == "4":
                    no_count = val[hlas_key]
                else:
                    yes_count = val[hlas_key]
            else:
                yes_count = val[1][hlas_key]
                no_count = val[0][hlas_key]

            results.append(
                OkrsokResult(f"{vys['OKRSEK'][obec_key]}-{vys['OKRSEK'][okrsek_key]}", int(yes_count), int(no_count)))
        else:
            for o in vys["OKRSEK"]:
                yes_count = 0
                no_count = 0
                val = davka["@OKRSKY_DAVKA"] == "1" and o or o["HLASY_OKRSEK"]
                hlas_key = "@HLASY"
                if not isinstance(val, list):
                    if val["@PORADOVE_CISLO"] == "4":
                        no_count = val[hlas_key]
                    else:
                        yes_count = val[hlas_key]
                else:
                    yes_count = val[1][hlas_key]
                    no_count = val[0][hlas_key]

                results.append(OkrsokResult(f"{o[obec_key]}-{o[okrsek_key]}", int(yes_count), int(no_count)))

        print(os.getcwd())
        with open(f"cz_prezident_2023/results-{i}.json", "w") as f:
            json.dump([result.to_dict() for result in results], f, indent=4)

    return results


def euro_2024():
    results = []

    for i in range(1, 61):
        url = f"https://www.volby.cz/pls/ep2024/vysledky_okrsky?davka={i}"
        response = requests.get(url)
        doc = xmltodict.parse(response.content)
        json_text = json.loads(json.dumps(doc))
        vys = json_text["VYSLEDKY_OKRSKY"]
        davka = vys["DAVKA"]
        obec_key = "@CIS_OBEC"
        okrsek_key = "@CIS_OKRSEK"
        print(f"Processing batch {i}...")

        def process_okrsok(okrsok):
            party_14_count = 0
            party_23_count = 0
            party_17_count = 0
            other_count = 0

            hlas_key = "@HLASY"
            estrana_key = "@ESTRANA"

            for val in okrsok.get("HLASY_OKRSEK", []):
                if val[estrana_key] == "14":
                    party_14_count = int(val[hlas_key])
                elif val[estrana_key] == "23":
                    party_23_count = int(val[hlas_key])
                elif val[estrana_key] == "17":
                    party_17_count = int(val[hlas_key])
                else:
                    other_count += int(val[hlas_key])

            total_votes = party_14_count + party_23_count + party_17_count + other_count
            if total_votes == 0:
                party_14_percent = party_23_percent = party_17_percent = other_percent = 0
            else:
                party_14_percent = (party_14_count / total_votes)
                party_23_percent = (party_23_count / total_votes)
                party_17_percent = (party_17_count / total_votes)
                other_percent = (other_count / total_votes)

            return OkrsokCZ2024Result(f"{okrsok[obec_key]}-{okrsok[okrsek_key]}", party_14_percent, party_23_percent,
                                      party_17_percent, other_percent, total_votes)

        if isinstance(vys["OKRSEK"], dict):
            results.append(process_okrsok(vys["OKRSEK"]))
        else:
            results.extend([process_okrsok(o) for o in vys["OKRSEK"]])

        with open(f"cz_euro_2024/results-{i}.json", "w") as f:
            json.dump([result.to_dict() for result in results], f, indent=4)
    return results


if __name__ == '__main__':
    euro_2024()