import argparse
import os
import shutil
import typing as T

import arrow
import requests

# Sample Python Client
from client import SkydioAPIClient

DEFAULT_OUTPUT_DIR = "/tmp/media_files/"


def dirname_for_flight(flight):
    # type: (T.Dict) -> str
    """Create a human-readable directory name for the flight"""
    return "/{}_{}_{}_{}/".format(
        flight["takeoff"], flight["vehicle_serial"], flight["user_email"], flight["flight_id"]
    )


def request_all_files_for_flight(client, flight_id):
    # type: (SkydioAPIClient, str) -> T.List[T.Dict]
    # Construct query
    query = {"flight_id": flight_id}  # type: T.Dict[str, T.Any]

    media_files = []
    current_page = 1
    total_pages = 1500000  # A really large number, will be updated with actual total pages

    while current_page < total_pages:

        query.update({"page_number": current_page})

        files_response = client.get(
            "/v0/media_files",
            params=query,
        )
        media_files.extend(files_response["data"]["files"])

        total_pages = files_response["data"]["pagination"]["total_pages"]
        current_page += 1

    return media_files


def request_flight(client, flight_id):
    # type: (SkydioAPIClient, str) -> T.Dict
    response = client.get("/v0/flight/" + flight_id)
    return response["data"]["flight"]


def query_flights(client, search_params):
    # type: (SkydioAPIClient, T.Dict[str, T.Any]) -> T.List[T.Dict]
    response = client.get(
        "/v0/flights",
        params=search_params,
    )
    return response["data"]["flights"]


def main():
    # type: () -> None
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--api-token", type=str, required=True, help="Api token for interacting with Skydio Cloud"
    )
    parser.add_argument(
        "-o",
        "--output_dir",
        type=str,
        default=DEFAULT_OUTPUT_DIR,
        help="Directory to put output directories",
    )

    # add sub-parsers for different actions
    subparsers = parser.add_subparsers(help="Action to execute", dest="command")

    # download all media for the specified user or vehicle and duration
    recent = subparsers.add_parser("recent", description="Download recent flights for user")
    recent.add_argument("hours", type=int, help="Last N hours of flights to retrieve")
    recent.add_argument("--user-email", type=str, help="User email of flights to download for")
    recent.add_argument(
        "--vehicle-serial", type=str, help="Vehicle serial of flights to download for"
    )

    # find and download files from specific flights
    by_id = subparsers.add_parser("flight_id", description="Download files from specific flight")
    by_id.add_argument("id", type=str, nargs="+", help="The flight IDs to download")

    args = parser.parse_args()

    client = SkydioAPIClient(args.api_token)

    if args.command == "recent":
        search_params = {"takeoff_since": arrow.utcnow().shift(hours=-1 * args.hours).isoformat()}
        if args.user_email:
            search_params["user_email"] = args.user_email
        if args.vehicle_serial:
            search_params["vehicle_serial"] = args.vehicle_serial
        flights = query_flights(client, search_params)
    elif args.command == "flight_id":
        # just query a list of flight ids
        flights = [request_flight(client, id_) for id_ in args.id]
    else:
        raise KeyError("Invalid command")

    if not flights:
        raise Exception("No flights found matching criteria")

    if not os.path.isdir(args.output_dir):
        print("Making dir {}".format(args.output_dir))
        os.makedirs(args.output_dir)

    for flight in flights:
        media_files = request_all_files_for_flight(client, flight["flight_id"])

        if not media_files:
            print("No media files found for flight {}".format(flight["flight_id"]))
            continue

        dirpath = args.output_dir.rstrip("/") + dirname_for_flight(flight)
        if not os.path.isdir(dirpath):
            print("Making dir {}".format(dirpath))
            os.mkdir(dirpath)
        else:
            print("Using existing dir {}".format(dirpath))

        print("Downloading {} media files to {}".format(len(media_files), dirpath))

        for f in media_files:
            filepath = dirpath + f["filename"]
            if os.path.isfile(filepath) and os.path.getsize(filepath) == f["size"]:
                print("File already downloaded, skipping {}".format(filepath))
                continue

            # download the file
            print("Downloading to {}".format(filepath))
            with requests.get(
                f["download_url"], stream=True, headers=client.api_token_header()
            ) as r:
                with open(filepath, "wb") as file_to_write:
                    shutil.copyfileobj(r.raw, file_to_write)


if __name__ == "__main__":
    main()