#!python
import argparse
import json
import os

import colorful

from taigapy import TaigaClient, DEFAULT_TAIGA_URL
from taigapy.custom_exceptions import (
    Taiga404Exception,
    TaigaDeletedVersionException,
    TaigaClientConnectionException,
)
from taigapy.utils import format_datafile_id


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
    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)

    return TaigaClient(url=url, cache_dir=cache_dir)


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)

    if args.format == "feather":
        _fetch = tc.get
        datafile_index = 2  # feather_path
    else:
        _fetch = tc.download_to_cache
        datafile_index = 1  # raw_path

    tc = _get_taiga_client(args)

    df_or_path = _fetch(
        id=args.data_file_id,
        name=args.name,
        version=args.version,
        file=args.file,
    )
    if df_or_path is None:
        d = {"error": True}
    else:
        if args.data_file_id is not None:
            query = args.data_file_id
        else:
            query = format_datafile_id(args.name, args.version, args.file)

        datafile = tc.cache._get_datafile_from_db(query, query)
        d = {
            "filename": datafile[datafile_index],
            "datafile_type": datafile.datafile_format,
            "error": False,
        }

    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", 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(
        "--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)
