from __future__ import annotations

import argparse
from os import environ
import threading
import json
from pathlib import Path
import sys
import traceback
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Iterator, List, Any, Optional, Dict, Tuple

from dbt.logger import GLOBAL_LOGGER as logger

from fal.node_graph import FalScript
from fal.utils import print_run_info, DynamicIndexProvider
from faldbt.parse import get_environment
from faldbt.project import DbtModel, FalDbt, NodeStatus

from datetime import datetime, timezone

SUCCESS = 0
FAILURE = 1


class Task:
    def set_run_index(self, index_provider: DynamicIndexProvider) -> None:
        self._run_index = index_provider.next()

    @property
    def run_index(self) -> int:
        run_index = getattr(self, "_run_index", None)
        assert run_index is not None
        return run_index

    def execute(self, args: argparse.Namespace, fal_dbt: FalDbt) -> int:
        raise NotImplementedError


class Status(Enum):
    PENDING = auto()
    RUNNING = auto()
    SKIPPED = auto()
    SUCCESS = auto()
    FAILURE = auto()


def _unique_id_to_model_name(unique_id: str) -> str:
    split_list = unique_id.split(".")
    # if its a unique id 'model.fal_test.model_with_before_scripts'
    return split_list[len(split_list) - 1]


def _unique_ids_to_model_names(id_list: List[str]) -> List[str]:
    return list(map(_unique_id_to_model_name, id_list))


def _mark_dbt_nodes_status_and_response(
    fal_dbt: FalDbt,
    status: NodeStatus,
    dbt_node: Optional[str] = None,
    adapter_response: Optional[dict] = None,
):
    for model in fal_dbt.models:
        if dbt_node is not None:
            if model.unique_id == dbt_node:
                model.status = status

                if adapter_response is not None:
                    model.adapter_response = adapter_response
        else:
            model.status = status


def _map_cli_output_model_results(
    run_results: Dict[Any, Any]
) -> Iterator[Tuple[str, NodeStatus, Optional[dict]]]:
    if not isinstance(run_results.get("results"), list):
        raise Exception("Could not read dbt run results")

    for result in run_results["results"]:
        if not result.get("unique_id") or not result.get("status"):
            continue

        yield result["unique_id"], NodeStatus(result["status"]), result.get(
            "adapter_response"
        )


def _run_script(script: FalScript) -> Dict[str, Any]:
    print_run_info([script])

    # DBT seems to be dealing with only UTC times
    # so we'll follow that convention.
    started_at = datetime.now(tz=timezone.utc)
    try:
        with _modify_path(script.faldbt):
            script.exec()
    except:
        logger.error("Error in script {}:\n{}", script.id, traceback.format_exc())
        # TODO: what else to do?
        status = NodeStatus.Fail
    else:
        status = NodeStatus.Success
    finally:
        logger.debug("Finished script {}", script.id)
        finished_at = datetime.now(tz=timezone.utc)

    return {
        "path": str(script.path),
        "unique_id": str(script.relative_path),
        "status": status,
        "thread_id": threading.current_thread().name,
        "is_hook": script.is_hook,
        "execution_time": (finished_at - started_at).total_seconds(),
        "timing": [
            {
                "name": "execute",
                # DBT's suffix for UTC is Z, but isoformat() uses +00:00. So
                # we'll manually cast it to the proper format.
                # https://stackoverflow.com/a/42777551
                "started_at": started_at.isoformat().replace("+00:00", "Z"),
                "finished_at": finished_at.isoformat().replace("+00:00", "Z"),
            }
        ],
    }


def run_script(script: FalScript, run_index: int) -> int:
    results = _run_script(script)
    run_results_file = Path(script.faldbt.target_path) / f"fal_results_{run_index}.json"
    with open(run_results_file, "w") as stream:
        json.dump(
            {
                "results": [results],
                "elapsed_time": results["execution_time"],
            },
            stream,
        )
    return SUCCESS if results["status"] == NodeStatus.Success else FAILURE


@contextmanager
def _modify_path(fal_dbt: FalDbt):
    sys.path.append(fal_dbt.scripts_dir)
    try:
        yield
    finally:
        sys.path.remove(fal_dbt.scripts_dir)


@dataclass
class DBTTask(Task):
    model_ids: List[str]

    def execute(self, args: argparse.Namespace, fal_dbt: FalDbt) -> int:
        from fal.cli.dbt_runner import dbt_run_through_python

        model_names = _unique_ids_to_model_names(self.model_ids)
        output = dbt_run_through_python(
            args, model_names, fal_dbt.target_path, self.run_index
        )

        for node, status, adapter_response in _map_cli_output_model_results(
            output.run_results
        ):
            _mark_dbt_nodes_status_and_response(fal_dbt, status, node, adapter_response)

        return output.return_code


@dataclass
class FalModelTask(DBTTask):
    bound_model: Optional[DbtModel] = None

    def set_run_index(self, index_provider: DynamicIndexProvider) -> None:
        super().set_run_index(index_provider)
        self._model_script_run_index = index_provider.next()

    @property
    def model_script_run_index(self) -> int:
        run_index = getattr(self, "_model_script_run_index", None)
        assert run_index is not None
        return run_index

    def execute(self, args: argparse.Namespace, fal_dbt: FalDbt) -> int:
        # Run the ephemeral model
        dbt_result = super().execute(args, fal_dbt)

        # And then run the Python script if it didn't fail.
        if dbt_result != SUCCESS:
            return dbt_result

        assert self.bound_model is not None
        bound_script = FalScript.model_script(fal_dbt, self.bound_model)
        script_result = run_script(bound_script, self.model_script_run_index)

        status = NodeStatus.Success if script_result == SUCCESS else NodeStatus.Error
        _mark_dbt_nodes_status_and_response(fal_dbt, status, self.bound_model.unique_id)
        return script_result


@dataclass
class FalLocalHookTask(Task):
    hook_path: Path
    bound_model: Optional[DbtModel] = None
    arguments: Optional[Dict[str, Any]] = None
    is_hook: bool = True

    def execute(self, args: argparse.Namespace, fal_dbt: FalDbt) -> int:
        script = self.build_fal_script(fal_dbt)
        return run_script(script, self.run_index)

    @classmethod
    def from_fal_script(cls, script: FalScript):
        return cls(
            script.path,
            script.model,
            script.hook_arguments,
            script.is_hook,
        )

    def build_fal_script(self, fal_dbt: FalDbt):
        return FalScript(
            fal_dbt,
            self.bound_model,
            str(self.hook_path),
            hook_arguments=self.arguments,
            is_hook=self.is_hook,
        )


@dataclass
class FalIsolatedHookTask(Task):
    hook_path: Path
    environment_name: str
    bound_model_name: str
    arguments: Optional[Dict[str, Any]] = None

    def execute(self, args: argparse.Namespace, fal_dbt: FalDbt) -> int:
        environment = get_environment(fal_dbt.project_dir, self.environment_name)
        with environment.setup() as hook_runner:
            return hook_runner(
                fal_dbt=fal_dbt,
                hook_path=self.hook_path,
                arguments=self.arguments,
                bound_model_name=self.bound_model_name,
                run_index=self.run_index,
                disable_logging=args.disable_logging,
            )


@dataclass
class TaskGroup:
    task: Task
    pre_hooks: List[Task] = field(default_factory=list)
    post_hooks: List[Task] = field(default_factory=list)
    dependencies: List[TaskGroup] = field(default_factory=list)
    status: Status = Status.PENDING

    def __post_init__(self):
        self._id = str(uuid.uuid4())

    @property
    def id(self) -> str:
        return self._id

    def iter_tasks(self) -> Iterator[Task]:
        yield from self.pre_hooks
        yield self.task
        yield from self.post_hooks
