#!/usr/bin/env python3
"""
Access the KM3NeT StreamDS DataBase service.

Usage:
    streamds
    streamds list
    streamds info STREAM
    streamds get [-f FORMAT] STREAM [PARAMETERS...]
    streamds upload [-q -x] CSV_FILE
    streamds (-h | --help)
    streamds --version

Options:
    STREAM      Name of the stream.
    PARAMETERS  List of parameters separated by space (e.g. detid=29).
    CSV_FILE    Whitespace separated data for the runsummary tables.
    -f FORMAT   Usually 'txt' for ASCII or 'text' for UTF-8 [default: txt].
    -q          Test run! When uploading, a TEST_ prefix will be added to the data.
    -x          Do not verify the SSL certificate.
    -h --help   Show this screen.

"""
import getpass
import json
import logging
import os
import requests

import km3db
import km3db.extras
from docopt import docopt


log = logging.getLogger("streamds")

RUNSUMMARY_URL = "https://km3netdbweb.in2p3.fr/jsonds/runsummarynumbers/i"
REQUIRED_COLUMNS = set(["run", "det_id", "source"])


def print_streams():
    """Print all available streams with their full description"""
    sds = km3db.StreamDS()
    sds.print_streams()


def print_info(stream):
    """Print the information about a stream"""
    sds = km3db.StreamDS()
    sds.help(stream)


def get_data(stream, parameters, fmt):
    """Retrieve data for given stream and parameters, or None if not found"""
    sds = km3db.StreamDS()
    if stream not in sds.streams:
        log.error("Stream '{}' not found in the database.".format(stream))
        return
    params = {}
    if parameters:
        for parameter in parameters:
            if "=" not in parameter:
                log.error(
                    "Invalid parameter syntax '{}'\n"
                    "The correct syntax is 'parameter=value'".format(parameter)
                )
                continue
            key, value = parameter.split("=")
            params[key] = value
    data = sds.get(stream, fmt, **params)
    if data is not None:
        try:
            print(data)
        except BrokenPipeError:
            pass
    else:
        sds.help(stream)


def available_streams():
    """Show a short list of available streams."""
    sds = km3db.StreamDS()
    print("Available streams: ")
    print(", ".join(sorted(sds.streams)))


def upload_runsummary(csv_filename, testrun=False, verify=False):
    """Reads the CSV file and uploads its contents to the runsummary table"""
    pd = km3db.extras.pandas()

    print("Checking '{}' for consistency.".format(csv_filename))
    if not os.path.exists(csv_filename):
        log.critical("{} -> file not found.".format(csv_filename))
        return
    try:
        df = pd.read_csv(csv_filename, delim_whitespace=True)
    except pd.errors.EmptyDataError as e:
        log.error(e)
        return

    cols = set(df.columns)

    if not REQUIRED_COLUMNS.issubset(cols):
        log.error(
            "Missing columns: {}.".format(
                ', '.join(str(c) for c in REQUIRED_COLUMNS - cols)
            )
        )
        return

    parameters = cols - REQUIRED_COLUMNS
    if len(parameters) < 1:
        log.error("No parameter columns found.")
        return

    if len(df) == 0:
        log.critical("Empty dataset.")
        return

    print(
        "Found data for parameters: {}.".format(
            ', '.join(str(c) for c in parameters)
        )
    )
    print("Converting CSV data into JSON")
    if testrun:
        log.warn("Test run: adding 'TEST_' prefix to parameter names")
        prefix = "TEST_"
    else:
        prefix = ""

    db = km3db.DBManager()    # noqa

    det_id_zero_mask = df['det_id'] == 0
    if sum(det_id_zero_mask) > 0:
        log.warning("Entries with 'det_id=0' found, removing them.")
        df = df[~det_id_zero_mask]
    df['det_id'] = df['det_id'].apply(km3db.tools.todetoid)
    print(df)
    data = convert_runsummary_to_json(df, prefix=prefix)
    print("We have {:.3f} MB to upload.".format(len(data) / 1024**2))

    print("Requesting database session.")
    session_cookie = db.session_cookie

    print("Uploading the data to the database.")
    r = requests.post(
        RUNSUMMARY_URL,
        cookies={"sid": session_cookie},
        files={'datafile': data},
        verify=verify
    )

    if r.status_code == 200:
        log.debug("POST request status code: {}".format(r.status_code))
        print("Database response:")
        db_answer = json.loads(r.text)
        for key, value in db_answer.items():
            print("  -> {}: {}".format(key, value))
        if db_answer['Result'] == 'OK':
            print("Upload successful.")
        else:
            log.critical("Something went wrong.")
    else:
        log.error("POST request status code: {}".format(r.status_code))
        log.critical("Something went wrong...")
        return


def convert_runsummary_to_json(
    df, comment='Uploaded via km3pipe.StreamDS', prefix='TEST_'
):
    """Convert a Pandas DataFrame with runsummary to JSON for DB upload"""
    data_field = []
    comment += ", by {}".format(getpass.getuser())
    for det_id, det_data in df.groupby('det_id'):
        runs_field = []
        data_field.append({"DetectorId": det_id, "Runs": runs_field})

        for run, run_data in det_data.groupby('run'):
            parameters_field = []
            runs_field.append({
                "Run": int(run),
                "Parameters": parameters_field
            })

            parameter_dict = {}
            for row in run_data.iterrows():
                for parameter_name in run_data.columns:
                    if parameter_name in REQUIRED_COLUMNS:
                        continue

                    if parameter_name not in parameter_dict:
                        entry = {'Name': prefix + parameter_name, 'Data': []}
                        parameter_dict[parameter_name] = entry
                    data_value = getattr(row[1], parameter_name)
                    try:
                        data_value = float(data_value)
                    except ValueError as e:
                        log.critical("Data values has to be floats!")
                        raise ValueError(e)
                    value = {
                        'S': str(getattr(row[1], 'source')),
                        'D': data_value
                    }
                    parameter_dict[parameter_name]['Data'].append(value)
            for parameter_data in parameter_dict.values():
                parameters_field.append(parameter_data)
    data_to_upload = {"Comment": comment, "Data": data_field}
    file_data_to_upload = json.dumps(data_to_upload)
    return file_data_to_upload



def main():
    args = docopt(__doc__)

    if args["info"]:
        print_info(args["STREAM"])
    elif args["list"]:
        print_streams()
    elif args['upload']:
        upload_runsummary(args['CSV_FILE'], args['-q'], args['-x'])
    elif args["get"]:
        get_data(args["STREAM"], args["PARAMETERS"], fmt=args["-f"])
    else:
        available_streams()