#!/usr/bin/env python

import os
import re
import pymongo
from bs4 import BeautifulSoup
from attrdict import AttrDict
import json
import datetime
import sys
from bson import json_util
import base64
import zlib
import urllib
import subprocess
import yaml

config = AttrDict({
  'mongo_host':None,
  'mongo_db':None,
  'mongo_collection':None,
  'collection':None,
  'data_dir':None
})

count = {}

here = os.path.dirname(os.path.realpath(__file__)) + '/..'

def main():
  set_config()
  if len(sys.argv) > 1:
    for arg in sys.argv[1:]:
      if re.search(r'html', arg, flags=re.I):
        from_html()
      elif re.search(r'routes', arg, flags=re.I):
        to_routes()
      elif re.search(r'timing', arg, flags=re.I):
        to_timing()
      else:
        raise Exception('do not know what ' + arg + ' is')
  else:
    raise Exception('need to know if I am scraping HTML, generating routes, or generating timing data')

def writeflush(s):
  sys.stdout.write(s)
  sys.stdout.flush()

def stamp_to_date(stamp):
  m = re.search(r'^(\d\d)(\d\d)(\d\d)-(\d\d)(\d\d)', stamp)
  if m:
    nums = [int(m.group(x)) for x in xrange(2, 6)]
    nums.insert(0, int('20'+m.group(1)))
    return datetime.datetime(*nums)
  else:
    return datetime.datetime(1970, 1, 1)

def from_html():
  records = []
  count = 0
  temp_file = '/tmp/{}.json'.format(config.mongo_collection)

  most_recent = [x for x in config.collection.aggregate([{'$group':{'_id':'$stamp', 'latest':{'$max':'$utc'}}}])]
  if len(most_recent) > 0:
    most_recent = most_recent[0]['latest']
  else:
    most_recent = datetime.datetime(1970,1,1)

  existing = {d:1 for d in config.collection.distinct('stamp')}
  print "found {} existing timestamps, most recent in mongo is {}".format(len(existing.keys()), most_recent)

  files = subprocess.check_output('cd {} && ls -1 | sort -r'.format(config.data_dir), shell=True)
  for filename in files.splitlines():
    if stamp_to_date(filename) < most_recent:
      break
    count += 1
    if filename not in existing:
      records.extend([p for p in parse(filename) if p is not None])
    writeflush('\r{}'.format(count))

  writeflush('\n')

  write_json(records, temp_file)
  import_cmd = 'mongoimport --jsonArray -h {} -d {} {}'.format(config.mongo_host, config.mongo_db, temp_file)
  os.system(import_cmd)

# yields multiple times, with value of either None or an AttrDict
def parse(filename):
  utc_dt = datetime.datetime.strptime('20'+filename, '%Y%m%d-%H%M')

  with open(config.data_dir + '/' + filename) as f:
    html = f.read()

  if len(html) > 10 or re.search(r'Travel Times Currently Unavailable', html):
    yield None

  soup = BeautifulSoup(html)
  prev = AttrDict({
    'start':'',
    'end':'',
    'via':''
  })

  for row in soup.select('tr')[1:]:
    data = AttrDict({
      'stamp':filename,
      'utc':utc_dt,
      'start':None,
      'end':None,
      'via':None,
      'distance':None,
      'average':None,
      'current':None,
      'hov':None
    })

    def number_parse(index):
      val = re.sub(r'\s', '', tds[index].get_text())
      try:
        val = float(val)
      except:
        val = 0.0
      return int(val) if float.is_integer(val) else val

    tds = row.select('td')
    m = re.search(r'(.*?) to (.*)', tds[1].get_text())
    express = True if re.search(r'express lanes', tds[0].get_text(), re.I) else False
    offset = -1 if express else 0

    data.start = m.group(1) if m else prev.start
    data.end = m.group(2) if m else prev.end

    if express:
      data.via = prev.via + ' (E)'
    else:
      data.via = ', '.join([t['alt'] for t in tds[0].select('img[alt]')])
      prev.via = data.via

    data.distance = number_parse(2+offset)
    data.average = number_parse(3+offset)
    data.current = number_parse(4+offset)
    data.hov = number_parse(5+offset)

    yield data
    prev.start = data.start
    prev.end = data.end

def write_json(data, filename, compress=False):
  with open(filename, 'w') as f:
    if compress:
      text = json.dumps(data, default=json_util.default)
      compressed = base64.b64encode(zlib.compress(urllib.quote(text), 9))
      f.write(compressed)
    else:
      json.dump(data, f, default=json_util.default)
    print filename

def to_routes():
  routes = {}
  cities = sorted([x for x in config.collection.distinct('start')])
  for start in cities:
    for end in cities:
      if start not in routes:
        routes[start] = {}
      if end not in routes[start]:
        routes[start][end] = None

      print "{}->{}".format(start, end)
      routes[start][end] = [x for x in config.collection.distinct('via', {'start':start, 'end':end})]

  write_json(routes, here+'/assets/routes.json')

def to_timing():
  json_data = {}
  with open(here+'/assets/routes.json', 'r') as f:
    routes = json.load(f)
    for start in routes.keys():
      for end in routes[start].keys():
        directory = 'assets/timing/{}/{}'.format(start,end)
        try:
          os.makedirs(directory)
        except:
          pass
        for route in routes[start][end]:
          filename = '{}/{}/{}.base64'.format(here, directory, route)
          json_data[filename] = []

  min_time = datetime.datetime.utcnow() - datetime.timedelta(weeks=8, days=2)
  all_times = config.collection.find(
    {'utc':{'$gt':min_time}}, {'utc':1, 'current':1, 'start':1, 'end':1, 'via':1}
  ).sort('utc', pymongo.ASCENDING)

  missing_routes = {}
  for t in all_times:
    filename = '{}/assets/timing/{}/{}/{}.base64'.format(here, t['start'], t['end'], t['via'])
    if filename in json_data:
      json_data[filename].append([
        int((t['utc'] - datetime.datetime(1970,1,1)).total_seconds()),
        t['current']
      ])
    else:
      route = '{}/{}/{}'.format(t['start'], t['end'], t['via'])
      missing_routes[route] = True

  for f in json_data.keys():
    write_json(json_data[f], f, True)

  for r in missing_routes.keys():
    #sys.stderr.write('missing route: {}\n'.format(r))
    pass

def set_config():
  global config

  with open(here+'/ship.conf', 'r') as f:
    opts = yaml.load(f)['mongo']

  config.mongo_host = opts[0]
  client = pymongo.MongoClient(config.mongo_host, socketKeepAlive=True)

  config.collection = client[opts[2]][opts[3]]
  config.data_dir = opts[1]
  config.mongo_db = opts[2]
  config.mongo_collection = opts[3]

if __name__ == '__main__':
  main()
