#!/usr/bin/env python
"""
usage: 
    ia [--version] [--help] <command> [<args>...]

options:
    --version  
    -h, --help  

commands:
    configure Configure `ia`.
    metadata  Retrieve and modify metadata for items on archive.org
    upload    Upload items to archive.org
    download  Download files from archive.org
    search    Search archive.org
    mine      Download item metadata concurrently.
    catalog   Retrieve information about your catalog tasks

See 'ia <command> --help' for more information on a specific command.
"""
from docopt import docopt
import sys
import json
import os
import yaml
import urlparse

import internetarchive



# ia_configure()
#_________________________________________________________________________________________
def ia_configure(argv):
    """
    Configure the `ia` CLI and internetarchive Python library.

    usage:
        ia configure [--help] [options...]

    options:
        -h, --help
        -c, --cookies  Add your IA cookies to your configuration file.
    """
    args = docopt(ia_configure.__doc__, argv=argv)
    sys.stdout.write(
        'Please visit https://archive.org/account/s3.php to retrieve your S3 keys\n\n')
    access_key = raw_input('Please enter your IA S3 access key: ')
    secret_key = raw_input('Please enter your IA S3 secret key: ')
    config = {
            's3': {
                'access_key': access_key,
                'secret_key': secret_key
            }
    }

    if args['--cookies']:
        config['cookies'] = {
            'logged-in-user': raw_input('Please enter your logged-in-user cookie: '),
            'logged-in-sig': raw_input('Please Enter your logged-in-sig cookie: ')
    }

    configfile = yaml.dump(config, default_flow_style=False)
    configdir = os.path.join(os.environ['HOME'], '.config')

    if not os.path.isdir(configdir) and not os.path.isfile(configdir):
        os.mkdir(configdir)
    elif os.path.isdir(configdir):
        filename = os.path.join(configdir, 'internetarchive.yml')
    else:
        filename = os.path.join(os.environ['HOME'], '.internetarchive.yml')
    if os.path.exists(filename):
        overwrite = raw_input('\nYou already have an ia config file: '
                              '{0} \n\nWould you like to overwrite it?'
                              '[y/n] '.format(filename).lower())
        if overwrite not in ['y', 'yes']:
            sys.stdout.write('\nExiting without overwriting config file!\n')
            sys.exit(1)
    with open(filename, 'wb') as fp:
        os.chmod(filename, 0o700)
        fp.write(configfile)

    sys.stdout.write('\nSuccessfully saved your new config to: {0}\n'.format(filename))
    sys.exit(0)


# ia_metadata()
#_________________________________________________________________________________________
def ia_metadata(argv):
    """
    Retrieve and modify metadata for items on archive.org.
    
    usage: 
        ia metadata [--help] <identifier> [--modify <key:value>... | options...]

    options:
        -h, --help
        -m, --modify <key:value>   Modify the metadata of an item.
        -e, --exists               Check if an item exists. Exit status is 0 if the item
                                   exists, and 1 if it does not.
        -f, --files                Return select file-level metadata.
        -F, --formats               
        -c, --files-count          Return the file-count of an item.
        -i, --item-size            Return the item-size.
        -s, --server               Return the server from which the given item is being 
                                   served.
        --dir                      Return the full item-directory path.
        --d1                       Return the primary server.
        --d2                       Return the secondary server.
        -t, --target <target>      Return specified target, only.
    """
    args = docopt(ia_metadata.__doc__, argv=argv)
    item = internetarchive.Item(args['<identifier>'])

    # Check existence of item.
    if args['--exists']:
        if item.exists:
            sys.stdout.write('{0} exists\n'.format(item.identifier))
            sys.exit(0)
        else:
            sys.stderr.write('{0} does not exist\n'.format(item.identifier))
            sys.exit(1)

    # Modify metadata.
    elif args['--modify']:
        metadata = {}
        changes = [x.split(':', 1) for x in args['--modify']]
        for k,v in changes:
            if not metadata.get(k):
                metadata[k] = v
            else:
                if type(metadata[k]) != list:
                    metadata[k] = [metadata[k]]
                metadata[k].append(v)
        response = item.modify_metadata(metadata)
        status_code = response['status_code']
        if not response['content']['success']:
            error_msg = response['content']['error']
            sys.stderr.write('error: {0} ({1})\n'.format(error_msg, status_code))
            sys.exit(1)
        sys.stdout.write('success: {0}\n'.format(response['content']['log']))

    # Get metadata.
    elif args['--files']:
        for f in item.files():
            files_md = [f.item.identifier, f.name, f.source, f.format, f.size, f.md5]
            sys.stdout.write('\t'.join([str(x) for x in files_md]) + '\n')
    elif args['--formats']:
        formats = set([f.format for f in item.files()])
        sys.stdout.write('\n'.join(formats))
    elif args['--files-count']:
        sys.stdout.write(str(item.metadata.get('files_count')))
    elif args['--server']:
        sys.stdout.write(str(item.metadata.get('server')))
    elif args['--dir']:
        sys.stdout.write(str(item.metadata.get('dir')))
    elif args['--d1']:
        sys.stdout.write(str(item.metadata.get('d1')))
    elif args['--d2']:
        sys.stdout.write(str(item.metadata.get('d2')))
    elif args['--item-size']:
        sys.stdout.write(str(item.metadata.get('item_size')))
    elif args['--target']:
        keys = [k.strip('/') for k in args['--target'][0].split('/') if k]
        for i, k in enumerate(keys):
            if i == 0:
                md = item.metadata.get(k)
            else:
                md = md.get(k)
        sys.stdout.write(str(md))
    else:
        sys.stdout.write(str(item.metadata))

    sys.exit(0)


# ia_upload()
#_________________________________________________________________________________________
def ia_upload(argv):
    """
    Upload items to archive.org.
    
    usage: 
        ia upload <identifier> <file>... [options...]

    options:

     -h, --help
     -n, --no-derive             Do not derive the item after files have been 
                                 uploaded.
     -d, --debug                 Return the headers to be sent to IA-S3. default: True
     -M, --multipart             Upload files to archive.org in parts, using 
                                 IA-S3 multipart.
     -i, --ignore-bucket         Destroy and respecify the metadata for a 
                                 given item.
     -m, --metadata=<key:value>  Metadata to add to the item. default: None
     -H, --header=<key:value>    default: None
    """
    args = docopt(ia_upload.__doc__, argv=argv)

    s3_headers = dict(h.split(':') for h in args['--header'] if args['--header'])
    s3_metadata = {}
    changes = [x.split(':', 1) for x in args['--metadata']]
    for k,v in changes:
        if not s3_metadata.get(k):
            s3_metadata[k] = v
        else:
            if type(s3_metadata[k]) != list:
                s3_metadata[k] = [s3_metadata[k]]
            s3_metadata[k].append(v)

    item = internetarchive.Item(args['<identifier>'])
    upload_status = item.upload(args['<file>'], metadata=s3_metadata, headers=s3_headers,
                                debug=args['--debug'], derive=args['--no-derive'], 
                                multipart=args['--multipart'],
                                ignore_bucket=args['--ignore-bucket'])
    if args['--debug']:
        sys.stdout.write('IA-S3 Headers:\n\n{0}\n'.format(upload_status))
        sys.exit(0)
    elif not upload_status:
        sys.stderr.write('error: upload failed!\n')
        sys.exit(1)
    else:
        sys.stdout.write('uploaded:\t{0}\n'.format(item.details_url))
        sys.exit(0)


# ia_download()
#_________________________________________________________________________________________
def ia_download(argv):
    """
    Download files from archive.org.

    usage: 
        ia download [--help] <identifier> [<file>...] [options...]  

    options:
        -h, --help  
        --version  
        -i, --ignore-existing      Clobber files already downloaded.
        -s, --source=<source>...   Only Download files matching given sources.
        -o, --original             Download only files with source=original.
        -g, --glob=<pattern>       Only download files whose filename matches the 
                                   given glob pattern.
        -f, --format=<format>...   Only download files of the specified format(s).
                                   You can use the following command to retrieve
                                   a list of file formats contained within a given 
                                   item:
                                    
                                       ia metadata --formats <identifier>

        -c, --concurrent           Download files concurrently using the Python 
                                   gevent networking library.
    """
    args = docopt(ia_download.__doc__, argv=argv)

    if '/' in args['<identifier>']:
        identifier = args['<identifier>'].split('/')[0]
    else:
        identifier = args['<identifier>']
    item = internetarchive.Item(identifier)

    if '/' in args['<identifier>'] or args['<file>']:
        if not args['<file>']:
            fname = args['<identifier>'].split('/')[-1]
            files = [fname]
        else:
            files = args['<file>']
        for f in files:
            fname = f.encode('utf-8')
            path = os.path.join(identifier, fname)
            sys.stdout.write('downloading: {0}\n'.format(fname))
            f = item.file(fname)
            f.download(file_path=path, ignore_existing=args['--ignore-existing'])
        sys.exit(0)

    if args['--format']:
        formats = args['--format']
    else:
        formats = None

    if args['--glob']:
        glob = args['--glob'][0]
    else:
        glob = None

    if args['--source']:
        ia_source = args['--source']
    elif args['--original']:
        ia_source = ['original']
    else:
        ia_source = None

    item.download(formats=formats, source=ia_source, concurrent=args['--concurrent'], 
                  glob_pattern=glob, ignore_existing=args['--ignore-existing'])
    sys.exit(0)


# ia_search()
#_________________________________________________________________________________________
def ia_search(argv):
    """
    Search archive.org.

    usage: 
        ia search [--help] <query>...
    """
    args = docopt(ia_search.__doc__, argv=argv)

    query = ' '.join(args['<query>'])
    search = internetarchive.Search(query)
    for result in search.results:
        sys.stdout.write(result['identifier'] + '\n')
    sys.exit(0)


# ia_mine()
#_________________________________________________________________________________________
def ia_mine(argv):
    """
    usage:
        ia mine <itemlist.txt> [options...]

    options:
        -h, --help  
        -v, --verbose  
        -f, --files  
        -w, --workers=<count>  [default: 20]
        -c, --cache            

    """
    args = docopt(ia_mine.__doc__, argv=argv)

    items = [i.strip() for i in open(args['<itemlist.txt>'])]
    workers = int(args.get('--workers', 20)[0])
    miner = internetarchive.Mine(items, workers=workers)
    for i, item in miner.items():
        if args['--cache']:
            sys.stdout.write('saving metadata for: {0}\n'.format(item.identifier))
            with open('{0}_meta.json'.format(item.identifier), 'w') as fp:
                json.dump(item.metadata, fp)
        else:
            sys.stdout.write(item.metadata)


# ia_catalog()
#_________________________________________________________________________________________
def ia_catalog(argv):
    """
    Retrieve information about your catalog tasks.

    usage: 
        ia catalog [--help] [options...]

    options:
        -h, --help
        -v, --verbose       
        -u, --url=<url>             ...
        -g, --green-rows            Return information about tasks that have not run.
        -b, --blue-rows             Return information about running tasks.
        -r, --red-rows              Return information about tasks that have failed.
        -f , --fields <field>...    Return only the specified fields. <field> may be one
                                    of the following: 
                                    
                                    identifier, server, command, time, submitter, args, 
                                    task_id, row_type
    """
    args = docopt(ia_catalog.__doc__, argv=argv)

    if args['--url']:
        parsed_url = urlparse.urlparse(args['--url'][0])
        params = urlparse.parse_qs(parsed_url.query)
        c = internetarchive.Catalog(params=params)
    else:
        c = internetarchive.Catalog()
    row_types = {
            c.GREEN: 'green',
            c.BLUE: 'blue',
            c.RED: 'red',
            c.BROWN: 'brown',
    }
    if args['--green-rows']:
        tasks = c.green_rows
    elif args['--blue-rows']:
        tasks = c.blue_rows
    elif args['--red-rows']:
        tasks = c.red_rows
    else:
        tasks = c.tasks
    for t in tasks:
        if args['--fields']:
            task_info = []
            for field in args['--fields']:
                info = eval('t.{0}'.format(field))
                task_info.append(info)
        else:
            task_info = [
                    t.identifier, t.task_id, t.server, t.time, t.command, 
                    row_types[t.row_type]
            ]
            if args['--verbose']:
                targs = '\t'.join(['{0}={1}'.format(k, v) for (k,v) in t.args.items()])
                task_info += [t.submitter, targs]
        sys.stdout.write('\t'.join([str(x) for x in task_info]) + '\n')
    sys.exit(0)


# main()
#_________________________________________________________________________________________
if __name__ == '__main__':
    args = docopt(__doc__, version=internetarchive.__version__, options_first=True)
    cmd = args['<command>']
    aliases = dict(
            md = 'metadata',
            up = 'upload',
            ca = 'catalog',
            se = 'search',
            mi = 'mine',
            do = 'download',
    )
    if cmd in aliases:
        cmd = aliases[cmd]
    argv = [cmd] + args['<args>']

    if cmd == 'configure':
        ia_configure(argv)
    elif cmd == 'metadata':
        ia_metadata(argv)
    elif cmd == 'upload':
        ia_upload(argv)
    elif cmd == 'download':
        ia_download(argv)
    elif cmd == 'search':
        ia_search(argv)
    elif cmd == 'mine':
        ia_mine(argv)
    elif cmd == 'catalog':
        ia_catalog(argv)
    else:
        sys.stderr.write(__doc__)
        sys.exit(1)
