import os
from collections import defaultdict
from functools import lru_cache
from typing import Dict, List, Optional, Tuple, Type

from lxml import etree
from lxml.etree import Element
from scenariogeneration import xosc

from scenario_gym.catalog_entry import Catalog, CatalogEntry
from scenario_gym.entity import DEFAULT_ENTITY_TYPES, Entity


def load_object(
    entry: Element,
    catalog: Optional[Catalog] = None,
    entity_types: Tuple[Type[Entity]] = DEFAULT_ENTITY_TYPES,
    catalog_objects: Optional[List[Type[CatalogEntry]]] = None,
) -> Optional[Entity]:
    """Try to load a catalog entry with given catalog objects."""
    if catalog_objects is None:
        catalog_objects = [Ent._catalog_entry_type() for Ent in entity_types]
    for Ent, Obj in zip(entity_types, catalog_objects):
        types = Obj.xosc_names if Obj.xosc_names is not None else [Obj.__name__]
        if entry.tag in types:
            obj = Obj.from_xml(entry, catalog=catalog)
            return Ent(obj)


@lru_cache(maxsize=None)
def read_catalog(
    catalog_file: str,
    entity_types: Optional[Tuple[Type[Entity]]] = None,
) -> Tuple[Catalog, Dict[str, Entity]]:
    """
    Read a catalog and return it's name and a dictionary of entities.

    Parameters
    ----------
    catalog_file : str
        Filepath of the catalog file.

    entity_types : Optional[Tuple[Type[CatalogObject]]]
        Tuple of extra subclasses of CatalogObject that will be used when reading
        catalogs. Can be used for reading custom objects from catalog files. Must
        be immutable or lru_cache

    Returns
    -------
    catalog : Catalog
        The catalog.

    entries : Dict[str, Entity]
        A dictionary mapping each entry name to an entity with that catalog entry.
        These can then be cloned to create entities using the chosen catalog entry.

    """
    if entity_types is None:
        entity_types = DEFAULT_ENTITY_TYPES
    else:
        entity_types = entity_types + DEFAULT_ENTITY_TYPES

    catalog_objects = [Ent._catalog_entry_type() for Ent in entity_types]

    et = etree.parse(catalog_file)
    osc_root = et.getroot()
    catalog_element = osc_root.find("Catalog")
    try:
        catalog_group_name = catalog_file.split(os.sep)[-3]
    except IndexError:
        catalog_group_name = "Catalog"
    catalog = Catalog(catalog_element.attrib["name"], catalog_group_name)
    entries = {}
    for element in catalog_element.getchildren():
        entry = load_object(
            element,
            catalog=catalog,
            entity_types=entity_types,
            catalog_objects=catalog_objects,
        )
        if entry is None:
            entry = Entity(CatalogEntry.from_xml(element, catalog=catalog))
        entries[entry.catalog_entry.catalog_entry] = entry
    return catalog, entries


def write_catalogs(dirname: str, entries: List[CatalogEntry]) -> None:
    """Create catalog files for a list of catalog entries."""
    c_to_e = defaultdict(list)
    for entry in entries:
        c_to_e[entry.catalog].append(entry)

    group_to_c = defaultdict(list)
    for catalog in c_to_e:
        group_to_c[catalog.group_name].append(catalog)

    for group_name, catalogs in group_to_c.items():
        os.mkdir(os.path.join(dirname, group_name))
        for catalog in catalogs:
            c_entries = c_to_e[catalog]
            catalog_type = f"{c_entries[0].catalog_type}Catalogs"
            catalog_dir = os.path.join(dirname, group_name, catalog_type)
            if not os.path.exists(catalog_dir):
                os.mkdir(catalog_dir)

            catalog_file = os.path.join(catalog_dir, f"{catalog.name}.xosc")

            catalog_obj = xosc.CatalogFile()
            catalog_obj.create_catalog(
                catalog_file,
                catalog.name,
                (
                    f"This {c_entries[0].catalog_type}.lower() catalog was "
                    "generated by scenario_gym."
                ),
                "scenario_gym",
            )

            for entry in c_entries:
                catalog_obj.add_to_catalog(entry.to_xosc())

            catalog_obj.dump()
