import os

import prefect
from prefect import Flow, task, Parameter, unmapped
from prefect.client import Client
from prefect.storage import GitHub
from prefect.utilities.graphql import with_args
from prefect.utilities.logging import get_logger
from prefect.serialization.run_config import RunConfigSchema

HANDLERS = {}


def handler(type):
    def f(cls):
        HANDLERS[type] = cls
        return cls

    return f


no_default = object()


logger = get_logger("registrar")


def get_field(config, key, typ, default=no_default):
    if key not in config:
        if default is not no_default:
            return default
        raise ValueError(f"`config` is missing required field `{key}`")

    value = config[key]
    if not isinstance(value, typ):
        raise TypeError(f"`{key}` field expects a `{typ.__name__}`, got `{value!r}`")
    return value


@handler("local")
class LocalHandler:
    def __init__(self, config):
        self.path = get_field(config, "path", str)

    def describe(self):
        return os.path.abspath(self.path)

    def get_storage(self, path):
        from prefect.storage import Local

        return Local(stored_as_script=True, path=path)

    def find_files(self):
        paths = os.listdir(self.path)
        return [os.path.join(self.path, p) for p in paths]

    def load_file(self, path):
        with open(path, "r") as fil:
            return fil.read()


@handler("github")
class GitHubHandler:
    def __init__(self, config):
        self.repo = get_field(config, "repo", str)
        self.path = get_field(config, "path", str)
        self.ref = get_field(config, "ref", str, default=None)
        self._client = None

    def describe(self):
        return "GitHub {repo=%r, path=%r, ref=%r}" % (self.repo, self.path, self.ref)

    def get_storage(self, path):
        from prefect.storage import GitHub

        return GitHub(repo=self.repo, path=path, ref=self.ref)

    @property
    def repo_client(self):
        from github import Github

        github = Github()
        if self._client is None:
            self._client = github.get_repo(self.repo)
        return self._client

    def find_files(self):
        objs = self.repo_client.get_contents(
            self.path, ref=self.ref or self.repo_client.default_branch
        )
        if not isinstance(objs, list):
            objs = [objs]
        return [o.path for o in objs]

    def load_file(self, path):
        obj = self.repo_client.get_contents(
            path, ref=self.ref or self.repo_client.default_branch
        )
        return obj.decoded_content.decode()


@task
def setup(config):
    if not isinstance(config, dict):
        raise ValueError("`config` expects a dict, got %s" % type(config))

    typ = get_field(config, "type", str).lower()

    if typ in HANDLERS:
        return HANDLERS[typ](config)
    else:
        raise ValueError(f"Unsupported `type`: {type}")


@task
def load_active_run_config():
    client = Client()
    query = {
        "query": {
            with_args("flow_run_by_pk", {"id": prefect.context.flow_run_id}): {
                "run_config": True
            }
        }
    }
    blob = client.graphql(query).data.flow_run_by_pk.run_config
    return RunConfigSchema().load(blob)


@task
def find_files(handler):
    logger.info("Looking for `.py` files in %s...", handler.describe())
    paths = handler.find_files()
    logger.info("Found %d files.", len(paths))
    return paths


@task
def register_flows(handler, path, run_config):
    logger.info("Loading flows from %r...", path)
    contents = handler.load_file(path)
    ns = {}
    exec(contents, ns)
    flows = [f for f in ns.values() if isinstance(f, Flow)]
    logger.info("Found %d flows in %r", len(flows), path)
    for f in flows:
        f.run_config = run_config
        f.storage = handler.get_storage(path)
        logger.info("Registering %r", f.name)
        f.register(
            project_name=prefect.context.project_name,
            idempotency_key=f.serialized_hash(),
            no_url=True,
        )
    logger.info("Successfully registered all flows in %r", path)


with Flow("registrar") as flow:
    config = Parameter("config")

    handler = setup(config)

    files = find_files(handler)

    run_config = load_active_run_config()

    register_flows.map(unmapped(handler), files, unmapped(run_config))


flow.storage = GitHub(repo="prefecthq/system-flows", path="flows/registrar.py")
