import uuid
import traceback
import logging
from inspect import signature
from datetime import datetime
from typing import Optional, Callable, Dict
from sqlalchemy.orm import Session
import sentry_sdk

from ..database import db_session, engine
from ..models import *  # pylint: disable=wildcard-import
from ..util.isoformat import isoformat
from ..util.jsonschema import JSONDict
from .. import config

logger = logging.getLogger("arlo.worker")


class UserError(Exception):
    pass


task_dispatch: Dict[str, Callable] = {}


# Decorator to register background task handlers. We use the handler's function
# name as the key.
def background_task(task_handler: Callable):
    task_dispatch[task_handler.__name__] = task_handler
    return task_handler


# All tasks should have election_id in the payload in order to easily identify their logs.
def create_background_task(
    task_handler: Callable, payload: JSONDict, db_session=db_session
) -> BackgroundTask:
    assert task_handler.__name__ in task_dispatch, (
        f"No task handler registered for {task_handler.__name__}."
        " Did you forget to use the @background_task decorator?"
    )

    task = BackgroundTask(
        id=str(uuid.uuid4()), task_name=task_handler.__name__, payload=payload
    )
    db_session.add(task)

    # For testing, we often prefer tasks to run immediately, instead of asynchronously.
    if config.RUN_BACKGROUND_TASKS_IMMEDIATELY:
        run_task(task)

    return task


def task_log_data(task: BackgroundTask) -> JSONDict:
    return dict(id=task.id, task_name=task.task_name, payload=task.payload)


def emit_progress_for_task(task_id: str):
    progress_session = Session(engine)

    def emit_progress(work_progress: int, work_total: int):
        task = progress_session.query(BackgroundTask).get(task_id)
        task.work_progress = work_progress
        task.work_total = work_total
        progress_session.commit()

    return emit_progress


# Currently, we assume that only one worker is consuming tasks at a time. There
# are no guards to prevent parallel workers from running the same task.
#
# Due to this constraint, functions in this file take an optional db_session
# argument in order for the tests to call them using isolated databases. In
# non-test environments, using the default global db_session is fine.
def run_task(task: BackgroundTask, db_session=db_session):
    task_handler = task_dispatch.get(task.task_name)
    assert task_handler, (
        f"No task handler registered for {task.task_name}."
        " Did you forget to use the @background_task decorator?"
    )

    logger.info(f"TASK_START {task_log_data(task)}")

    task.started_at = datetime.now(timezone.utc)

    db_session.commit()

    task_args = dict(task.payload)
    # Inject emit_progress for handlers that want to record task progress
    if "emit_progress" in signature(task_handler).parameters:
        task_args["emit_progress"] = emit_progress_for_task(task.id)

    try:
        task_handler(**task_args)

        task.completed_at = datetime.now(timezone.utc)

        db_session.commit()

        logger.info(f"TASK_COMPLETE {task_log_data(task)}")

    except Exception as error:
        db_session.rollback()

        task.completed_at = datetime.now(timezone.utc)

        # Some exceptions stringify nicely, some don't (e.g. StopIteration) so
        # we just print the exception class name.
        task.error = str(error) or str(error.__class__.__name__)

        db_session.commit()

        log_data = {
            **task_log_data(task),
            "error": task.error,
        }

        if isinstance(error, UserError):
            logger.info(f"TASK_USER_ERROR {log_data}")
        else:
            log_data["traceback"] = str(traceback.format_tb(error.__traceback__))
            logger.error(f"TASK_ERROR {log_data}")
            sentry_sdk.capture_exception(error)


def run_new_tasks(db_session=db_session):
    # Cleanup any tasks that failed to finish processing last time the worker was run
    stuck_tasks = (
        db_session.query(BackgroundTask)
        .filter(BackgroundTask.started_at.isnot(None))
        .filter(BackgroundTask.completed_at.is_(None))
        .all()
    )
    for task in stuck_tasks:
        task.started_at = None
        logger.info(f"TASK_RESET {task_log_data(task)}")

    db_session.commit()

    # Find and run new tasks
    for task in (
        db_session.query(BackgroundTask)
        .filter_by(started_at=None)
        .order_by(BackgroundTask.created_at)
        .all()
    ):
        run_task(task, db_session)


def serialize_background_task(task: Optional[BackgroundTask]) -> Optional[JSONDict]:
    if task is None:
        return None

    if task.error:
        status = ProcessingStatus.ERRORED
    elif task.completed_at:
        status = ProcessingStatus.PROCESSED
    elif task.started_at:
        status = ProcessingStatus.PROCESSING
    else:
        status = ProcessingStatus.READY_TO_PROCESS

    json: JSONDict = {
        "status": status,
        "startedAt": isoformat(task.started_at),
        "completedAt": isoformat(task.completed_at),
        "error": task.error,
    }

    if task.work_total is not None:
        json["workProgress"] = task.work_progress
        json["workTotal"] = task.work_total

    return json
