"""
Utility functions and classes associated with uploading data to GCP, so far.
"""
import json
import os
import re

from google.cloud import bigquery

from simeon.exceptions import (
    BigQueryNameException, MissingSchemaException
)


SEGMENTS = {
    'log': 'TRACKING-LOGS',
    'email': 'EMAIL',
    'sql': 'SQL',
    'rdx': 'RDX',
    'cold': 'COLD',
}
SCHEMA_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), 'schemas'
)
SCHEMAS = {}


def course_to_gcs_folder(course_id: str, file_type: str, bucket: str) -> str:
    """
    Use the given course ID to make a Google Cloud Storage path

    :type course_id: str
    :param course_id: edX course ID to format into a GCS path
    :type file_type: str
    :param file_type: One of sql, log, email, rdx
    :type bucket: str
    :param bucket: A GCS bucket name
    :rtype: str
    :return: A nicely formatted GCS path
    """
    segment = SEGMENTS.get(file_type)
    if not segment:
        raise ValueError(
            'file_type is not one of these: {ft}'.format(
                ft=', '.join(SEGMENTS)
            )
        )
    if not bucket.startswith('gs://'):
        bucket = 'gs://{b}'.format(b=bucket)
    dirname = course_id.replace('/', '__').replace('.', '_')
    return '{b}/{s}/{d}'.format(b=bucket, s=segment, d=dirname)


def local_to_gcs_path(fname: str, file_type: str, bucket: str) -> str:
    """
    Convert the local file name into a GCS path

    :type fname: str
    :param fname: A local file name
    :type file_type: str
    :param file_type: One of sql, log, email, rdx, cold
    :type bucket: str
    :param bucket: A GCS bucket name
    :rtype: str
    :return: A nicely formatted GCS path
    """
    segment = SEGMENTS.get(file_type)
    if not segment:
        msg = 'file_type is not one of these: {s}'
        raise ValueError(msg.format(s=', '.join(SEGMENTS)))
    if not bucket.startswith('gs://'):
        bucket = 'gs://{b}'.format(b=bucket)
    dname, bname = os.path.split(os.path.abspath(os.path.expanduser(fname)))
    if segment == 'COLD':
        gcs_file = bname
    else:
        gcs_file = '{d}/{f}'.format(
            d=os.path.basename(dname).replace('.', '_').replace('-', '_'),
            f=bname
        )
    return '{b}/{s}/{f}'.format(b=bucket, s=segment, f=gcs_file)


def course_to_bq_dataset(course_id: str, file_type: str, project: str) -> str:
    """
    Make a fully qualified BigQuery dataset name with the given info

    :type course_id: str
    :param course_id: edX course ID to format into a GCS path
    :type file_type: str
    :param file_type: One of sql, log, email, rdx
    :type project: str
    :param project: A GCP project ID
    :rtype: str
    :return: BigQuery dataset name with components separated by dots
    """
    if file_type not in SEGMENTS:
        msg = 'file_type is not one of these: {s}'
        raise ValueError(msg.format(s=', '.join(SEGMENTS)))
    if file_type in ('sql',):
        suffix = '_latest'
    elif file_type == 'log':
        suffix = '_logs'
    else:
        suffix = ''
    dataset = '{d}{s}'.format(
        d=course_id.replace('/', '__').replace('.', '_').replace('-', '_'),
        s=suffix
    )
    return '{p}.{d}'.format(p=project, d=dataset)


def local_to_bq_table(fname: str, file_type: str, project: str) -> str:
    """
    Use the given local file to make a fully qualified BigQuery table name

    :type fname: str
    :param fname: A local file name
    :type file_type: str
    :param file_type: One of sql, log, email, rdx
    :type project: str
    :param project: A GCP project ID
    :rtype: str
    :return: BigQuery dataset name with components separated by dots
    """
    if file_type not in SEGMENTS:
        raise ValueError(
            'file_type is not one of these: {s}'.format(
                s=', '.join(SEGMENTS)
            )
        )
    dname, bname = os.path.split(os.path.abspath(os.path.expanduser(fname)))
    if file_type in ('sql',):
        table = bname.split('.', 1)[0].replace('-', '_')
        suffix = '_latest'
    elif file_type in ('log',):
        suffix = '_logs'
        table = 'tracklog_' + ''.join(re.findall(r'\d+', bname))
    else:
        table = bname.split('.', 1)[0].replace('-', '_')
        suffix = ''
    dataset = os.path.basename(dname).replace('.', '_').replace('-', '_')
    if not table.replace('tracklog_', '') or not dataset:
        msg = (
            'A BigQuery table name for file type {t} '
            'could not be constructed with file name {f}. '
            'Please make sure that you are using the right file type.'
        )
        raise BigQueryNameException(msg.format(f=fname, t=file_type))
    return '{p}.{d}{s}.{t}'.format(
        d=dataset, s=suffix, p=project, t=table
    )


def dict_to_schema_field(schema_dict: dict):
    """
    Make a SchemaField from a schema directory

    :type schema_dict: dict
    :param schema_dict: One of the objects in the schema JSON file
    :rtype: bigquery.SchemaField
    :returns: A SchemaField matching the given dictionary's name, type, etc.
    """
    if schema_dict.get('field_type') != 'RECORD':
        return bigquery.SchemaField(**schema_dict)
    fields = []
    for subfield in schema_dict.get('fields', []):
        fields.append(dict_to_schema_field(subfield))
    schema_dict['fields'] = fields
    field = bigquery.SchemaField(**schema_dict)
    return field


def get_bq_schema(table: str, schema_dir: str=SCHEMA_DIR):
    """
    Given a bare table name (without leading project or dataset),
    make a list of bigquery.SchemaField objects to act as the table's schema.

    :type table: str
    :param table: A BigQuery (bare) table name
    :type schema_dir: str
    :param schema_dir: Directory where schema JSON file is looked up
    :rtype: Tuple[List[bigquery.SchemaField], str]
    :return: A 2-tuple with list of bigquery.SchemaField objects and
        a description text for the target table
    :raises: MissingSchemaException
    """
    bname = table.split('.')[-1]
    if all(k in bname for k in ('track', 'log')):
        bname = 'tracking_log'
    targets = ('{t}.json', 'schema_{t}.json')
    for target in targets:
        sfile = os.path.join(schema_dir, target.format(t=bname))
        if not os.path.exists(sfile):
            continue
        out = []
        with open(sfile) as fh:
            schema = json.load(fh)
            for field in schema.get(bname, []):
                out.append(dict_to_schema_field(field))
        if not out:
            msg = (
                'The schema file {f!r} does not contain an object '
                'with a name matching the given table {t}'
            )
            raise MissingSchemaException(msg.format(f=sfile, t=table))
        return out, schema.get('description')
    msg = 'No JSON schema file found for {t} in directory {d}'
    raise MissingSchemaException(msg.format(t=table, d=schema_dir))


def make_bq_load_config(
    table: str, schema_dir=SCHEMA_DIR, append: bool=False,
    create: bool=True, file_format: str='json', delim=',', max_bad_rows=0,
):
    """
    Make a bigquery.LoadJobConfig object and description of a table

    :type table: str
    :param table: Fully qualified table name
    :type schema_dir: str
    :param schema_dir: The directory where schema files live
    :type append: bool
    :param append: Whether to append the loaded to the table
    :type create: bool
    :param create: Whether to create the target table if it does not exist
    :type file_format: str
    :param file_format: One of sql, json, csv, txt
    :type delim: str
    :param delim: The delimiter of the file being loaded
    :rtype: Tuple[bigquery.LoadJobConfig, str]
    :return: A 2-tuple with a bigquery.LoadJobConfig object and
        a description text for the destination table
    """
    schema_dir = schema_dir or SCHEMA_DIR
    schema, desc = get_bq_schema(table, schema_dir=schema_dir)
    if 'json' in file_format.lower():
        format_ = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
        delim = None
        skips = None
    else:
        format_ = bigquery.SourceFormat.CSV
        skips = 1
    if create:
        create = bigquery.CreateDisposition.CREATE_IF_NEEDED
    else:
        create = bigquery.CreateDisposition.CREATE_NEVER
    if append:
        append = bigquery.WriteDisposition.WRITE_APPEND
    else:
        append = bigquery.WriteDisposition.WRITE_TRUNCATE
    if 'json' in file_format.lower():
        config = bigquery.LoadJobConfig(
            schema=schema, source_format=format_,
            create_disposition=create, write_disposition=append,
            max_bad_records=max_bad_rows, ignore_unknown_values=True,
            destination_table_description=desc,
        )
        return config, desc
    config = bigquery.LoadJobConfig(
        schema=schema, source_format=format_,
        create_disposition=create, write_disposition=append,
        field_delimiter=delim, skip_leading_rows=skips,
        max_bad_records=max_bad_rows, ignore_unknown_values=True,
        allow_quoted_newlines=True, destination_table_description=desc,
    )
    return config, desc


def make_bq_query_config(append: bool=False, plain=True, table=None):
    """
    Make a bigquery.QueryJobConfig object to tie to a query to be sent
    to BigQuery for secondary table generation

    :type append: bool
    :param append: Whether to append the loaded to the table
    :type plain: bool
    :param plain: Make an empty QueryJobConfig object
    :type table: Union[None, str]
    :param table: Fully qualified name of a destination table
    :rtype: bigquery.QueryJobConfig
    :returns: Make a bigquery.QueryJobConfig object
    """
    if plain:
        return bigquery.job.QueryJobConfig()
    if append:
        append = bigquery.WriteDisposition.WRITE_APPEND
    else:
        append = bigquery.WriteDisposition.WRITE_TRUNCATE
    config = bigquery.job.QueryJobConfig()
    if table is not None:
        config.destination = bigquery.Table.from_string(table)
    config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED
    config.write_disposition = append
    config.allow_large_results = True
    config.use_legacy_sql = False
    return config


def sqlify_bq_field(field, named=True):
    """
    Convert a bigquery.SchemaField object into a DDL
    column definition.

    :type field: bigquery.SchemaField
    :param field: A SchemaField object to convert to a DDL statement
    :type named: bool
    :param named: Whether the returned statement start with the field name
    :rtype: str
    :returns: A SQL column's DDL statement
    """
    nullability = '' if field.is_nullable else 'NOT NULL'
    if 'INTEGER' in field.field_type:
        type_ = 'INT64'
    elif 'FLOAT' in field.field_type:
        type_ = 'FLOAT64'
    else:
        type_ = field.field_type
    if type_ != 'RECORD':
        if field.mode != 'REPEATED':
            return '{n} {t} {m} OPTIONS(description="""{d}""")'.format(
                n=field.name if named else '',
                t=type_,
                m=nullability,
                d=field.description or '',
            )
        else:
            return '{n} ARRAY<{t}> {m} OPTIONS(description="""{d}""")'.format(
                n=field.name if named else '',
                t=type_,
                m=nullability,
                d=field.description or '',
            )
    if type_ == 'RECORD' and field.mode != 'REPEATED':
        return '{n} STRUCT<{t}> {m} OPTIONS(description="""{d}"""'.format(
            t=',\n\t'.join(sqlify_bq_field(f) for f in field.fields),
            n=field.name if named else '',
            m=nullability,
            d=field.description or '',
        )
    field = bigquery.SchemaField(
        name=field.name, field_type=field.field_type,
        mode='NULLABLE' if field.is_nullable else 'REQUIRED',
        description=field.description, fields=field.fields,
    )
    return '{n} ARRAY<{t}> {m} OPTIONS(description="""{d}""")'.format(
        t=sqlify_bq_field(field, False).lstrip(' '),
        n=field.name,
        m=nullability,
        d=field.description or '',
    )
