from __future__ import annotations

import logging
from pathlib import Path
from typing import List

import typer
from PIL import Image

from yarrtist.cli.globals import CONTEXT_SETTINGS, OPTIONS, LogLevel
from yarrtist.plotter.core_functions import (
    plot_fe_config,
    plot_module_histo,
)
from yarrtist.utils.utils import (
    create_img_grid,
    data_from_config,
    fill_pdf_summary,
    format_data,
    get_chip_types_configs_from_scanlog,
    get_configs_from_connectivity,
    get_geomId_sn_from_config,
    load_data,
)

logging.basicConfig(format="[%(levelname)s] %(message)s")
log = logging.getLogger("YARRtist")

app = typer.Typer(context_settings=CONTEXT_SETTINGS)


def config_summary_perchip(configs, absolute_path, summary_title):
    all_imgs = []

    for chip in [next(iter(c.values()))["config"] for c in configs.get("chips", [])]:
        images = []
        _, name = get_geomId_sn_from_config(chip)
        log.info(f"Plotting FE {name}")
        try:
            plotters = plot_fe_config(chip, f"{name}")
            for k in plotters:
                if k != "TDAC_1d":
                    fill_pdf_summary(plotters[k], images)
            chip_img = create_img_grid(images, (2, 2), f"{name}")
            all_imgs.append(chip_img)
            tdac_img = []
            fill_pdf_summary(plotters["TDAC_1d"], tdac_img)
            img = create_img_grid(tdac_img, (1, 1), "")
            all_imgs.append(img)
        except Exception as e:
            log.warning(f"Failed plotting: {e}")

    if all_imgs:
        summary = [p.convert("RGB") for p in all_imgs]
        summary[0].save(
            f"{absolute_path}/{summary_title}.pdf",
            save_all=True,
            append_images=all_imgs[1:],
        )
        log.info(f"Config summary saved in {absolute_path}/{summary_title}.pdf")
    else:
        log.info("Config summary is empty")


def config_summary(configs, module_type, absolute_path, summary_title):
    all_imgs = []

    all_ch_names = [next(iter(chip.keys())) for chip in configs.get("chips", [])]

    all_data = [
        data_from_config(next(iter(c.values()))["config"])
        for c in configs.get("chips", [])
    ]
    all_data = list(map(list, zip(*all_data)))

    plot = {}
    plot["data"] = {}

    for n in range(len(all_data)):
        for geomId, test in enumerate(all_data[n]):
            format_data(plot, test, geomId)
        try:
            plotter = plot_module_histo(plot, test, module_type, "", all_ch_names)
            fill_pdf_summary(plotter, all_imgs)
        except Exception as e:
            log.warning(f"Failed plotting: {e}")

    if all_imgs:
        summary = [Image.open(p).convert("RGB") for p in all_imgs]
        summary[0].save(
            f"{absolute_path}/{summary_title}.pdf",
            save_all=True,
            append_images=summary[1:],
        )
        log.info(f"Config summary saved in {absolute_path}/{summary_title}.pdf")
    else:
        log.info("Config summary is empty")


@app.command()
def main(
    connectivity_files: List[Path] = OPTIONS["connectivity_files"],
    scan_directory: Path = OPTIONS["scan_directory_conn_plot"],
    per_module: bool = OPTIONS["per_module"],
    relative_path: Path = OPTIONS["path"],
    verbosity: LogLevel = OPTIONS["verbosity"],
):
    log.setLevel(verbosity.value)

    if scan_directory:
        all_configs = get_chip_types_configs_from_scanlog(
            load_data(f"{scan_directory}/scanLog.json"), scan_directory, relative_path
        )
        all_paths = [scan_directory] * len(all_configs)
    elif connectivity_files:
        all_configs = []
        for conn_file in connectivity_files:
            all_configs.append(
                get_configs_from_connectivity(
                    load_data(conn_file), Path(conn_file).parent, relative_path
                )
            )
        all_paths = [Path(conn).parent for conn in connectivity_files]
    else:
        err_msg = "Please provide at least an option for the connectivity"
        raise typer.BadParameter(err_msg)

    log.info(f"Loaded {len(all_configs)} connectivity")

    for m in all_configs:
        chip_type = m.get("chipType")
        m["chips"] = []
        for c in m.get("configs"):
            full_config = load_data(f"{c}")
            geomId, name = get_geomId_sn_from_config(full_config, chip_type)
            m["chips"].append({f"{name}": {"id": geomId, "config": full_config}})
        if per_module:
            m["chips"] = sorted(m["chips"], key=lambda d: next(iter(d.values()))["id"])

    for i, (config, path) in enumerate(zip(all_configs, all_paths)):
        chip_type = m.get("chipType")
        if chip_type not in ["RD53B", "ITKPIXV2"]:
            log.warning(f"Plotting not available for the moment for {chip_type}")
            continue
        if not per_module:
            log.info(f"Plotting connectivity {i + 1}, divided per chip")
            config_summary_perchip(config, path, f"config_summary_perchip_{i}")
        else:
            log.info(f"Plotting connectivity of module {i + 1}")
            n_chips = len(config.get("chips", []))
            if n_chips not in [3, 4]:
                log.error("Number of chips does not match a known module type")
            else:
                module_type = "Triplet" if n_chips == 3 else "Quad"
                log.info(f"Module is a {module_type}")
                config_summary(config, module_type, path, f"config_summary_{i}")
