import logging
from pathlib import Path
from typing import Optional
from typing import Sequence

from ewoks import BaseInputModel
from ewoks import Task

from .. import hdf5
from ..fit import CDI_align_estimate_COR
from ..plot import plot_CDI_align
from ..utils import logger


class Inputs(BaseInputModel):
    filename: str
    scan_nos: Sequence[int]
    axis_name: str
    rc_axis_name: str
    counter: str
    norm: str = "log"
    output_folder: str
    retry_timeout: Optional[int] = None
    retry_period: Optional[int] = None


class CdiPlot(Task, input_model=Inputs, output_names=["saved_figs", "saved_hdf5"]):
    def run(self):
        inputs = Inputs(**self.get_input_values())

        logger.setLevel(logging.INFO)
        handler = logging.FileHandler(Path(inputs.output_folder) / "cdi_plot.log")
        logger.addHandler(handler)

        try:
            xs, intensities, intensity_fits, labels, offset_x, offset_y, offset_fit = (
                CDI_align_estimate_COR(
                    inputs.filename,
                    inputs.scan_nos,
                    inputs.axis_name,
                    inputs.rc_axis_name,
                    inputs.counter,
                    retry_timeout=inputs.retry_timeout,
                    retry_period=inputs.retry_period,
                )
            )
        finally:
            logger.removeHandler(handler)
            handler.close()

        motor_pos_fig, offset_fig = plot_CDI_align(
            xs,
            intensities,
            intensity_fits,
            labels,
            offset_x,
            offset_y,
            offset_fit,
            inputs.axis_name,
            inputs.counter,
            norm=inputs.norm,
        )

        suffix = f"{inputs.scan_nos[0]}-{inputs.scan_nos[-1]}"

        output_filename = hdf5.save_fits(
            Path(inputs.output_folder) / f"{Path(inputs.filename).stem}.h5",
            suffix,
            xs,
            intensities,
            intensity_fits,
            labels,
            offset_x,
            offset_y,
            offset_fit,
            inputs.model_dump(),
        )
        self.outputs.saved_hdf5 = f"{output_filename}::/{suffix}"

        # Save fig for now (to be replaced by push to the logbook)
        motor_pos_fig_path = str(Path(inputs.output_folder) / f"motor_pos_{suffix}.png")
        motor_pos_fig.savefig(motor_pos_fig_path)
        offset_fig_path = str(Path(inputs.output_folder) / f"offset_{suffix}.png")
        offset_fig.savefig(offset_fig_path)

        self.outputs.saved_figs = (motor_pos_fig_path, offset_fig_path)
