#!/usr/bin/env python3
import requests
from matplotlib import pyplot as plt
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from datetime import datetime
from io import StringIO

# In order of what we know:

class VisuWindow:
	def __enter__(self):
		plt.close()
		plt.clf()
	def __exit__(self, _et, _ev, _tb):
		plt.show(block=False)

# Alternative: plot to file
class VisuFile:
	def __init__(self, filename):
		self.filename = filename
	
	def __enter__(self):
		plt.clf()
	def __exit__(self, _et, _ev, _tb):
		plt.savefig(self.filename, format='pdf')

class VisuInter:
	def __init__(self):
		plt.ion()
	def __enter__(self):
		plt.clf()
	def __exit__(self, _et, _ev, _tb):
		plt.show(block=False)

def get_xml():
	URL=r'https://www.volby.cz/appdata/kz2024/odata/vysledky.xml'
	response = requests.get(URL)
	return StringIO(response.content.decode())

@dataclass
class Vysledek:
	kandidati: dict[str, int]
	timestamp: datetime
	ucast: dict

def get_results() -> Vysledek:
	NS = r'{http://www.w3.org/namespace/}' # fuck you xml
	xml = get_xml()
	root = ET.parse(xml).getroot()
	kandidati = {}

	timestamp = datetime.fromisoformat(root.attrib[r'DATUM_CAS_GENEROVANI'])
	for kandidat in root.findall(rf'./{NS}KRZAST[@CIS_KRZAST="1"]/{NS}STRANA'):
		ka = kandidat.attrib
		jmeno = f'{ka[r"NAZ_STR"]}'
		hlasy = kandidat.findall(rf'./{NS}HODNOTY_STRANA')[0].attrib[r'HLASY']
		kandidati[jmeno] = hlasy
	
	for ucast in root.findall(rf'./{NS}KRZAST[@CIS_KRZAST="1"]/{NS}UCAST'):
		fucast = ucast.attrib
	
	return Vysledek(kandidati=kandidati, timestamp=timestamp, ucast=fucast)

def fill_data(vysl) -> tuple[datetime, dict[str, int]]:
	"""Dopočítá ve výsledku zbytek dat."""
	print(vysl)
	VOLICU_CELKEM = 1_044_056
	vysl.kandidati[r'NEPLATNÉ'] = int(vysl.ucast[r'ODEVZDANE_OBALKY']) - int(vysl.ucast[r'PLATNE_HLASY'])
	vysl.kandidati[r'NEZAPOČÍTANÉ'] = max(0, VOLICU_CELKEM - int(vysl.ucast[r'ZAPSANI_VOLICI']))
	vysl.kandidati[r'NEVOLILI'] = int(vysl.ucast[r'ZAPSANI_VOLICI']) - int(vysl.ucast[r'ODEVZDANE_OBALKY'])
	# Na ztracené obálky (rozdíl vydaných obálek a odevzdaných) kašleme, to by mělo být velmi zanedbatelné číslo
	return vysl.timestamp, vysl.kandidati

def plot(ts, kand):
	with VisuInter():
		labels, values = list(zip(*kand.items()))
		plt.pie(values, autopct='%1.3f %%', labels=labels)
		plt.text(0, -1.2, f'Data z {ts}', ha='center')

def visu_once():
	while True:
		try:
			v = get_results()
			break
		except OSError:
			print('fuu')
			plt.pause(1)
			continue
	
	ts, k = fill_data(v)
	plot(ts, k)

def loop(refresh=30):
	while True:
		visu_once()
		plt.pause(refresh)

loop()