from typing import Any

from db_schema.schemas.apps import DatabricksExecution, DBXJobStatusEnum
from flask_admin.contrib.sqla.filters import BaseSQLAFilter
from sqlalchemy.orm import Query, Session

from delphi_slz_admin.modelviews import BaseFilterById, BaseModelView


class FilterUnitOfWorkId(BaseFilterById):
    def __init__(self, **kwargs: Any) -> None:
        self.column = DatabricksExecution.unit_of_work_id
        self.name = 'Unit of Work ID'
        super().__init__(self.column, self.name, **kwargs)


class FilterStatus(BaseSQLAFilter):
    def __init__(self, **kwargs: Any) -> None:
        self.column = DatabricksExecution.status
        self.name = 'Status'
        self.options = [(item.value, item.value) for item in DBXJobStatusEnum]
        super().__init__(self.column, self.name, options=self.options, **kwargs)

    def apply(self, query: Query, value: str, alias: Any = None) -> Query:
        return query.filter(self.column == value)

    def operation(self) -> str:
        return 'equals'


class AppsDbxExecutionModelView(BaseModelView):
    column_list = [
        'databricks_execution_id', 'unit_of_work_id', 'dbx_job_id', 'dbx_run_id',
        'sf_execution_name', 'status', 'started_at', 'completed_at', 'contexts'
    ]
    column_sortable_list = [
        'databricks_execution_id', 'unit_of_work_id', 'dbx_job_id', 'dbx_run_id',
        'sf_execution_name', 'status', 'started_at', 'completed_at'
    ]
    column_filters = [
        FilterUnitOfWorkId(),
        FilterStatus(),
    ]
    column_default_sort = [('started_at', True)]

    def __init__(self, session: Session, **kwargs: Any) -> None:
        self.model = DatabricksExecution
        self.name = 'Databricks Executions'
        self.endpoint = 'dbx_executions'
        super().__init__(self.model, session, name=self.name, endpoint=self.endpoint, **kwargs)