#!/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/pls/prez2023/vysledky'
	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.volby.cz/prezident/}' # 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}CR/{NS}KANDIDAT'):
		ka = kandidat.attrib
		jmeno = f'{ka[r"JMENO"]} {ka[r"PRIJMENI"]}'
		hlasy = int(ka.get(r'HLASY_1KOLO', 0))
		kandidati[jmeno] = hlasy
	
	for ucast in root.findall(rf'./{NS}CR/{NS}UCAST'):
		if ucast.attrib[r'KOLO'] == '1':
			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."""
	VOLICU_CELKEM = 8_500_000 # Odhad!!
	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():
		order = (
			'Petr Pavel',
			'Danuše Nerudová',
			'Andrej Babiš',
			'Jaroslav Bašta',
			'Marek Hilšer',
			'Pavel Fischer',
			'Karel Diviš',
			'Tomáš Zima',
			'NEPLATNÉ',
			'NEVOLILI',
			'NEZAPOČÍTANÉ',
			)
		x = [kand[i] for i in order]
		labels = order
		label = f'Data z {ts}'
		colors = {
			'Petr Pavel': '#627210',
			'Danuše Nerudová': '#811367',
			'Andrej Babiš': '#262161',
			'Jaroslav Bašta': '#B51119',
			'Marek Hilšer': '#CA834E',
			'Pavel Fischer': '#244C76',
			'Karel Diviš': '#3B6E5D',
			'Tomáš Zima': '#E5DE1A',
			'NEPLATNÉ': '#000000',
			'NEVOLILI': '#666666',
			'NEZAPOČÍTANÉ': '#CCCCCC',
			}
		plt.pie(x, labels=labels, colors=[colors[i] for i in order], autopct='%d %%')
		plt.text(0, -1.2, label, ha='center')

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

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

loop()