from __future__ import annotations

import logging
import shutil
import sys
import tempfile
import time
import warnings
from collections.abc import Generator
from collections.abc import Sequence
from pathlib import Path

import redis
from blissdata.beacon.data import BeaconData
from ewoksjob.client import submit
from ewoksutils.task_utils import task_inputs

from blissoda.utils.directories import get_workflows_dir

from ..import_utils import unavailable_class
from ..import_utils import unavailable_module
from ..version_utils import assert_has_minimal_version

try:
    import gevent
except ImportError as ex:
    gevent = unavailable_module(ex)

try:
    from bliss import setup_globals  # type: ignore
except ImportError as ex:
    setup_globals = unavailable_class(ex)

try:
    from bliss.scanning.chain import AcquisitionChain  # type: ignore
except ImportError as ex:
    AcquisitionChain = unavailable_class(ex)

try:
    from bliss.scanning.chain import ChainIterationPreset  # type: ignore
except ImportError as ex:
    ChainIterationPreset = unavailable_class(ex)

try:
    from bliss.scanning.chain import ChainPreset  # type: ignore
except ImportError as ex:
    ChainPreset = unavailable_class(ex)


logger = logging.getLogger(__name__)

_MOTOR_POLL_INTERVAL = 0.2  # s to wait for filter/atten motor move
_CURSOR_READ_TIMEOUT = 0.5  # s to wait for a new blissdata frame
_REDIS_XREAD_BLOCK_MS = 100  # ms to block waiting for stream events
_ID31_ATTENUATOR_BITS = 32
_COLORS = {
    "red": "\033[31m",
    "green": "\033[32m",
    "yellow": "\033[33m",
    "blue": "\033[34m",
    "reset": "\033[0m",
}


def _color(text: str, color: str) -> str:
    """Wrap text in ANSI color codes only when stderr is a TTY."""
    if not hasattr(sys.stderr, "isatty") or not sys.stderr.isatty():
        return text
    return f"{_COLORS.get(color, '')}{text}{_COLORS['reset']}"


def _timestamp() -> str:
    now = time.time()
    ms = int((now - int(now)) * 1000)
    s = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(now)) + ",%03d" % ms
    return s


def _log_atten_action(
    level: int,
    symbol: str,
    color: str,
    frame_nb: int,
    imax: float,
    threshold: float,
    comparator: str,
    atten_msg: str,
) -> None:
    """Emit a log message for insertion/removal of filters or attenuators."""
    logger.log(
        level,
        "[%s] %s #%2.0f max: %2.2e %s %2.2e, %s",
        _timestamp(),
        _color(symbol, color),
        frame_nb,
        imax,
        comparator,
        threshold,
        _color(atten_msg, color),
    )


class StopIntegrateSum(ChainPreset):  # pyright: ignore[reportGeneralTypeIssues]
    def __init__(
        self,
        workflow_threshold: float,
        workflow_path: str | Path,
        detector_name: str,
        detector_saturation: float,
        pyfai_config_path: str | Path,
        attenuator_name: str = "filt_wheel",
        frame_target_max: float | None = None,
        frame_target_min: float | None = None,
        ewoksjob_queue: str = "mt",
    ) -> None:
        """An acquisition chain preset for automatic attenuation and scan stopping in
        XRPD measurements.

        To be used with:

        - a continous scan mimicking detector accumulation, e.g
          ``timescan(0.2, npoints=0)``.
        - an Ewoks workflow containing an ``IntegrateAndSendMsg`` task inheriting from
          ``ewoksxrpd.tasks.integrate._BaseIntegrateSingleGeometry``, and sending
          a {'type':'stop'} event via the `blissoda:feedback:stop_task` stream upon
          reaching some chosen `workflow_threshold`.

        Attach this preset to a scan to:

        - insert / remove a filter automatically based on specified min and max
          frame counts;
        - stop the scan based on the output of an Ewoks workflow triggered at scan
          start, collecting scan statistics.

        :param workflow_threshold: The "statistic" computed by the Ewoks workflow, upon
            reaching which the workflow sends a signal to stop the scan.
        :param workflow_path: Path to the Ewoks workflow `.json`.
        :param detector_name: Name of the detector used in the scan, whose `.image`
            counter (i.e. raw frames) will be fed to the Ewoks workflow.
        :param detector_saturation: Max of detector dynamic range. Damage above.
        :param pyfai_config_path: Path to the PyFAI configuration `.json`.
        :param attenuator_name: Name of the filter or attenuator Bliss object to
            insert/remove dynamically.
        :param frame_target_max: Maximum number of desired counts in a frame.
            Select attenuation to get closest to this value, and above
            `frame_target_min`. If None, set to `0.9 * detector_saturation`.
        :param frame_target_min: Minimum number of desired counts in a frame.
            Select attenuation to be above this value, and closest to
            `frame_target_max`. If None, set to `0.1 * detector_saturation`.
        :param ewoksjob_queue: Name of the Celery queue to route the Ewoks tasks
            submitted by the preset to. Default: 'mt'.
        """
        assert_has_minimal_version("bliss", "2.2")

        self._detector_saturation = detector_saturation
        self._workflow_threshold = workflow_threshold
        self._target_max = (
            frame_target_max
            if frame_target_max is not None
            else detector_saturation * 0.9
        )
        self._target_min = (
            frame_target_min
            if frame_target_min is not None
            else detector_saturation * 0.1
        )

        try:
            self._filter = getattr(setup_globals, attenuator_name)
        except AttributeError as ex:
            raise RuntimeError(
                f"The attenuator {attenuator_name} is not available "
                f"in setup_globals: {ex}."
            )
        self._filter_name = attenuator_name

        self._filename = None
        self._scan_number = None
        self._scan_key = None
        self._detector_name = detector_name

        self.pyfai_config = pyfai_config_path
        if self.pyfai_config is None:
            raise NotImplementedError(
                "pyfai_config cannot be None: workflow needs detector geometry params."
                "Could be supplied via _integration_options. TBD."
            )
        self._energy = None
        self._workflow_path = workflow_path
        self._integration_options = dict(
            method=("no", "csr", "cython"),
            unit="2th_deg",
            integrator_name="sigma_clip_ng",
            extra_options={"max_iter": 5, "thres": 3},
            error_model="azimuthal",
            correct_solid_angle=True,
            polarization_factor=0.97,
        )

        self._redis_url = BeaconData().get_redis_data_db()
        self._redis = redis.Redis.from_url(self._redis_url)
        self._stream_name = "blissoda:feedback:stop_task"
        self._last_id = "0-0"

        self._future = None
        self._events_buffer = []
        self._greenlet = None
        self._cursor = None
        self._queue = ewoksjob_queue

        self._transmissions_cache: tuple[tuple, Sequence] | tuple[()] = ()

    def __repr__(self) -> str:
        return (
            "StopSumPreset\n"
            f"  Detector:     {self._detector_name}\n"
            f"  Attenuator:   {self._filter_name}\n"
            f"  PyFAI config: {self.pyfai_config}\n"
            f"  Integration options: (precedence over config)\n"
            + "\n".join(
                [f"     {k}: {v}" for k, v in self._integration_options.items()]
            )
        )

    @property
    def pyfai_integration_options(self) -> dict:
        return self._integration_options

    @pyfai_integration_options.setter
    def pyfai_integration_options(self, value: dict) -> None:
        self._integration_options = value

    @property
    def pyfai_config(self) -> str | Path | None:
        return self._pyfai_config

    @pyfai_config.setter
    def pyfai_config(self, value: str | Path) -> None:
        if value is None:
            self._pyfai_config = None
            return
        p = Path(value)
        if not p.exists():
            raise FileNotFoundError(f"pyFAI config file does not exist: {p}")
        self._pyfai_config = str(p.absolute())

    def prepare(self, chain: AcquisitionChain) -> None:  # type: ignore
        self._cache_scan_info(chain)

    def start(self, chain: AcquisitionChain) -> None:  # type: ignore
        self._start_workflow()

    def stop(self, chain: AcquisitionChain) -> None:  # type: ignore
        self._stop_workflow()

    def _cache_scan_info(self, chain: AcquisitionChain) -> None:  # type: ignore
        """Called at scan.prepare()."""
        self._scan = chain.scan
        self._filename = chain.scan.scan_info["filename"]
        self._scan_number = chain.scan.scan_info["scan_nb"]
        self._scan_key = chain.scan._scan_data.key
        self._energy = chain.scan.scan_info["positioners"]["positioners_start"][
            "energy"
        ]

    def _get_workflow_task_inputs(self) -> list[dict]:
        if not self._pyfai_config:
            warnings.warn("No pyFAI config file specified")

        # pyfai config task inputs
        # out: energy / geometry (dict) / detector name(str)
        # detector config (dict) / mask (array|None) / ff (array|none)
        # integration options
        inputs = task_inputs(
            task_identifier="PyFaiConfig",
            inputs=dict(
                energy=self._energy,  # priority 1
                filename=self._pyfai_config,  # priority 4
                integration_options=self._integration_options,  # priority 2
            ),
        )

        # integration task inputs
        # uses iter_bliss_data
        inputs += task_inputs(
            task_identifier="IntegrateAndSendMsg",
            inputs=dict(
                scan_key=self._scan_key,  # checked first
                filename=self._filename,  # fall back to fname/scan
                scan_number=self._scan_number,
                detector_name=self._detector_name,
                response_stream_name=self._stream_name,
                acc_profile_threshold=self._workflow_threshold,
                redis_url=self._redis_url,
                retry_timeout=1000,  # 1s
            ),
        )

        # write frame sum task inputs
        # reads frames from output loopscan and writes a single summed frame
        # in *_sum.h5, then fed to the xrpd workflow
        inputs += task_inputs(
            task_identifier="WriteSumFrame",
            inputs=dict(
                filename=self._filename,
                scan_number=self._scan_number,
                detector_name=self._detector_name,
            ),
        )

        return inputs

    def _resolve_workflow_path(self) -> str:
        """Ensure the workflow file is in the proposal's workflows directory.

        Follows the same pattern as the blissoda xrpd processor: copy the
        source workflow into ``<processed>/workflows/`` (derived from the
        dataset filename) so that the Ewoks worker can always reach it via
        an absolute path on the shared filesystem.
        """
        src = Path(self._workflow_path).absolute()
        if not src.is_file():
            raise FileNotFoundError(f"Ewoks workflow not found: {src}")

        if self._filename:
            dst_dir = Path(get_workflows_dir(self._filename))
        else:
            dst_dir = Path(tempfile.gettempdir())
        dst = dst_dir / src.name

        if src != dst and not dst.is_file():
            dst_dir.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(src, dst)

        return str(dst)

    def _start_workflow(self) -> None:
        """Called at scan.start(). Also starts _background_event_collector."""
        if not self._filename and not self._scan_key:
            raise RuntimeError("Neither scan key nor filename are known")

        # clear previous stream
        self._redis.delete(self._stream_name)
        self._last_id = "0-0"

        # submit workflow
        workflow = self._resolve_workflow_path()
        inputs = self._get_workflow_task_inputs()
        self._future = submit(
            args=(workflow,),
            kwargs=dict(
                inputs=inputs,
                convert_destination="last_stop_scan.json",
            ),
            queue=self._queue,
        )
        logger.info(
            "[%s] Workflow %s started for fname=%s key=%s, queue=%s",
            _timestamp(),
            self._future.uuid,
            self._filename,
            self._scan_key,
            self._queue,
        )

        # start background greenlet
        self._greenlet = gevent.spawn(self._background_event_collector)

    def _background_event_collector(self) -> None:
        """Continuously fetch and buffer feedback events."""
        while True:
            try:
                data = self._redis.xread(
                    {self._stream_name: self._last_id},
                    count=None,  # read all available
                    block=_REDIS_XREAD_BLOCK_MS,
                )

                if not data:
                    continue

                for _, events in data:
                    for event_id, event in events:
                        self._last_id = event_id
                        self._events_buffer.append(event)

            except Exception as ex:
                logger.exception("Error reading Redis stream: %s", ex)

    def _stop_workflow(self) -> None:
        """Called at scan.stop(). Only really stops _background_event_collector.
        The workflow needs to wait for scan.stop() to return in order to sum frames.
        """
        if self._greenlet is not None:
            self._greenlet.kill(block=True)
            self._greenlet = None

        # keep check only to catch exceptions: the WriteSumFrame task will
        # not fire until the data is flushed, which only happens at scan.stop()
        if self._future is None:
            return
        if self._future.done():
            try:
                self._future.result()
            except Exception as ex:
                logger.error("Stop scan workflow failed: %s", ex)

    def _get_filter_transmissions(self) -> Sequence:
        """Return a list of calculated transmission values for every filter or
        attenuator position.  Results are cached per preset instance and cleared
        when the energy or filter object is changed.
        """
        import numpy
        from bliss.common.auto_filter.filterset import FilterSet  # type: ignore
        from id31.attenuator import Attenuator  # pyright: ignore[reportMissingImports]
        from id31.attenuator import SiO2trans  # pyright: ignore[reportMissingImports]

        cache_key = (id(self._filter), self._energy)
        if self._transmissions_cache and self._transmissions_cache[0] == cache_key:
            return self._transmissions_cache[1]

        if isinstance(self._filter, FilterSet):
            result = [
                f.get("transmission_calc", None) for f in self._filter.get_filters()
            ]
        elif isinstance(self._filter, Attenuator):
            result = SiO2trans(self._energy, numpy.arange(_ID31_ATTENUATOR_BITS))
        else:
            raise RuntimeError(f"Filter {self._filter} is not supported.")

        self._transmissions_cache = (cache_key, result)
        return result

    def _find_best_filter_position(
        self, frame_max_intensity: float, current_atten_pos: int
    ) -> int | None:
        """Find the filter/attenuator position whose predicted intensity is closest
        to `_target_max` while staying below `_detector_saturation`.

        Returns `None` when the calculation cannot be performed
        (e.g. missing transmission data).
        """
        try:
            # get current transmission
            transmissions = self._get_filter_transmissions()
            t_current = transmissions[current_atten_pos]
            if t_current is None or t_current <= 0:
                return None

            # get unattenuated count rate
            count_rate = frame_max_intensity / t_current

            # find which transmission gives the cts closest to _target_max
            best_pos = None
            best_deviation = float("inf")
            for idx, t in enumerate(transmissions):
                if t is None or t < 0:
                    continue
                predicted = count_rate * t
                # only accept positions that give max cts below _detector_saturation
                if predicted < self._detector_saturation:
                    deviation = abs(predicted - self._target_max)
                    if deviation < best_deviation:
                        best_deviation = deviation
                        best_pos = idx
            return best_pos
        except Exception as ex:
            logger.warning("Cannot compute optimal filter position: %s", ex)
            return None

    def _predict_intensity_at(
        self, current_intensity_max: float, current_atten: int, target_atten: int
    ) -> float:
        """Predict the max pixel intensity at attenuator `target_atten` given a
        measurement at attenuator `current_atten`.  Returns `inf` when the prediction
        cannot be made, treated as saturation (conservative).
        """
        try:
            transmissions = self._get_filter_transmissions()
            t_current = transmissions[current_atten]
            t_target = transmissions[target_atten]
            if t_current and t_current > 0 and t_target is not None:
                return current_intensity_max * t_target / t_current
        except Exception:
            pass
        return float("inf")

    def _apply_filter(self, target_atten: int) -> None:
        """Move the filter/atten with blocking."""
        from bliss.common.auto_filter.filterset import FilterSet  # type: ignore

        try:
            if isinstance(self._filter, FilterSet):
                self._filter.set_filter(target_atten)
                while self._filter.rotation_axis.is_moving:
                    gevent.sleep(_MOTOR_POLL_INTERVAL)
            else:  # it's Attenuator
                self._filter.bits = target_atten
                # TODO proper way to block atten.bits?
                while self._filter.bits != target_atten:
                    gevent.sleep(_MOTOR_POLL_INTERVAL)
            logger.debug(
                "[%s] %s moved to position %d",
                _timestamp(),
                self._filter_name,
                target_atten,
            )
        except ValueError as err:
            logger.error(
                "[%s] Could not set %s to %d: %s",
                _timestamp(),
                self._filter_name,
                target_atten,
                err,
            )

    def get_iterator(
        self,
        chain: AcquisitionChain,  # type: ignore
    ) -> Generator[ChainIterationPreset, None, None]:  # type: ignore
        """Consume feedback events each iteration."""

        iter_nb = 0
        while True:
            yield ScanPointIterator(self, chain, iter_nb)
            iter_nb += 1


class ScanPointIterator(ChainIterationPreset):  # type: ignore
    def __init__(
        self,
        parent: StopIntegrateSum,
        chain: AcquisitionChain,  # type: ignore
        iter_nb: int,
    ):
        self.iter = iter_nb
        self.parent = parent
        self.chain = chain

    # executed synchronously
    def prepare(self) -> None:
        """Adjust attenuator/filter and stop scan if "stop" message received, at the
        preparation of each scan point. Skip the fist point.
        """
        if self.iter == 0:
            return

        self._adjust_filter()
        self._check_workflow_events()

    def start(self) -> None:
        """Start blissdata stream cursor at the start of the *first* scan point."""
        if self.iter == 0:
            stream = self.parent._scan.streams[f"{self.parent._detector_name}:image"]
            self.parent._cursor = stream.cursor()

    def _adjust_filter(self) -> None:
        """Read the frame acquired in the previous iter and adjust the
        attenuator/filter position.

        If:
        1. **Saturating** (max >= ``detector_saturation``): increase attenuation.
           Use transmission ratios to jump to the optimal position;
           fall back to +1.
        2. **Too dim** (max < ``frame_target_min``): decrease attenuation.
           Predict intensity at the target to avoid saturating;
           stay put if it would.
        3. **Acceptable range**: do nothing.
        """
        if self.parent._cursor is None:
            raise RuntimeError(
                "No blissdata Cursor present, did the first scan point run?"
            )

        view = self.parent._cursor.read(timeout=_CURSOR_READ_TIMEOUT, last_only=True)
        last_frame = view.get_data()
        imax_last_frame = float(last_frame.max())

        from bliss.common.auto_filter.filterset import FilterSet  # type: ignore

        if isinstance(self.parent._filter, FilterSet):
            current_atten = self.parent._filter.get_filter()
            n_filters = len(self.parent._filter.get_filters())
        else:  # Attenuator
            current_atten = self.parent._filter.bits
            n_filters = _ID31_ATTENUATOR_BITS
        max_atten = n_filters - 1

        # last frame was saturating: *increase* atten
        if (
            imax_last_frame >= self.parent._detector_saturation
            and current_atten < max_atten
        ):
            best = self.parent._find_best_filter_position(
                imax_last_frame, current_atten
            )
            if best is not None and best > current_atten:
                target_atten = best
            else:
                target_atten = current_atten + 1  # fallback: +1
            target_atten = min(target_atten, max_atten)

            _log_atten_action(
                logging.WARNING,
                "↓",
                "red",
                self.iter - 1,
                imax_last_frame,
                self.parent._detector_saturation,
                "≥",
                "atten: %d → %d" % (current_atten, target_atten),
            )
            self.parent._apply_filter(target_atten)

        # last frame was too dim: *decrease* atten
        elif imax_last_frame < self.parent._target_min and current_atten > 0:
            best = self.parent._find_best_filter_position(
                imax_last_frame, current_atten
            )
            if best is not None and best < current_atten:
                target_atten = best
            else:
                target_atten = current_atten - 1  # fallback: step down
            target_atten = max(target_atten, 0)

            # if removing atten would fry detector: do nothing
            predicted = self.parent._predict_intensity_at(
                imax_last_frame, current_atten, target_atten
            )
            if predicted >= self.parent._detector_saturation:
                _log_atten_action(
                    logging.WARNING,
                    "⨉",
                    "yellow",
                    self.iter - 1,
                    imax_last_frame,
                    self.parent._target_min,
                    "<",
                    "atten: %d (keep)" % current_atten,
                )
            else:
                _log_atten_action(
                    logging.WARNING,
                    "↑",
                    "blue",
                    self.iter - 1,
                    imax_last_frame,
                    self.parent._target_min,
                    "<",
                    "atten: %d → %d" % (current_atten, target_atten),
                )
                self.parent._apply_filter(target_atten)

        else:
            _log_atten_action(
                logging.INFO,
                "✓",
                "green",
                self.iter - 1,
                imax_last_frame,
                self.parent._target_max,
                "<",
                "atten: %d (in range)" % current_atten,
            )

    def _check_workflow_events(self) -> None:
        """Drain buffered workflow events and stop the scan if requested."""
        events_to_process, self.parent._events_buffer = (
            self.parent._events_buffer,
            [],
        )
        for event in events_to_process:
            if event.get(b"type") == b"stop":
                self.chain.scan.stop(wait=False)
                break
