#!/usr/bin/env python

# Downloads REMIB data for all nodes

# Usage: env [start=...] [n=...] self

import csv

import itertools

import os.path

import StringIO

import sys

import urllib2

sys.path.append(os.path.dirname(__file__)+"/../../lib")

import exc

import opts

import profiling

import streams

import strings

import timeout

import util

# Config

timeout = 20 # sec

max_consec_empty_responses = 10

alphabet = map(chr, xrange(ord('A'), ord('Z')+1))

class InputException(Exception): pass

class EmptyResponseException(InputException): pass

def is_ignore(line):

    line = strings.remove_line_ending(line)

    return line == '' or line.startswith('\t') or line.find(',') < 0

def main():

    # Get config from env vars

    start = util.cast(int, opts.get_env_var('start', 1))

    end = util.cast(int, util.none_if(opts.get_env_var('n', None), u''))

    if end != None: end += start

    log_ = open(sys.argv[0]+'.log', 'a')

    def clear_line(): log_.write('\n')

    log_indent = 0

    def log(msg, line_ending='\n'):

        log_.write(('    '*log_indent)+msg+line_ending)

    os.chdir(os.path.dirname(__file__)) # dir of output files

    # Get by family ('familia') because that is the most general level at which

    # an identification can be made. This assumes all records have a family.

    url_template = ('http://www.conabio.gob.mx/remib/cgi-bin/'

        'remib_distribucion.cgi?lengua=EN&niveltax=familia&taxon=[prefix]%25&'

        'pais=Todos&pais_otro=&estado=100&formato=csv&mapa=no&mapabase=estados'

        '&coleccion=id%3D[node_id]')

    if end != None: node_ids = xrange(start, end)

    else: node_ids = itertools.count(start)

    consec_empty_responses = 0

    for node_id in node_ids:

        if consec_empty_responses > max_consec_empty_responses: break

        log('Processing node #'+str(node_id)+'...')

        log_indent += 1

        profiler = profiling.ItersProfiler(start_now=True, iter_text='row')

        filename = 'node.'+str(node_id)+'.csv'

        out = streams.LineCountStream(open(filename, 'w'))

        def log_ex(e):

            clear_line()

            log('! Output line '+str(out.line_num)+': '+exc.str_(e))

        start_line_num = out.line_num

        node_url_template = url_template.replace('[node_id]', str(node_id))

        for prefix_chars in itertools.product(alphabet, repeat=2):

            prefix = ''.join(prefix_chars)

            log('Processing prefix '+prefix+'...')

            row_ct = 0

            def print_status(line_ending='\n'):

                log('Processed '+str(row_ct)+' row(s)', line_ending)

            log_indent += 1

            url = node_url_template.replace('[prefix]', prefix)

            stream = streams.StreamIter(streams.TimeoutInputStream(

                urllib2.urlopen(url), timeout))

            is_empty_response = False

            try:

                util.skip(stream, is_ignore) # skip header

                try: metadata_row = csv.reader(stream).next()

                except StopIteration: raise EmptyResponseException()

                if metadata_row[0] != 'COLLECTION': raise InputException(

                    'Invalid metadata row: '+str(metadata_row))

                # Copy lines

                for line in stream:

                    if is_ignore(line):

                        error = strings.remove_prefix('\t\t', line)

                        if len(error) != len(line): raise InputException(error)

                        break

                    out.write(line)

                    row_ct += 1

                    if row_ct % 100 == 0: print_status('\r')

                        # CR at end so next print overwrites msg

            except EmptyResponseException, e: # must come before InputException

                is_empty_response = True

                log_ex(e)

                break # assume node doesn't exist, so abort node

            except InputException, e: log_ex(e)

            except timeout.TimeoutException, e:

                log_ex(e)

                break # assume node is down, so abort node

            finally: # still run if break is called

                stream.close()

                profiler.add_iters(row_ct)

                print_status()

                log_indent -= 1

                if is_empty_response: consec_empty_responses += 1

                else: consec_empty_responses = 0 # reset count

        profiler.stop()

        log(profiler.msg())

        log_indent -= 1

main()