from abc import abstractmethod

import sqlalchemy as db
from dagster import check
from dagster.core.definitions.run_request import InstigatorType
from dagster.core.errors import DagsterInvariantViolationError
from dagster.core.scheduler.instigation import (
    InstigatorState,
    InstigatorTick,
    TickData,
    TickStatsSnapshot,
    TickStatus,
)
from dagster.serdes import deserialize_json_to_dagster_namedtuple, serialize_dagster_namedtuple
from dagster.utils import utc_datetime_from_timestamp

from .base import ScheduleStorage
from .schema import JobTable, JobTickTable


class SqlScheduleStorage(ScheduleStorage):
    """Base class for SQL backed schedule storage"""

    @abstractmethod
    def connect(self):
        """Context manager yielding a sqlalchemy.engine.Connection."""

    def execute(self, query):
        with self.connect() as conn:
            result_proxy = conn.execute(query)
            res = result_proxy.fetchall()
            result_proxy.close()

        return res

    def _deserialize_rows(self, rows):
        return list(map(lambda r: deserialize_json_to_dagster_namedtuple(r[0]), rows))

    def all_instigator_state(self, repository_origin_id=None, instigator_type=None):
        check.opt_inst_param(instigator_type, "instigator_type", InstigatorType)
        base_query = db.select([JobTable.c.job_body, JobTable.c.job_origin_id]).select_from(
            JobTable
        )

        if repository_origin_id:
            query = base_query.where(JobTable.c.repository_origin_id == repository_origin_id)
        else:
            query = base_query

        if instigator_type:
            query = query.where(JobTable.c.job_type == instigator_type.value)

        rows = self.execute(query)
        return self._deserialize_rows(rows)

    def get_instigator_state(self, origin_id):
        check.str_param(origin_id, "origin_id")

        query = (
            db.select([JobTable.c.job_body])
            .select_from(JobTable)
            .where(JobTable.c.job_origin_id == origin_id)
        )

        rows = self.execute(query)
        return self._deserialize_rows(rows[:1])[0] if len(rows) else None

    def add_instigator_state(self, state):
        check.inst_param(state, "state", InstigatorState)
        with self.connect() as conn:
            try:
                conn.execute(
                    JobTable.insert().values(  # pylint: disable=no-value-for-parameter
                        job_origin_id=state.instigator_origin_id,
                        repository_origin_id=state.repository_origin_id,
                        status=state.status.value,
                        job_type=state.instigator_type.value,
                        job_body=serialize_dagster_namedtuple(state),
                    )
                )
            except db.exc.IntegrityError as exc:
                raise DagsterInvariantViolationError(
                    f"InstigatorState {state.instigator_origin_id} is already present in storage"
                ) from exc

        return state

    def update_instigator_state(self, state):
        check.inst_param(state, "state", InstigatorState)
        if not self.get_instigator_state(state.instigator_origin_id):
            raise DagsterInvariantViolationError(
                "InstigatorState {id} is not present in storage".format(
                    id=state.instigator_origin_id
                )
            )

        with self.connect() as conn:
            conn.execute(
                JobTable.update()  # pylint: disable=no-value-for-parameter
                .where(JobTable.c.job_origin_id == state.instigator_origin_id)
                .values(
                    status=state.status.value,
                    job_body=serialize_dagster_namedtuple(state),
                )
            )

    def delete_instigator_state(self, origin_id):
        check.str_param(origin_id, "origin_id")

        if not self.get_instigator_state(origin_id):
            raise DagsterInvariantViolationError(
                "InstigatorState {id} is not present in storage".format(id=origin_id)
            )

        with self.connect() as conn:
            conn.execute(
                JobTable.delete().where(  # pylint: disable=no-value-for-parameter
                    JobTable.c.job_origin_id == origin_id
                )
            )

    def _add_filter_limit(self, query, before=None, after=None, limit=None):
        check.opt_float_param(before, "before")
        check.opt_float_param(after, "after")
        check.opt_int_param(limit, "limit")

        if before:
            query = query.where(JobTickTable.c.timestamp < utc_datetime_from_timestamp(before))
        if after:
            query = query.where(JobTickTable.c.timestamp > utc_datetime_from_timestamp(after))
        if limit:
            query = query.limit(limit)
        return query

    def get_ticks(self, origin_id, before=None, after=None, limit=None):
        check.str_param(origin_id, "origin_id")
        check.opt_float_param(before, "before")
        check.opt_float_param(after, "after")
        check.opt_int_param(limit, "limit")

        query = (
            db.select([JobTickTable.c.id, JobTickTable.c.tick_body])
            .select_from(JobTickTable)
            .where(JobTickTable.c.job_origin_id == origin_id)
            .order_by(JobTickTable.c.timestamp.desc())
        )

        query = self._add_filter_limit(query, before=before, after=after, limit=limit)

        rows = self.execute(query)
        return list(
            map(lambda r: InstigatorTick(r[0], deserialize_json_to_dagster_namedtuple(r[1])), rows)
        )

    def create_tick(self, tick_data):
        check.inst_param(tick_data, "tick_data", TickData)

        with self.connect() as conn:
            try:
                tick_insert = (
                    JobTickTable.insert().values(  # pylint: disable=no-value-for-parameter
                        job_origin_id=tick_data.instigator_origin_id,
                        status=tick_data.status.value,
                        type=tick_data.instigator_type.value,
                        timestamp=utc_datetime_from_timestamp(tick_data.timestamp),
                        tick_body=serialize_dagster_namedtuple(tick_data),
                    )
                )
                result = conn.execute(tick_insert)
                tick_id = result.inserted_primary_key[0]
                return InstigatorTick(tick_id, tick_data)
            except db.exc.IntegrityError as exc:
                raise DagsterInvariantViolationError(
                    f"Unable to insert InstigatorTick for job {tick_data.instigator_name} in storage"
                ) from exc

    def update_tick(self, tick):
        check.inst_param(tick, "tick", InstigatorTick)

        with self.connect() as conn:
            conn.execute(
                JobTickTable.update()  # pylint: disable=no-value-for-parameter
                .where(JobTickTable.c.id == tick.tick_id)
                .values(
                    status=tick.status.value,
                    type=tick.instigator_type.value,
                    timestamp=utc_datetime_from_timestamp(tick.timestamp),
                    tick_body=serialize_dagster_namedtuple(tick.tick_data),
                )
            )

        return tick

    def purge_ticks(self, origin_id, tick_status, before):
        check.str_param(origin_id, "origin_id")
        check.inst_param(tick_status, "tick_status", TickStatus)
        check.float_param(before, "before")

        utc_before = utc_datetime_from_timestamp(before)

        with self.connect() as conn:
            conn.execute(
                JobTickTable.delete()  # pylint: disable=no-value-for-parameter
                .where(JobTickTable.c.status == tick_status.value)
                .where(JobTickTable.c.timestamp < utc_before)
                .where(JobTickTable.c.job_origin_id == origin_id)
            )

    def get_tick_stats(self, origin_id):
        check.str_param(origin_id, "origin_id")

        query = (
            db.select([JobTickTable.c.status, db.func.count()])
            .select_from(JobTickTable)
            .where(JobTickTable.c.job_origin_id == origin_id)
            .group_by(JobTickTable.c.status)
        )

        rows = self.execute(query)

        counts = {}
        for status, count in rows:
            counts[status] = count

        return TickStatsSnapshot(
            ticks_started=counts.get(TickStatus.STARTED.value, 0),
            ticks_succeeded=counts.get(TickStatus.SUCCESS.value, 0),
            ticks_skipped=counts.get(TickStatus.SKIPPED.value, 0),
            ticks_failed=counts.get(TickStatus.FAILURE.value, 0),
        )

    def wipe(self):
        """Clears the schedule storage."""
        with self.connect() as conn:
            # https://stackoverflow.com/a/54386260/324449
            conn.execute(JobTable.delete())  # pylint: disable=no-value-for-parameter
            conn.execute(JobTickTable.delete())  # pylint: disable=no-value-for-parameter
