#!/usr/bin/env python3

"""Corn jobs 🌽

Supporting the core app with asynchronous maintenance tasks.

Mostly related to monitoring and managing Printer state.
"""

from contextlib import closing
from datetime import datetime, timedelta
import logging
from pathlib import Path
from threading import Event, Thread
from time import sleep
from typing import Callable
from urllib import parse as urlparse

from octorest import OctoRest as _OR
from requests import Response
from requests.exceptions import (
    ConnectionError,
    HTTPError,
    Timeout,
)
from tentacles.store import Store


class OctoRest(_OR):
    def _get(self, path, params=None):
        url = urlparse.urljoin(self.url, path)
        response = self.session.get(url, params=params, timeout=(0.1, 1.0))
        self._check_response(response)

        return response.json()

    def _check_response(self, response: Response):
        response.raise_for_status()
        return response


log = logging.getLogger(__name__)

SHUTDOWN = Event()


def corn_job(every: timedelta):
    def _decorator(f):
        def _helper(*args, **kwargs):
            last = None
            while not SHUTDOWN.is_set():
                if not last or (datetime.now() - last) > every:
                    log.debug(f"Ticking job {f.__name__}")
                    try:
                        last = datetime.now()
                        f(*args, **kwargs)
                    except Exception:
                        log.exception(f"Error while procesing task {f.__name__}")
                else:
                    sleep(1)

        return _helper

    return _decorator


def poll_printers(db_factory: Callable[[], Store]) -> None:
    """Poll printers for their status."""

    with closing(db_factory()) as db:
        for printer in db.list_printers():
            id, name, url, api_key, last_poll, status = printer
            mapped_job = db.fetch_job_by_printer(id)
            try:
                client = OctoRest(url=url, apikey=api_key)
                printer_job = client.job_info()
                try:
                    printer_state = client.printer().get("state").get("flags", {})
                except HTTPError:
                    printer_state = {"disconnected": True}

                if printer_state.get("error"):
                    # If there's a mapped job, we manually fail it so that
                    # there's no possibility of a sync problem between the
                    # polling tasks. This violates separation of concerns a bit,
                    # but appears required for correctness.
                    if mapped_job:
                        db.finish_job(mapped_job.id, "error")

                    if status != "error":
                        print(f"Printer {printer.id} -> error")
                        db.update_printer_status(id, "error")

                elif printer_state.get("disconnected"):
                    if status != "disconnected":
                        print(f"Printer {printer.id} -> disconnected")
                        db.update_printer_status(id, "disconnected")

                elif printer_state.get("printing"):
                    if status != "running":
                        print(f"Printer {printer.id} -> running")
                        db.update_printer_status(id, "running")

                elif printer_job.get("state").lower() == "connecting":
                    if status != "connecting":
                        print(f"Printer {printer.id} -> connecting")
                        db.update_printer_status(id, "connecting")

                elif printer_state.get("ready"):
                    if status != "idle":
                        print(f"Printer {printer.id} -> idle")
                        db.update_printer_status(id, "idle")

                else:
                    raise Exception(f"Indeterminate state {printer_job!r} {printer_state!r}")

            except (ConnectionError, Timeout):
                db.update_printer_status(id, "error")

            except HTTPError as e:
                assert isinstance(e.response, Response)
                if e.response.status_code in [403, 401] or "error" in printer_job:
                    db.update_printer_status(id, "error")
                elif e.response.json().get("error") == "Printer is not operational":
                    db.update_printer_status(id, "disconnected")
                else:
                    raise Exception(
                        f"Indeterminate state {e.response.status_code}, {e.response.json()!r}"
                    )


def assign_jobs(db_factory: Callable[[], Store]) -> None:
    """Assign jobs to printers. Uploading files and job state management is handled separately."""

    with closing(db_factory()) as db:
        for printer_id in db.list_idle_printers():
            if job_id := db.poll_job_queue():
                db.assign_job(job_id, printer_id)
                print(f"Mapped job {job_id} to printer {printer_id}")


def push_jobs(db_factory: Callable[[], Store]) -> None:
    """Ensure that job files are uploaded and started to the assigned printer."""

    with closing(db_factory()) as db:
        for job in db.list_mapped_jobs():
            printer = db.fetch_printer(job.printer_id)
            file = db.fetch_file(job.user_id, job.file_id)
            if not file:
                log.error(f"Job {job.id} no longer maps to a file")
                db.delete_job(job.user_id, job.id)

            try:
                client = OctoRest(url=printer.url, apikey=printer.api_key)
                printer_job = client.job_info()
                try:
                    printer_state = client.printer().get("state").get("flags", {})
                except HTTPError:
                    printer_state = {"error": printer_job.get("error")}

                if printer_state.get("error"):
                    print(f"Printer {printer.id} is in error, can't push")
                    continue

                try:
                    client.upload(file.path)
                except HTTPError as e:
                    if e.response.status_code == 409:
                        pass
                    else:
                        raise

                client.select(Path(file.path).name)
                client.start()
                db.start_job(job.id)
            except TimeoutError:
                pass

            except Exception:
                log.exception("Oop")


def revoke_jobs(db_factory: Callable[[], Store]) -> None:
    """Ensure that job files are uploaded and started to the assigned printer.

    Note that this will ALSO cancel jobs out of the print queue.

    """

    with closing(db_factory()) as db:
        for job in db.list_cancelled_jobs():
            if job.printer_id:
                printer = db.fetch_printer(job.printer_id)
                try:
                    print(f"Cancelling running job {job.id}")
                    client = OctoRest(url=printer.url, apikey=printer.api_key)
                    try:
                        client.cancel()
                    except HTTPError as e:
                        if e.response.status_code == 409:
                            pass
                        else:
                            raise

                    print(f"Job {job.id} -> cancelled")
                    db.finish_job(job.id, "cancelled")
                except TimeoutError:
                    pass

                except Exception:
                    log.exception("Oop")

            else:
                print(f"Unmapped job {job.id} became cancelled")
                db.finish_job(job.id, "cancelled")



def pull_jobs(db_factory: Callable[[], Store]) -> None:
    """Poll the state of mapped printers to control jobs."""

    with closing(db_factory()) as db:
        for job in db.list_running_jobs():
            printer = db.fetch_printer(job.printer_id)
            try:
                client = OctoRest(url=printer.url, apikey=printer.api_key)
                job_state = client.job_info()
                try:
                    printer_state = client.printer().get("state").get("flags", {})
                except HTTPError:
                    printer_state = {"disconnected": True, "error": True}

                if printer_state.get("printing"):
                    pass

                elif job_state.get("progress", {}).get("completion", 0.0) == 100.0:
                    print(f"Job {job.id} has succeeded")
                    db.finish_job(job.id, "success")

                elif printer_state.get("error"):
                    print(f"Job {job.id} has failed")
                    db.finish_job(job.id, "failed")

                elif printer_state.get("cancelling"):
                    print(f"Job {job.id} has been acknowledged as cancelled")
                    db.finish_job(job.id, "cancelled")

                else:
                    print(f"Job {job.id} is in a weird state {job_state.get('progress')!r} {printer_state!r}")

            except TimeoutError:
                pass

            except Exception:
                log.exception("Oop")


@corn_job(timedelta(seconds=5))
def run_worker(db_factory):
    poll_printers(db_factory)
    assign_jobs(db_factory)
    push_jobs(db_factory)
    revoke_jobs(db_factory)
    pull_jobs(db_factory)

def create_workers(db_factory: Callable[[], Store]) -> Event:
    Thread(target=run_worker, args=[db_factory]).start()

    return SHUTDOWN
