#!/usr/bin/env python

import os
import sys
import yaml
import shutil
import argparse
from subprocess import call

from daggit.core.base.utils import nostdout
from daggit.core.base.config import ORCHESTRATOR

with nostdout():
    from airflow.models import DagModel
    from airflow import settings
    from airflow.utils import db
    from airflow import models


# improve orchestrator factory and replace if based switching
def _unpause(dag):

    if ORCHESTRATOR == 'airflow':

        dag = dag

        session = settings.Session()
        dm = session.query(DagModel).filter(
            DagModel.dag_id == dag.dag_id).first()
        dm.is_paused = False
        session.commit()

        msg = "Dag: {}, paused: {}".format(dag, str(dag.is_paused))
        print(msg)


def init(graph):
    if ORCHESTRATOR == 'airflow':
        os.path.expanduser('~')
        airflow_home = os.getenv('AIRFLOW_HOME', os.path.expanduser('~/airflow'))
        dag_location = os.path.join(airflow_home, 'dags')
        pkgdir = sys.modules['daggit'].__path__[0]
        airflow_dag_file = pkgdir + '/runtime/airflow_dag.py'
        try:
            os.remove(os.path.join(dag_location,'airflow_dag.py'))
        except OSError:
            pass
        try:
            os.remove(os.path.join(airflow_home, 'airflow.db'))
        except OSError:
            pass
        if not os.path.exists(dag_location):
            os.makedirs(dag_location)
        shutil.copy2(airflow_dag_file, dag_location)

        graph_location = os.path.join(os.getcwd(), graph)

        if os.path.isfile(os.path.join(os.getcwd(), graph)):
            pass
        else:
            raise FileNotFoundError("The path to the dag file could not be resolved")
        # edit today
        graph_name = os.path.join(dag_location, 'graph_location') 
        # f = open(os.path.join(dag_location, 'graph_location'), "w+")
        # f.write(graph_location)
        # f.close()
        # db.initdb()
        if os.path.exists(graph_name) and os.path.getsize(graph_name)>0:
            file = open(graph_name, "r")
            read_file = file.readlines()
            new_lines = []
            for line in read_file:
                # Strip white spaces
                line = line.strip()
                if line == graph_location:
                    pass
                else:
                    new_lines.append(line)
            with open(graph_name, "w") as fp:
                for n in new_lines:
                    fp.write(n+"\n")

        f = open(graph_name, "a+")
        f.write(graph_location+"\n")
        f.close()
        db.initdb()


def run(dag_id, clear_all, clear_failed):

    if ORCHESTRATOR == 'airflow':
        if clear_all:
            call(["airflow", "clear", "-c", dag_id])
        if clear_failed:
            call(["airflow", "clear", "-c", "-f", dag_id])
        print("Running----")
        dbag = models.DagBag()
        mlwb = dbag.get_dag(dag_id)
        #_unpause(mlwb)
        print("mlwb:", mlwb)
        print("---start running")
        mlwb.run()
        print("---stop running")
    

if __name__ == '__main__':

    from sys import argv

    parser = argparse.ArgumentParser(prog='daggit')
    parser.add_argument("-v", "--version", help="MLWB version"
                        , action='store_true', dest='version', default=False)
    subparsers = parser.add_subparsers(dest='cmd', help='sub-commands')

    parser_init = subparsers.add_parser('init', help='Initiates a DAG.')
    parser_init.add_argument('graph', help='YAML file location.')

    parser_run = subparsers.add_parser('run', help='Runs a DAG.')
    parser_run.add_argument('dag_id', help='Experiment name (DAG ID)')
    parser_run.add_argument('--clear_all', action='store_true', dest='clear_all', default=False
                            , help='Clears state of all the tasks before running the DAG.')
    parser_run.add_argument('--clear_failed', action='store_true', dest='clear_failed', default=False
                            , help='Clears state of failed tasks before running the DAG.')

    args = parser.parse_args()


    if args.cmd == "init":
        init(args.graph)

    if args.cmd == "run":
        run(args.dag_id, args.clear_all, args.clear_failed)

    if args.version:
        print("version - 0.5.0")
