#!/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
from functools import cache
import logging
import os
from pathlib import Path
from pprint import pformat
from typing import Callable
from urllib import parse as urlparse

from cherrypy.process.plugins import Monitor
from fastmail import FastMailSMTP
from flask import Flask as App, render_template
from gcode import analyze_gcode_file
from octorest import OctoRest as _OR
from requests import Response
from requests.exceptions import (
    ConnectionError,
    HTTPError,
    Timeout,
)
from tentacles.db import Db


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

        return response.json()

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

    def files_info(self, *args, **kwargs):
        try:
            return super().files_info(*args, **kwargs)
        except HTTPError as e:
            if e.response.status_code == 404:
                return {}
            else:
                raise e


@cache
def get_client(url, key):
    return OctoRest(url=url, apikey=key)


log = logging.getLogger(__name__)


def poll_printers(app: App, db: Db) -> None:
    """Poll printers for their status."""

    for printer in db.list_printers():
        mapped_job = db.fetch_job_by_printer(pid=printer.id)

        def _set_status(status: str):
            if printer.status != status:
                log.info(f"Printer {printer.id} {printer.status} -> {status}")
                db.update_printer_status(pid=printer.id, status=status)

        def _bed_clear():
            if not (
                snapshots := client._post(
                    "/api/plugin/bedready", json={"command": "list_snapshots"}
                )
            ):
                return True  # Assume the bed is ready

            status: dict = client._post(
                "/api/plugin/bedready",
                json={
                    "command": "check_bed",
                    "reference": snapshots[0],
                },
            )

            if "error" in status:
                log.error(f"Printer {printer.id} failed to report bed readiness with {status!r}")
                return False
            
            similarity = status.get("similarity", 0.0) ** 3 # Cube it
            return similarity > 0.9

        printer_job = {}
        try:
            client = get_client(printer.url, printer.api_key)
            printer_job: dict = client.job_info()
            try:
                printer_state: dict = client.printer().get("state").get("flags", {})
            except HTTPError:
                printer_state: dict = {"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(jid=mapped_job.id, state="error")

                _set_status("error")

            elif printer_state.get("disconnected"):
                _set_status("disconnected")

            elif printer_state.get("printing"):
                _set_status("running")
                if mapped_job:
                    db.start_job(jid=mapped_job.id)

            elif printer_job.get("state", "").lower() == "connecting":
                _set_status("connecting")

            elif printer_state.get("ready"):
                if _bed_clear():
                    _set_status("idle")
                else:
                    _set_status("occupied")

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

        except (ConnectionError, Timeout):
            _set_status("error")

        except HTTPError as e:
            assert isinstance(e.response, Response)
            if e.response.status_code in [403, 401] or "error" in printer_job:
                _set_status("error")

            elif e.response.json().get("error") == "Printer is not operational":
                _set_status("disconnected")

            else:
                raise Exception(
                    f"Indeterminate state {e.response.status_code}, {e.response.json()!r}"
                )


def assign_jobs(app: App, db: Db) -> None:
    """Assign jobs to printers. Uploading files and job state management is handled separately."""

    for job in db.list_job_queue(uid=None):
        # FIXME: Jobs which have been mapped are still in the "queue" until they finish
        # Ignore such as they have already been mapped
        if job.printer_id:
            continue

        idle = list(db.list_idle_printers())
        for printer in idle:
            if (
                job.analysis_id is not None
                and printer.limit_x >= job.max_x
                and printer.limit_y >= job.max_y
                and printer.limit_z >= job.max_z
                and printer.limit_hotend >= job.max_end
                and printer.limit_bed >= job.max_bed
                and printer.nozzle_diameter == job.nozzle_diameter
                and printer.filament_id == job.filament_id
            ):
                db.assign_job(jid=job.id, pid=printer.id)
                log.info(f"Mapped job {job.id} to printer {printer.id}")
                break
        else:
            if not idle:
                log.info(f"Could not map job {job!r}")


def push_jobs(app: App, db: Db) -> None:
    """Ensure that job files are uploaded and started to the assigned printer."""

    for job in db.list_mapped_jobs():
        printer = db.fetch_printer(pid=job.printer_id)
        file = db.fetch_file(uid=job.user_id, fid=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 = get_client(printer.url, 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"):
                log.warn(f"Printer {printer.id} is in error, can't push")
                continue

            try:
                if not client.files_info("local", Path(file.path).name):
                    client.upload(file.path)
                else:
                    log.info("Don't need to upload the job!")

            except HTTPError as e:
                if e.response.status_code == 409:
                    pass
                else:
                    raise

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

        except Exception:
            log.exception("Oop")


def revoke_jobs(app: App, db: Db) -> 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.

    """

    for job in db.list_canceling_jobs():
        if job.printer_id:
            printer = db.fetch_printer(pid=job.printer_id)

            log.info(f"Cancelling running job {job.id}")
            client = get_client(printer.url, printer.api_key)
            try:
                client.cancel()
            except HTTPError as e:
                if e.response.status_code == 409:
                    pass
                else:
                    raise

            log.info(f"Job {job.id} -> cancelled")
            db.finish_job(jid=job.id, state="cancelled")

        else:
            log.info(f"Unmapped job {job.id} became cancelled")
            db.finish_job(jid=job.id, state="cancelled")


def pull_jobs(app: App, db: Db) -> None:
    """Poll the state of mapped printers to control jobs."""

    for job in db.list_running_jobs():
        printer = db.fetch_printer(pid=job.printer_id)
        client = get_client(printer.url, 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 job_state.get("progress", {}).get("completion", 0.0) == 100.0:
            # Debounce jobs which JUST started from immediately completing due to a data race
            start_date = datetime.fromisoformat(job.started_at)
            runtime = datetime.utcnow() - start_date
            if runtime >= timedelta(seconds=15):  # 3 polling cycles
                log.info(f"Job {job.id} has succeeded")
                db.finish_job(jid=job.id, state="success")

        elif printer_state.get("error"):
            log.warn(f"Job {job.id} has failed")
            db.finish_job(jid=job.id, state="failed")

        elif printer_state.get("cancelling"):
            log.info(f"Job {job.id} has been acknowledged as cancelled")
            db.finish_job(jid=job.id, state="cancelled")

        elif printer_state.get("printing"):
            db.start_job(jid=job.id)

        else:
            log.warn(
                f"Job {job.id} is in a weird state; assuming error\n{pformat(job_state.get('progress'))}\n{pformat(printer_state)}"
            )
            db.finish_job(jid=job.id, state="failed")


def send_emails(app, db: Db):
    with closing(
        FastMailSMTP(
            app.config.get("fastmail", {}).get("username"),
            app.config.get("fastmail", {}).get("key"),
        )
    ) as fm:
        for message in db.poll_email_queue():
            fm.send_message(
                from_addr="root@tirefireind.us",
                to_addrs=[message.to],
                subject=message.subject,
                msg=message.body,
            )
            db.send_email(eid=message.id)


def analyze_files(app: App, db: Db):
    for file in db.list_unanalyzed_files():
        p = Path(file.path)
        if not p.is_file():
            log.error(f"Deleting missing file {file.id}!")
            db.delete_file(uid=file.user_id, fid=file.id)
            continue

        record = analyze_gcode_file(p)
        if not record:
            log.error(
                f"Unable to analyze {file.path} ({file.filename} owned by {file.user_id})!"
            )

            db.delete_file(uid=file.user_id, fid=file.id)
            db.delete_file_analysis_by_fid(fid=file.id)
            db.delete_jobs_by_fid(uid=file.user_id, fid=file.id)

            if os.path.exists(file.path):
                os.unlink(file.path)

            with app.app_context():
                db.create_email(
                    uid=file.user_id,
                    subject=f"Failure to print: {file.filename}",
                    body=render_template(
                        "analysis_failure_email.html.j2",
                        username=db.fetch_user(uid=file.user_id).name,
                        filename=file.filename,
                    ),
                )

            continue

        log.info(f"Analyzed {file}")
        db.create_analysis(
            file_id=file.id,
            max_x=record.max_x,
            max_y=record.max_y,
            max_z=record.max_z,
            max_end=record.max_end,
            max_bed=record.max_bed,
            nozzle=record.nozzle,
            filament_id=db.create_filament(name=record.filament).id,
        )


def debug_queue(app: App, db: Db):
    output = ["---"]

    for job in db.list_job_queue(uid=None):
        output.append("Job " + repr(job))

    for printer in db.list_idle_printers():
        output.append("Printer " + repr(printer))

    for unanalyzed in db.list_unanalyzed_files():
        output.append("Unanalyzed file " + repr(unanalyzed))

    print("\n".join(output))


def toil(*fs):
    def _helper(*args, **kwargs):
        for f in fs:
            f(*args, **kwargs)

    _helper.__name__ = "toil"
    return _helper


class Worker(Monitor):
    def __init__(
        self,
        bus,
        app: App,
        db_factory: Callable[[App], Db],
        callback: Callable[[App, Db], None],
        **kwargs,
    ):
        self._app = app
        self._db_factory = db_factory
        self._callback = callback
        super().__init__(
            bus, self.callback, **kwargs, name=f"Async {callback.__name__}"
        )

    def callback(self):
        try:
            with closing(self._db_factory(self._app)) as db:
                self._callback(self._app, db)
        except:
            log.exception("Worker tick failed; exception swallowed")
