#!python
import argparse
import colorful
import json
import os
from taigapy import TaigaClient, DEFAULT_TAIGA_URL
from taigapy.custom_exceptions import (
    Taiga404Exception,
    TaigaDeletedVersionException,
    TaigaClientConnectionException,
)


def _get_taiga_client(args: argparse.Namespace):
    """Get TaigaClient based on args from either `fetch` or `dataset_meta`"""
    url = DEFAULT_TAIGA_URL
    cache_dir = None
    token_path = None
    if args.taiga_url is not None:
        url = args.taiga_url
    if args.data_dir is not None:
        cache_dir = os.path.expanduser(args.data_dir)
        token_path = os.path.join(cache_dir, "token")
        if not os.path.exists(token_path):
            raise Exception("No Taiga token found at {}".format(args.data_dir))

    return TaigaClient(url=url, cache_dir=cache_dir, token_path=token_path)


def fetch(args):
    if args.data_file_id is None and args.name is None:
        raise Exception("data_file_id or name must be set")

    tc = _get_taiga_client(args)

    try:
        filename, datafile_type = tc.download_to_cache_for_fetch(
            datafile_id=args.data_file_id,
            dataset_name=args.name,
            dataset_version=args.version,
            datafile_name=args.file,
            force_fetch=args.force_fetch,
            force_convert=args.force_convert,
            file_format=args.format,
            quiet=args.quiet,
            encoding="utf8",
        )
        d = {"filename": filename, "datafile_type": datafile_type, "error": False}
    except (
        Taiga404Exception,
        TaigaDeletedVersionException,
        TaigaClientConnectionException,
        ValueError,
    ) as e:
        print(colorful.red(str(e)))
        d = {"error": True}

    if args.write_filename is not None:
        with open(args.write_filename, "w+") as f:
            j = json.dump(d, f)
            f.close()
    else:
        print(d)


def dataset_meta(args):
    tc = _get_taiga_client(args)
    metadata = tc.get_dataset_metadata(
        args.dataset_name, version=args.version, version_id=args.version_id
    )
    if args.write_filename is not None:
        with open(args.write_filename, "w+") as f:
            j = json.dump(metadata, f)
            f.close()
    else:
        print(metadata)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--taiga-url", help="Override default Taiga url (https://cds.team/taiga)"
    )

    parser.add_argument(
        "--data-dir", help="Path to where token file lives and cached downloaded files"
    )

    subparsers = parser.add_subparsers(title="commands", required=True, dest="command")

    # fetch command parser
    parser_fetch = subparsers.add_parser(
        "fetch", help="Download a Taiga file into the cache directory"
    )
    parser_fetch.add_argument(
        "data_file_id",
        nargs="?",
        default=None,
        help="Taiga ID or datafile ID. If not set, NAME must be set",
    )
    parser_fetch.add_argument(
        "--name", help="Dataset name. Must be set if data_file_id is not set."
    )
    parser_fetch.add_argument("--version", help="Dataset version")
    parser_fetch.add_argument("--file", help="Datafile name")
    parser_fetch.add_argument(
        "--force-fetch",
        action="store_true",
        help="If set, will bypass local cache and try to redownload from Taiga",
    )
    parser_fetch.add_argument(
        "--quiet", action="store_true", help="If set, do not print progress"
    )
    parser_fetch.add_argument(
        "--force-convert",
        action="store_true",
        help="Ask Taiga to convert this file again (Implies --force-fetch)",
    )
    parser_fetch.add_argument(
        "--format",
        choices=["raw", "feather"],
        default="feather",
        help="Format to store file. If Taiga file is a raw file, choose raw. Otherwise, the default is feather.",
    )
    parser_fetch.add_argument(
        "--write-filename",
        help="If set, will write the full path and Taiga file type of the cached file to WRITE_FILENAME. Otherwise, will write to stdout",
    )
    parser_fetch.set_defaults(func=fetch)

    # dataset-meta command parser
    parser_dataset_meta = subparsers.add_parser(
        "dataset-meta",
        help="Fetch the metadata about a dataset (or dataset version if version-id is provided).",
    )
    parser_dataset_meta.add_argument(
        "dataset_name", nargs="?", default=None, help="Dataset name or ID"
    )
    parser_dataset_meta.add_argument("--version", help="Dataset version")
    parser_dataset_meta.add_argument("--version-id", help="Dataset version ID")
    parser_dataset_meta.add_argument(
        "--write-filename",
        help="If set, will write the metadata to WRITE_FILENAME. Otherwise, will write to stdout",
    )
    parser_dataset_meta.set_defaults(func=dataset_meta)

    args = parser.parse_args()

    args.func(args)
