#!/usr/bin/env python

import argparse
import logging
import os
import sys

import obspy

from obspy.core import UTCDateTime
try: 
    # more recent versions of ObsPy
    from obspy.clients.fdsn import Client
except:
    # this works for 0.10.2 version
    from obspy.fdsn import Client

if __name__ == "__main__":
    """
        Simple program to query a FDSN Station Webservice and load basic 
        meta-data into a PostgreSQL AQMS database.

        Populates D_Abbreviation, Station_Data, Channel_Data, and Simple_Response
        If these tables do not exist yet, they will be created.

        If no database is available, the script can be run with the -f flag
        to write the meta-data to a StationXML file instead.
    """
    POPULATE = True

    parser = argparse.ArgumentParser(description="Retrieves FDSN StationXML from \
        a fdsn webservice (default=IRIS) and populates (PostgreSQL) AQMS tables \
        station_data, channel_data, simple_response, and associated dictionary \
        tables (d_abbreviation, d_unit, d_format). Database connection \
        parameters have to be set with environment variables DB_NAME, DB_HOST, \
        DB_PORT, DB_USER, and optionally, DB_PASSWORD. \
        The tables will be created if they do not exist yet. \
        Alternatively it can be run with the -f flag to save the meta-data to a \
        StationXML file instead. Logs are written to fdsnws-station2aqms.log. \
        See https://github.com/jrhartog/fdsnws-station2aqms") 

    # required argument
    help_text = "Specify a FDSN or Virtual network code, wildcards are allowed"
    parser.add_argument("network",help=help_text)

    # optional arguments
    help_text = "Save the inventory to a file instead, provide filename"
    parser.add_argument("-f","--filename", help=help_text)
    help_text = "Specify a station code, wildcards are allowed"
    parser.add_argument("-s","--station",help=help_text)
    help_text = "Specify a channel code, wildcards are allowed"
    parser.add_argument("-c","--channel",help=help_text)
    help_text = "Specify a location code, wildcards are allowed"
    parser.add_argument("-l","--location",help=help_text)
    help_text = "Request metadata for all times, default is active channels only!"
    parser.add_argument("-a","--all",help=help_text,action="store_true")
    help_text = "Specify Webservice to query (default=IRIS)"
    parser.add_argument("-ws","--webservice",help=help_text,default="IRIS", choices=["BGR",
    "EMSC","ETH","GEONET","GFZ","INGV","IPGP","IRIS","ISC","KOERI","LMU","NCEDC",
    "NIEP","NOA","ODC","ORFEUS","RESIF","SCEDC","USGS","USP"])
    help_text = "Specify level of information (default=response)"
    parser.add_argument("-level","--level",help=help_text,default="response",choices=["station","channel","response"])
    help_text = "Be more verbose in logfile"
    parser.add_argument("-v","--verbose",help=help_text,action="store_true")
    
    args = parser.parse_args()
    print args

    client = Client(args.webservice)
    logfile = "fdsnws-station2aqms_{}.log".format(datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S'))
    logging.basicConfig(filename=logfile, level=logging.WARNING)

    kwargs = {}
    kwargs["network"] = args.network
    if not args.all:
        kwargs["endafter"] = UTCDateTime.now()
    if args.station:
        kwargs["station"] = args.station
    if args.channel:
        kwargs["channel"] = args.channel
    if args.location:
        kwargs["location"] = args.location
    if args.level:
        kwargs["level"] = args.level

    if args.verbose:
        logging.getLogger().setLevel(logging.INFO)


    logging.info("Retrieving station information from {} webservice\n".format(args.webservice))
    logging.info("Request parameters:\n")
    for key, value in kwargs.iteritems():
        logging.info("\t{}={}".format(key,value))
    
    # retrieve the requested inventory
    try:
        inventory = client.get_stations(**kwargs)
    except Exception as e:
        print("No data available at {}: {}".format(args.webservice,e))
        sys.exit()

    logging.info("Retrieved inventory: \n {}".format(inventory)) 


    # write to file if so requested by user
    if args.filename:
        POPULATE = False
        logging.info("Writing inventory to file {}".format(args.filename))
        inventory.write(args.filename, format="STATIONXML")

    if POPULATE:
        try:
            # import database specific stuff here so that script can be run
            # with the -f flag even when these are not available.
            from sqlalchemy import engine_from_config
            from sqlalchemy.orm import sessionmaker

            from aqms_ir.configure import configure
            from aqms_ir.inv2schema import inventory2db, print_metrics
            from aqms_ir.schema import Base


            # create configured engine instance, configure() uses environment
            # variables to create postgresql connect URI
            logging.info("Configuring database connection parameters...")
            config = configure()
            logging.info("Configured database connection parameters: {}".format(config))

            engine = engine_from_config(config, prefix='sqlalchemy.')

            # create a configured "Session" class
            Session = sessionmaker(bind=engine)

            logging.getLogger('sqlalchemy.engine').setLevel(logging.WARNING)

            db_name = config["sqlalchemy.url"].split("/")[-1]
            db_user = config["sqlalchemy.url"].split("/")[2].split(":")[0]
            logging.info("Populating database {} as user {}\n".format(db_name,db_user))
            # This command will create the database tables if they do not exist yet.
            Base.metadata.create_all(engine)

            session = Session()
            inventory2db(session,inventory)
            session.close()
        
            print_metrics()
            logging.info("Finished successfully.")

        except Exception as e:
            #logging.error("Error connecting to database {} as {} ({})".format(db_name,db_user,e))
            logging.error("Error: {}".format(e))
            if not args.filename:
                logging.error("Saving inventory to inventory.xml file instead.")
                inventory.write("inventory.xml", format="STATIONXML")
