from typing import List
import warnings
from visivo.models.base.parent_model import ParentModel
from visivo.models.models import local_merge_model

from visivo.models.project import Project
from visivo.models.trace import Trace
from visivo.logging.logger import Logger
from time import time
from concurrent.futures import Future, ThreadPoolExecutor
import queue
from visivo.query.jobs.job import CachedFuture, Job, JobResult

from visivo.query.jobs.run_csv_script_job import jobs as csv_script_jobs
from visivo.query.jobs.run_trace_job import jobs as run_trace_jobs
from visivo.query.jobs.run_local_merge_job import jobs as run_local_merge_jobs
from visivo.query.target_job_tracker import TargetJobTracker

warnings.filterwarnings("ignore")


class Runner:
    def __init__(
        self,
        project: Project,
        output_dir: str,
        threads: int = 8,
        soft_failure=False,
        run_only_changed=False,
        name_filter: str = None,
    ):
        self.project = project
        self.output_dir = output_dir
        self.run_only_changed = run_only_changed
        self.threads = threads
        self.soft_failure = soft_failure
        self.name_filter = name_filter
        self.dag = project.dag()
        self.errors = []
        self.jobs: List[Job] = []

    def run(self):
        complete = False
        target_job_tracker = TargetJobTracker()
        start_time = time()
        self.jobs = self._all_jobs()
        with ThreadPoolExecutor(max_workers=self.threads) as executor:
            while True:
                complete = self.update_job_queue(target_job_tracker)
                if complete:
                    break

                try:
                    job = target_job_tracker.get_next_job()
                except queue.Empty:
                    continue

                if job.done():
                    pass
                elif target_job_tracker.is_accepting_job(job):
                    Logger.instance().info(job.start_message())
                    job.set_future(executor.submit(job.action, **job.kwargs))
                    job.future.add_done_callback(self.job_callback)
                else:
                    target_job_tracker.return_to_queue(job)

        if len(self.errors) > 0 and self.soft_failure:
            Logger.instance().error(
                f"\nRefresh failed in {round(time()-start_time, 2)}s with {len(self.errors)} query error(s)."
            )
        elif len(self.errors) > 0 and not self.soft_failure:
            Logger.instance().error(
                f"\nRun failed in {round(time()-start_time, 2)}s with {len(self.errors)} query error(s)"
            )
            exit(1)
        else:
            Logger.instance().info(f"\nRun finished in {round(time()-start_time, 2)}s")

    def update_job_queue(self, target_job_tracker: TargetJobTracker) -> bool:
        all_dependencies_completed = True
        for job in self.jobs:
            job_item_children = ParentModel.all_descendants(
                dag=self.dag, from_node=job.item
            )
            dependencies = list(
                filter(
                    lambda j: job != j and job_item_children.has_node(j.item),
                    self.jobs,
                )
            )
            incomplete_dependencies = list(
                filter(
                    lambda d: not target_job_tracker.is_job_name_done(d.name),
                    dependencies,
                )
            )
            dependencies_completed = len(incomplete_dependencies) == 0
            if not dependencies_completed:
                all_dependencies_completed = False

            if dependencies_completed and not target_job_tracker.is_job_name_enqueued(
                job.name
            ):
                if not job.output_changed and self.run_only_changed:
                    job.future = CachedFuture()
                target_job_tracker.track_job(job)

        return all_dependencies_completed and target_job_tracker.empty()

    def job_callback(self, future: Future):
        job_result: JobResult = future.result(timeout=1)
        if job_result.success:
            Logger.instance().success(str(job_result.message))
        else:
            Logger.instance().error(str(job_result.message))
            self.errors.append(str(job_result.message))

    def _all_jobs(self) -> List[Job]:
        jobs = []
        jobs = jobs + run_trace_jobs(
            dag=self.dag,
            output_dir=self.output_dir,
            project=self.project,
            name_filter=self.name_filter,
        )
        jobs = jobs + csv_script_jobs(
            dag=self.dag,
            output_dir=self.output_dir,
            project=self.project,
            name_filter=self.name_filter,
        )
        jobs = jobs + run_local_merge_jobs(
            dag=self.dag,
            output_dir=self.output_dir,
            project=self.project,
            name_filter=self.name_filter,
        )
        return jobs
