from __future__ import annotations

import datetime
import glob
import logging
import os
import subprocess
import threading
import time
import traceback
from datetime import timezone
from multiprocessing import Manager, Process
from typing import TYPE_CHECKING

from jobflow import JobStore, initialize_logger
from jobflow.core.flow import get_flow
from monty.design_patterns import singleton
from monty.os import cd
from monty.serialization import dumpfn, loadfn
from monty.shutil import decompress_file

from jobflow_remote.jobs.batch import LocalBatchManager
from jobflow_remote.jobs.data import (
    BATCH_INFO_FILENAME,
    IN_FILENAME,
    OUT_FILENAME,
    JobDoc,
)
from jobflow_remote.remote.data import get_job_path, get_store_file_paths
from jobflow_remote.utils.log import initialize_remote_run_log

if TYPE_CHECKING:
    from pathlib import Path

    from jobflow.core.job import Job

logger = logging.getLogger(__name__)


@singleton
class JfrState:
    """State of the current job being executed."""

    job_doc: JobDoc = None

    def reset(self):
        """Reset the current state."""
        self.job_doc = None


CURRENT_JOBDOC: JfrState = JfrState()


def run_remote_job(run_dir: str | Path = ".") -> None:
    """Run the job."""
    initialize_remote_run_log()

    start_time = datetime.datetime.now(timezone.utc)
    with cd(run_dir):
        error = None
        try:
            dumpfn({"start_time": start_time}, OUT_FILENAME)
            in_data = loadfn(IN_FILENAME)

            job: Job = in_data["job"]
            store = in_data["store"]
            job_doc_dict = in_data.get("job_doc", None)
            if isinstance(job.function, dict):
                raise RuntimeError(  # noqa: TRY004,TRY301
                    f"The function in the Job could not be deserialized: {job.function}.\n"
                    "Check if this function is actually available in the worker's python environment"
                )
            if job_doc_dict:
                job_doc_dict["job"] = job
                JfrState().job_doc = JobDoc.model_validate(job_doc_dict)

            store.connect()

            initialize_logger()
            try:
                response = job.run(store=store)
            finally:
                # some jobs may have compressed the jfremote and store files while being
                # executed, try to decompress them if that is the case and files need to be
                # decompressed.
                decompress_files(store)

            # Close the store explicitly, as minimal stores may require it.
            try:
                store.close()
            except Exception:
                logger.exception("Error while closing the store")

            # The output of the response has already been stored in the store.
            response.output = None

            # Convert to Flow the dynamic responses before dumping the output.
            # This is required so that the response does not need to be
            # deserialized and converted to Flows by the runner.
            if response.addition:
                response.addition = get_flow(response.addition)
            if response.detour:
                response.detour = get_flow(response.detour)
            if response.replace:
                response.replace = get_flow(response.replace)

            output = {
                "response": response,
                "error": error,
                "start_time": start_time,
                "end_time": datetime.datetime.now(timezone.utc),
            }
            dumpfn(output, OUT_FILENAME)
        except Exception:
            # replicate the dump to catch potential errors in
            # serializing/dumping the response.
            error = traceback.format_exc()
            output = {
                "response": None,
                "error": error,
                "start_time": start_time,
                "end_time": datetime.datetime.now(timezone.utc),
            }
            dumpfn(output, OUT_FILENAME)
        finally:
            JfrState().reset()


def ping(start_time, interval=600, filename=BATCH_INFO_FILENAME):
    while True:
        dumpfn(
            {
                "start_time": start_time,
                "last_ping_time": datetime.datetime.now(timezone.utc),
            },
            fn=filename,
        )
        time.sleep(interval)


PING_TIME = 600


def run_batch_jobs(
    base_run_dir: str | Path,
    files_dir: str | Path,
    batch_uid: str,
    max_time: float | None = None,
    max_wait: float = 60,
    max_jobs: int | None = None,
    parallel_jobs: int | None = None,
    sleep_time: float = None,
    batch_info_fname: str | Path = BATCH_INFO_FILENAME,
) -> None:
    # Here we assume that we are in the batch work directory where a batch process is executed/submitted
    start_time = datetime.datetime.now(timezone.utc)
    dumpfn({"start_time": start_time}, batch_info_fname)

    threading.Thread(
        target=ping,
        args=(start_time, PING_TIME, batch_info_fname),  # dump every 600 seconds
        daemon=True,
    ).start()

    parallel_jobs = parallel_jobs or 1

    if parallel_jobs == 1:
        run_single_batch_jobs(
            base_run_dir=base_run_dir,
            files_dir=files_dir,
            batch_uid=batch_uid,
            max_time=max_time,
            max_wait=max_wait,
            max_jobs=max_jobs,
            sleep_time=sleep_time,
        )
    else:
        with Manager() as manager:
            multiprocess_lock = manager.Lock()
            parallel_ids = manager.dict()
            batch_manager = LocalBatchManager(
                files_dir=files_dir,
                batch_uid=batch_uid,
                multiprocess_lock=multiprocess_lock,
            )
            processes = [
                Process(
                    target=run_single_batch_jobs,
                    args=(
                        base_run_dir,
                        files_dir,
                        batch_uid,
                        max_time,
                        max_wait,
                        max_jobs,
                        batch_manager,
                        parallel_ids,
                        sleep_time,
                    ),
                )
                for _ in range(parallel_jobs)
            ]
            for p in processes:
                p.start()
                time.sleep(0.5)

            for p in processes:
                p.join()
    dumpfn(
        {
            "start_time": start_time,
            "last_ping_time": datetime.datetime.now(timezone.utc),
            "end_time": datetime.datetime.now(timezone.utc),
        },
        batch_info_fname,
    )


def run_single_batch_jobs(
    base_run_dir: str | Path,
    files_dir: str | Path,
    batch_uid: str,
    max_time: float | None = None,
    max_wait: float = 60,
    max_jobs: int | None = None,
    batch_manager: LocalBatchManager | None = None,
    parallel_ids: dict | None = None,
    sleep_time: float = None,
) -> None:
    initialize_remote_run_log()

    # TODO the ID should be somehow linked to the queue job
    if not batch_manager:
        batch_manager = LocalBatchManager(files_dir=files_dir, batch_uid=batch_uid)

    if parallel_ids:
        parallel_ids[os.getpid()] = False

    t0 = time.time()
    wait = 0.0
    sleep_time = sleep_time or 10.0
    count = 0
    while True:
        if max_time and max_time < time.time() - t0:
            logger.info("Stopping due to max_time")
            return

        if max_wait and wait > max_wait:
            # if many jobs run in parallel do not shut down here, unless all
            # the other jobs are also stopped
            if parallel_ids:
                for pid, pid_is_running in parallel_ids.items():
                    if pid_is_running:
                        try:
                            os.kill(pid, 0)  # throws OSError if the process is dead
                        except OSError:  # means this process is dead!
                            parallel_ids[pid] = False
                if not any(parallel_ids.values()):
                    logger.info(
                        f"No jobs available for more than {max_wait} seconds and all other jobs are stopped. Stopping."
                    )
                    return
            else:
                logger.info(
                    f"No jobs available for more than {max_wait} seconds. Stopping."
                )
                return

        if max_jobs and count >= max_jobs:
            logger.info(f"Maximum number of jobs reached ({max_jobs}). Stopping.")
            return

        job_str = batch_manager.get_job()
        if not job_str:
            time.sleep(sleep_time)
            wait += sleep_time
        else:
            wait = 0.0
            count += 1
            job_id, _index = job_str.split("_")
            index: int = int(_index)
            logger.info(f"Starting job with id {job_id} and index {index}")
            job_path = get_job_path(job_id=job_id, index=index, base_path=base_run_dir)
            if parallel_ids:
                parallel_ids[os.getpid()] = True
            try:
                with cd(job_path):
                    result = subprocess.run(
                        ["bash", "submit.sh"],  # noqa: S603, S607
                        check=True,
                        text=True,
                        capture_output=True,
                    )
                    if result.returncode:
                        logger.warning(
                            f"Process for job with id {job_id} and index {index} finished with an error"
                        )
                batch_manager.set_job_finished(job_id, index)
            except Exception:
                logger.exception(
                    "Error while running job with id {job_id} and index {index}"
                )
            else:
                logger.info(f"Completed job with id {job_id} and index {index}")

        if parallel_ids:
            parallel_ids[os.getpid()] = False


def decompress_files(store: JobStore) -> None:
    file_names = [OUT_FILENAME]
    file_names.extend(os.path.basename(p) for p in get_store_file_paths(store))

    for fn in file_names:
        # If the file is already present do not decompress it, even if
        # a compressed version is present.
        if os.path.isfile(fn):
            continue
        for f in glob.glob(fn + ".*"):
            decompress_file(f)
