#!/usr/bin/env python3
# Naplní databázi školami a obcemi, v nichž školy sídlí
# Používá extra/skoly/parsed/*.tsv
#
# Pozor, zrada: rejstřík škol je sice rozdělený do okresů dle NUTS/LAU,
# ale školy tam řadí podle úřadu, u nějž je škole registrovaná, což vůbec
# nemusí odpovídat skutečnému sídlu školy. Proto si poněkud magicky pomáháme
# číselníkem obcí z RUIANu.

import argparse
from typing import List, Dict, DefaultDict
import sys
from pathlib import Path
import re
import csv
from collections import defaultdict

import mo.db as db
import mo.util

mo.util.init_standalone()
session = db.get_session()
new_town_cnt = 0
processed_school_cnt = 0
new_school_cnt = 0


def import_schools(path: Path, nuts: str):
    # XXX: Rejstřík škol používá několik chybných/obsoletních NUTS kódů :(
    nuts = re.sub('^CZ011', 'CZ010', nuts)
    nuts = re.sub('^CZ021', 'CZ020', nuts)
    nuts = re.sub('^CZ061', 'CZ063', nuts)
    nuts = re.sub('^CZ062', 'CZ064', nuts)
    nuts = re.sub('^CZ081', 'CZ080', nuts)

    with path.open('r') as file:
        columns = parse_header(file.readline())
        for line in file:
            f = line.split('\t')
            red_izo = f[columns['Red IZO']]
            ico = f[columns['IČO']]
            druh = f[columns['Druh školy/zařízení']]

            # Address of legal entity
            nazev = f[columns['Název']]
            misto = f[columns['Místo']]
            ulice = f[columns['Ulice']]
            cp = f[columns['Č.p.']]
            co = f[columns['Č.o.']]

            # Address of school building
            misto2 = f[columns['_Místo']]
            ulice2 = f[columns['_Ulice']]
            cp2 = f[columns['_Č.p.']]
            co2 = f[columns['_Č.o.']]

            addr = make_address(misto, ulice, cp, co)
            addr2 = make_address(misto2, ulice2, cp2, co2)
            # if addr != addr2:
            #     print(f"WARNING: Škola má dvě různé adresy: <{addr}> != <{addr2}>", file=sys.stderr)

            town = lookup_town(misto2, nuts)

            if druh.startswith('Základní škola'):
                is_zs = True
            elif druh.startswith('Střední škola'):
                is_zs = False
            else:
                assert False, f"Neznámý druh školy: {druh}"

            school = (session.query(db.School)
                      .join(db.Place)
                      .filter(db.Place.level == 4)
                      .filter(db.Place.parent == town.place_id)
                      .filter(db.School.red_izo == red_izo)
                      .filter(db.School.address == addr2)
                      .first())
            if school:
                assert school.official_name == nazev
                if is_zs:
                    school.is_zs = True
                else:
                    school.is_ss = True
            else:
                place = db.Place(
                    level=4,
                    parent=town.place_id,
                    name=nazev,
                    type=db.PlaceType.school)
                school = db.School(
                    place=place,
                    red_izo=red_izo,
                    ico=ico,
                    official_name=nazev,
                    address=addr2,
                    is_zs=is_zs,
                    is_ss=not is_zs)
                session.add(school)
                global new_school_cnt
                new_school_cnt += 1

            global processed_school_cnt
            processed_school_cnt += 1


def parse_header(header: str) -> Dict[str, int]:
    columns = {}
    i = 0
    for col in header.split('\t'):
        if col.endswith(':'):
            col = col[:-1]
        while col in columns:
            col = '_' + col
        columns[col] = i
        i += 1
    return columns


def make_address(misto: str, ulice: str, cp: str, co: str) -> str:
    if cp and co:
        c = f"{cp}/{co}"
    else:
        c = cp or co

    if ulice:
        if c:
            return f"{ulice} {c}, {misto}"
        else:
            return f"{ulice}, {misto}"
    else:
        return misto


def lookup_town(misto: str, region_nuts: str) -> db.Place:
    ruian_nuts = ruian_obec_to_okres_nuts[misto]
    region = None

    if region_nuts in ruian_nuts:
        nuts = region_nuts
    elif not ruian_nuts:
        if misto.startswith('Praha '):
            # XXX: Pražské obvody nejsou v RUIANu
            region = session.query(db.Place).filter_by(level=2, name=misto).first()
            assert region
        else:
            nuts = region_nuts
            print(f"WARNING: Obec {misto} není v RUIAN", file=sys.stderr)
    elif len(ruian_nuts) == 1:
        nuts = ruian_nuts[0]
        print(f"WARNING: Obec {misto} je podle rejstříku v okrese {region_nuts}, ale pod RUIAN v {nuts} => preferuji RUIAN", file=sys.stderr)
    else:
        nuts = region_nuts
        print(f"WARNING: Obec {misto} je podle rejstříku v okrese {region_nuts}, podle RUIAN je na výběr {ruian_nuts} => dořešit ručně!", file=sys.stderr)

    if not region:
        region = session.query(db.Place).filter_by(level=2, nuts=nuts).first()
        assert region

    town = session.query(db.Place).filter_by(level=3, parent=region.place_id, name=misto).first()
    if town is None:
        town = db.Place(level=3, parent=region.place_id, name=misto, type=db.PlaceType.region)
        session.add(town)
        session.flush()
        global new_town_cnt
        new_town_cnt += 1
    return town


def load_ruian_csv(name):
    with open(name) as file:
        reader = csv.reader(file, delimiter=';')
        rows = list(reader)
        columns = {}
        i = 0
        for h in rows[0]:
            columns[h] = i
            i += 1
        return columns, rows[1:]


ruian_obec_to_okres_nuts: DefaultDict[str, List[str]] = defaultdict(list)


def load_ruian():
    ocols, okresy = load_ruian_csv('extra/ruian/UI_OKRES.csv')
    okres_by_id: Dict[int, List[str]] = {}
    for o in okresy:
        id = int(o[ocols['KOD']])
        assert id not in okres_by_id
        okres_by_id[id] = o

    mcols, mesta = load_ruian_csv('extra/ruian/UI_OBEC.csv')
    for m in mesta:
        jmeno = m[mcols['NAZEV']]
        oid = int(m[mcols['OKRES_KOD']])
        okres = okres_by_id[oid]
        # print(f"{jmeno} -> {okres}")
        ruian_obec_to_okres_nuts[jmeno].append(okres[ocols['NUTS_LAU']])

parser = argparse.ArgumentParser(description='Importuje školy z naparsovaného Rejstříku škol')
parser.add_argument('-n', '--dry-run', default=False, action='store_true', help='pouze ukáže, co by bylo provedeno')

args = parser.parse_args()

load_ruian()

for path in Path('extra/skoly/parsed').glob('*.tsv'):
    m = re.fullmatch(r'^[A-Z]-(CZ\w+)\.tsv', path.name)
    assert m is not None
    nuts = m[1]
    import_schools(path, nuts)

if not args.dry_run:
    session.commit()
print(f"Importováno {processed_school_cnt} škol.")
print(f"Založeno {new_school_cnt} nových škol a {new_town_cnt} nových obcí.")