import logging
import time
import traceback
import typing as t

from sarus_data_spec import typing as st
from sarus_data_spec.bounds import Bounds
from sarus_data_spec.constants import BOUNDS_TASK
from sarus_data_spec.manager.cache_utils import lru_caching
from sarus_data_spec.manager.computations.local.base import LocalComputation
from sarus_data_spec.manager.ops.processor.routing import TransformedDataset
from sarus_data_spec.status import (
    DataSpecErrorStatus,
    error,
    processing,
    ready,
)

logger = logging.getLogger(__name__)


class BoundsComputation(LocalComputation[st.Bounds]):
    """Class responsible to compute bounds"""

    task_name = BOUNDS_TASK

    async def prepare(self, dataspec: st.DataSpec) -> None:
        dataspec = t.cast(st.Dataset, dataspec)
        processing(
            dataspec=dataspec,
            manager=self.computing_manager(),
            task=BOUNDS_TASK,
        )

        try:
            logger.info(f"STARTED BOUNDS {dataspec.uuid()}")
            start = time.perf_counter()
            bounds = await TransformedDataset(dataset=dataspec).bounds()

        except DataSpecErrorStatus as exception:
            error(
                dataspec=dataspec,
                manager=self.computing_manager(),
                task=self.task_name,
                properties={
                    "message": traceback.format_exc(),
                    "relaunch": str(exception.relaunch),
                },
            )

            raise

        except Exception:
            error(
                dataspec=dataspec,
                manager=self.computing_manager(),
                task=self.task_name,
                properties={
                    "message": traceback.format_exc(),
                    "relaunch": str(False),
                },
            )

            raise DataSpecErrorStatus((False, traceback.format_exc()))
        else:
            end = time.perf_counter()
            logger.info(
                f"FINISHED BOUNDS {dataspec.uuid()} ({end-start:.2f}s)"
            )
            ready(
                dataspec=dataspec,
                manager=self.computing_manager(),
                task=BOUNDS_TASK,
                properties={"uuid": bounds.uuid()},
            )

    async def result_from_stage_properties(
        self,
        dataspec: st.DataSpec,
        properties: t.Mapping[str, str],
        **kwargs: t.Any,
    ) -> st.Bounds:
        return t.cast(
            Bounds,
            dataspec.storage().referrable(properties["uuid"]),
        )

    @lru_caching("computation", use_first_arg=True)
    async def task_result(
        self, dataspec: st.DataSpec, **kwargs: t.Any
    ) -> st.Bounds:
        return await super().task_result(dataspec, **kwargs)