import re
import inspect
import jsonschema
import numpy as np
from os import path as os_path
from asdf.tags.core import ndarray

from stdatamodels import validate
from stdatamodels import properties
from stdatamodels import schema as mschema
from . import (ImageModel, IFUImageModel)

from ..pipeline import (GuiderPipeline, DarkPipeline, Ami3Pipeline,
                        Tso3Pipeline, Coron3Pipeline,
                        Detector1Pipeline, Image2Pipeline,
                        Image3Pipeline, Spec2Pipeline, Spec3Pipeline)

DEBUG = 0

comment = re.compile(r'#.*')
continued = re.compile(r'\\\s*$')
assignment = re.compile(r'=(?!=)')

single_quote = re.compile(r"'[^']+'")
double_quote = re.compile(r'"[^"]+"')
triple_quote = re.compile(r'""".*?"""', flags=re.DOTALL)

relative_import = re.compile(r'\s*from\s+\.+[\.\w]*\s+import')

def main(instrument, mode, level):
    """
    Main procedure for make_header
    """
    cls = get_class(instrument, mode)
    im = cls()

    defaults = set_default_values(im, instrument, mode, level)
    pipeline = choose_pipeline(defaults)

    add_metadata(im, defaults, pipeline)
    fill_extensions(im, defaults)

    return im

def add_metadata(im, defaults, pipeline):
    """
    Add metadata values to match values used in pipeline
    """
    keywords = undefined_metadata(pipeline, 'meta')
    for keyword in keywords:
        if keyword == 'meta':
            pass
        elif keyword == 'meta.model_type':
            pass
        elif keyword == 'meta.wcs':
            set_value(im, defaults, 'meta.wcsinfo')
        elif not another_instrument(defaults, im, keyword):
            set_value(im, defaults, keyword)
            additional_metadata(im, defaults, keyword)

def additional_metadata(im, defaults, keyword):
    """
    Add additional metadata items from the same group
    """
    groups = []

    groups.append(set(('meta.subarray.name',
                      'meta.subarray.xstart',
                      'meta.subarray.ystart',
                      'meta.subarray.xsize',
                      'meta.subarray.ysize')))

    groups.append(set(('meta.instrument.name',
                      'meta.instrument.band',
                      'meta.instrument.channel',
                      'meta.instrument.module',
                      'meta.instrument.detector',
                      'meta.instrument.filter',
                      'meta.instrument.grating',
                      'meta.instrument.pupil',
                      'meta.exposure.readpatt')))

    groups.append(set(('meta.asn.pool_name',
                      'meta.asn.table_name')))

    groups.append(set(('meta.target.catalog_name',
                      'meta.target.ra',
                      'meta.target.dec')))

    groups.append(set(('meta.exposure.start_time',
                      'meta.exposure.mid_time',
                      'meta.exposure.end_time',
                      'meta.exposure.exposure_time')))

    groups.append(set(('meta.exposure.frame_divisor',
                      'meta.exposure.frame_time',
                      'meta.exposure.groupgap',
                      'meta.exposure.group_time',
                      'meta.exposure.integration_time',
                      'meta.exposure.nframes',
                      'meta.exposure.ngroups',
                      'meta.exposure.nints',
                      'meta.exposure.nresets_at_start',
                      'meta.exposure.nresets_between_ints')))

    for group in groups:
        if keyword in group:
            for additional_key in group:
                if not another_instrument(defaults, im, additional_key):
                    if additional_key != keyword:
                        set_value(im, defaults, additional_key)

def another_instrument(defaults, im, keyword):
    """
    Return True if metadata item is for
    another instrument or exposure type
    """
    modes = ((('NIRCAM',), 'NIRCAM'),
             (('NIRSPEC',), 'NIRSPEC'),
             (('NIRSPEC',), 'IRS2'),
             (('MIRI',), 'MIRI'),
             (('FGS',), 'FGS'),
             (('NIRISS',), 'NIRISS'),
             (('NRC_FOCUS',), 'WFSC'),
             (('NRS_IFU',), 'IFU'),
             (('NRS_MSASPEC',), 'MSA'),
             (('MIR_MRS',), 'MRS'),
             (('MIR_LRS-FIXEDSLIT', 'MIR_LRS-SLITLESS'), 'LRS'),
            )

    instrument = defaults['meta.instrument.name']
    exposure_type = defaults['meta.exposure.type']

    schema = im._schema
    for attr in keyword.split('.'):
        schema = properties._get_schema_for_property(schema, attr)
        if schema:
            title = schema.get('title', '').upper()
            for mode in modes:
                (mode_name, mode_value) = mode
                if not (instrument in mode_name or
                        exposure_type in mode_name):
                    if title.find(mode_value) >= 0:
                        return True
        else:
            break

    return False

def build_exposure_type(instrument, mode):
    """
    Build the exposure type from the instrument name and mode
    """
    abbreviations = {'FGS':'FGS', 'MIRI':'MIR', 'NIRISS':'NIS',
                     'NIRCAM':'NRC', 'NIRSPEC':'NRS'}

    spec_mode = {'FGS':'', 'MIRI':'MRS', 'NIRISS':'WFSS',
                 'NIRCAM':'WFSS', 'NIRSPEC':'FIXEDSLIT'}

    abbrev = abbreviations[instrument]
    if mode[0:4] == 'IMAG':
        mode = 'IMAGE'
        mode = '_'.join((abbrev, mode))
    elif mode[0:4] == 'SPEC':
        mode = spec_mode[instrument]
        mode = '_'.join((abbrev, mode))

    return mode

def choose_pipeline(defaults):
    """
    Select the pipeline class from the exposure type and level
    """
    mode = {'FGS_ACQ1':1, 'FGS_ACQ2':1, 'FGS_DARK':2,
            'FGS_FINEGUIDE':1, 'FGS_FOCUS':6, 'FGS_ID-IMAGE':1,
            'FGS_ID-STACK':1, 'FGS_IMAGE':6, 'FGS_INTFLAT':0,
            'FGS_SKYFLAT':6, 'FGS_TRACK':1, 'MIR_IMAGE':6,
            'MIR_TACQ':0, 'MIR_LYOT':5, 'MIR_4QPM':5,
            'MIR_LRS-FIXEDSLIT':9, 'MIR_LRS-SLITLESS':4, 'MIR_MRS':9,
            'MIR_DARK':2, 'MIR_FLAT-IMAGE':6, 'MIR_FLATIMAGE':6,
            'MIR_FLAT-MRS':6, 'MIR_FLATMRS':6, 'MIR_CORONCAL':7,
            'NIS_AMI':3, 'NIS_DARK':2, 'NIS_FOCUS':6, 'NIS_IMAGE':6,
            'NIS_LAMP':0, 'NIS_SOSS':4, 'NIS_TACQ':6,
            'NIS_TACONFIRM':6, 'NIS_WFSS':9, 'N/A':6, 'ANY':6,
            'NRC_IMAGE':6, 'NRC_GRISM':9, 'NRC_TACQ':6,
            'NRC_CORON':5, 'NRC_FOCUS':6, 'NRC_DARK':2, 'NRC_FLAT':6,
            'NRC_LED':0, 'NRC_TACONFIRM':6, 'NRC_TSIMAGE':4,
            'NRC_WFSS':9, 'NRC_TSGRISM':4, 'NRS_AUTOFLAT':0,
            'NRS_AUTOWAVE':0, 'NRS_BRIGHTOBJ':4, 'NRS_CONFIRM':0,
            'NRS_DARK':2, 'NRS_FIXEDSLIT':9, 'NRS_FOCUS':0, 'NRS_IFU':9,
            'NRS_IMAGE':6, 'NRS_LAMP':0, 'NRS_MIMF':0, 'NRS_MSASPEC':9,
            'NRS_MSATA':0, 'NRS_TACONFIRM':0, 'NRS_TACQ':0, 'NRS_TASLIT':9,
            'NRS_WATA':0}

    pipelines = [None, GuiderPipeline, DarkPipeline,
                 Ami3Pipeline, Tso3Pipeline, Coron3Pipeline,
                 Detector1Pipeline, Image2Pipeline, Image3Pipeline,
                 Detector1Pipeline, Spec2Pipeline, Spec3Pipeline]

    exposure_type = defaults['meta.exposure.type']
    index = mode[exposure_type]
    if index == 0:
        raise ValueError("Pipeline not defined for mode " +
                         exposure_type)
    if index > 5:
        level = defaults['meta.pipeline_level']
        index += level - 1

    return pipelines[index]

def drop_prefix(string, name):
    """
    Drop the model name preceding the metadata keyword
    """
    parts = name.split('.')
    for i, part in enumerate(parts):
        if part == string:
            return '.'.join(parts[i:])
    return ''

def drop_suffix(name):
    """
    Drop a method name from the keyword name
    """
    parts = name.split('.')
    while parts:
        if parts[-1][0] != '_' and parts[-1][-1] != '(':
            break
        parts.pop()

    return '.'.join(parts)

def fill_extensions(im, defaults):
    """
    Set image extensions to their default values
    """
    def fill_hdu(subschema, path, combiner, shape, recurse):
        if 'fits_hdu' not in subschema:
            return

        dtype = subschema.get('datatype')
        if dtype is None:
            return

        ndim = subschema.get('ndim')
        if ndim and ndim != len(shape):
            return

        dtype = ndarray.asdf_datatype_to_numpy_dtype(dtype)

        keyword = '.'.join(path)
        default = subschema.get('default', 0.0)
        im[keyword] = np.full(shape, default, dtype=dtype)

    shape = get_shape(im, defaults)
    mschema.walk_schema(im._schema, fill_hdu, shape)

def get_class(instrument, mode):
    """
    Get the model class from the exposure type
    """
    instrument = instrument.upper()
    mode = mode.upper()
    exposure_type = build_exposure_type(instrument, mode)

    if exposure_type in ('NRS_IFU', 'MIR_MRS'):
        cls = IFUImageModel
    else:
        cls = ImageModel
    return cls

def get_package(obj):
    """
    Get the package name associated with an object
    """
    if inspect.isclass(obj):
        path = obj.__module__
    elif inspect.isfunction(obj):
        path = obj.__module__
    elif inspect.ismodule(obj):
        path = obj.__name__
    else:
        path = ''
    path = path.split('.')

    # Do not pop the module name if the source
    # comes from an `_init_.py` file.
    source_file = inspect.getfile(obj)
    if os_path.basename(source_file) != '__init__.py':
        path.pop()

    return '.'.join(path)

def get_shape(im, defaults):
    """
    Get the shape of the image
    """
    shape = im.shape
    if not shape or shape[0] == 0:
        xsize = defaults['meta.subarray.xsize']
        ysize = defaults['meta.subarray.ysize']
        shape = (xsize, ysize)

    return shape

def get_subschema(im, keyword):
    """
    Get the subschema for a keyword or None if not found
    """
    schema = im._schema
    for attr in keyword.split('.'):
        schema = properties._get_schema_for_property(schema, attr)
        if schema is None:
            break
    return schema

def import_objects(package, line):
    """
    Run an import command and save any objects generated by it
    """
    line = line.lstrip()
    local_variables = {}
    global_variables = {'__package__':package}
    try:
        exec(line, global_variables, local_variables) # nosec
    except ImportError:
        pass

    objects = []
    for value in local_variables.values():
        if (inspect.isclass(value) or
            inspect.isfunction(value) or
            inspect.ismodule(value)):
            objects.append(value)
    return objects

def log_file(filename):
    """
    Log files checked to a log file
    """
    if DEBUG:
        path = filename.split('/')
        try:
            jpos = path.index('jwst')
        except ValueError:
            jpos = 0
        filename = '/'.join(path[jpos:])
        print("Searching " + '/'.join(path[jpos:]))

def log_results(title, strings):
    """
    Write intermediate results to a log file
    """
    if DEBUG:
        print("\n=== %s ===\n" % title)
        for s in strings:
            print(s)

def next_line(fd):
    """
    Read the next line from a file, including any continuations
    """
    qcount = 0
    pcount = 0
    more = True
    long_line = ''

    while more:
        more = False
        line = fd.readline()
        if line:
            qcount += line.count('"""')
            qcount = qcount % 2
            if qcount > 0:
                more = True

            else:
                line = unquote(line)
                if continued.search(line):
                    more = True

                pcount += line.count('(')
                pcount -= line.count(')')
                if pcount > 0:
                    more = True

        long_line += line

    return unquote(long_line)

def search_file(source_file, package, string, seen):
    """
    Search a file for lines containing a string
    """
    matches = []
    with open(source_file, 'r') as fd:
        while True:
            line = next_line(fd)
            if not line:
                break

            import_match = relative_import.match(line)
            if import_match:
                for obj in import_objects(package, line):
                    new_matches = search_source_lines(obj, string, seen)
                    matches.extend(new_matches)
            elif line.find(string) >= 0:
                matches.append(line)
    return matches

def search_source_lines(obj, string, seen):
    """
    Search all files contained in a pipeline for a string
    """
    matches = []
    package = get_package(obj)
    if package:
        source_file = inspect.getfile(obj)
        if source_file not in seen:
            seen.add(source_file)
            if source_file.endswith('.py'):
                log_file(source_file)
                new_matches = search_file(source_file, package,
                                          string, seen)
                matches.extend(new_matches)
    return matches

def set_default_values(im, instrument, mode, level):
    """
    Set the default values used to set header keywords
    """
    defaults = {}
    set_standard_defaults(im, defaults, instrument, mode, level)
    set_pipeline_defaults(im, defaults)
    set_wcs_defaults(im, defaults)
    return defaults

def set_pipeline_defaults(im, defaults):
    """
    Set header keyword defaults associated with a pipeline
    """
    generic_defaults = {'meta.telescope': 'JWST',
                        'meta.subarray.name': 'FULL',
                        'meta.subarray.xstart': 1,
                        'meta.subarray.ystart': 1,
                        'meta.target.catalog_name': 'SMC',
                        'meta.target.ra': 5.3196,
                        'meta.target.dec': -72.98605,
                        'meta.target.type': 'FIXED',
                        'meta.target.source_type': 'EXTENDED'
                        }

    nircam_defaults = {'meta.subarray.xsize': 2048,
                       'meta.subarray.ysize': 2048,
                       'meta.subarray.fastaxis': 1,
                       'meta.subarray.slowaxis': -2,
                       'meta.instrument.channel': 'SHORT',
                       'meta.instrument.module': 'B',
                       'meta.instrument.detector': 'NRCB1',
                       'meta.instrument.filter': 'F115W',
                       'meta.instrument.pupil': 'CLEAR'
                      }

    nirspec_defaults = {'meta.subarray.xsize': 2048,
                        'meta.subarray.ysize': 2048,
                        'meta.subarray.fastaxis': 1,
                        'meta.subarray.slowaxis': 2,
                        'meta.instrument.detector': 'NRS1',
                        'meta.instrument.filter': 'F100LP',
                        'meta.instrument.grating': 'MIRROR',
                        'meta.exposure.readpatt': 'NRSRAPID'
                       }

    miri_defaults = {'meta.subarray.xsize': 1032,
                     'meta.subarray.ysize': 1024,
                     'meta.subarray.fastaxis': 1,
                     'meta.subarray.slowaxis': 2,
                     'meta.subarray.name': 'FULL',
                     'meta.instrument.channel': '34',
                     'meta.instrument.band': 'SHORT',
                     'meta.instrument.detector': 'MIRIMAGE',
                     'meta.instrument.filter': 'N/A',
                     }

    fgs_defaults = {'meta.subarray.xsize': 2048,
                    'meta.subarray.ysize': 2048,
                    'meta.subarray.fastaxis': 2,
                    'meta.subarray.slowaxis': -1,
                    'meta.instrument.detector': 'GUIDER2'
                    }

    niriss_defaults = {'meta.subarray.xsize': 2048,
                       'meta.subarray.ysize': 2048,
                       'meta.subarray.fastaxis': -2,
                       'meta.subarray.slowaxis': -1,
                       'meta.instrument.detector': 'NIS',
                       'meta.instrument.filter': 'F380M'
                       }

    any_defaults = {}

    instrument_defaults = {'NIRCAM': nircam_defaults,
                           'NIRSPEC': nirspec_defaults,
                           'MIRI': miri_defaults,
                           'FGS': fgs_defaults,
                           'NIRISS': niriss_defaults,
                           'ANY': any_defaults
                           }

    mir_slitless_defaults = {'meta.subarray.name': 'SLITLESSPRISM',
                             'meta.subarray.xstart': 1,
                             'meta.subarray.ystart': 321,
                             'meta.subarray.xsize': 68,
                             'meta.subarray.ysize': 1024
                            }

    mir_fixedslit_defaults = {'meta.instrument.detector': 'MIRIFULONG'}

    mir_mrs_defaults = {'meta.instrument.detector': 'MIRIFULONG'}

    exposure_type_defaults = {'MIR_LRS-SLITLESS': mir_slitless_defaults,
                              'MIR_LRS-FIXEDSLIT': mir_fixedslit_defaults,
                              'MIR_MRS': mir_mrs_defaults
                             }


    defaults.update(generic_defaults);
    instrument = defaults['meta.instrument.name']
    defaults.update(instrument_defaults[instrument])
    exposure_type = defaults['meta.exposure.type']
    exposure_deltas = exposure_type_defaults.get(exposure_type)
    if exposure_deltas is not None:
        defaults.update(exposure_deltas)

    shape = im.shape
    if shape:
        if shape[0] != 0:
            defaults['meta.subarray.xsize'] = shape[0]
        if shape[1] != 0:
            defaults['meta.subarray.ysize'] = shape[1]

def set_standard_defaults(im, defaults, instrument, mode, level):
    """
    Set keyword defaults from instrument, mode, and level
    """
    instrument = instrument.upper()
    if validate_value(im, 'meta.instrument.name', instrument):
        defaults['meta.instrument.name'] = instrument
    else:
        raise ValueError('Unrecognized instrument name: ' + instrument)

    if defaults['meta.instrument.name'] == 'ANY':
        mode = 'ANY'
    else:
        mode = mode.upper()
    if validate_value(im, 'meta.exposure.type', mode):
        defaults['meta.exposure.type'] = mode
    else:
        value = build_exposure_type(instrument, mode)
        if not validate_value(im, 'meta.exposure.type', value):
            value = None

        if value is None:
            raise ValueError('Unrecognized mode: ' + mode)
        else:
            defaults['meta.exposure.type'] = value

    try:
        value = int(level)
    except ValueError:
        match = re.search(r'(\d+)', level)
        if match:
            value = int(match.group(1))
        else:
            value = 0
    if value < 1 or value > 3:
        raise ValueError('Unrecognized level: ' + level)
    else:
        defaults['meta.pipeline_level'] = value

def set_value(im, defaults, keyword):
    """
    Set a value for a header keyword from defaults and data type
    """
    schema = get_subschema(im, keyword)
    val = None
    if schema:
        if keyword in defaults:
            val = defaults[keyword]
        elif 'default' in schema:
            val = schema['default']
        elif 'enum' in schema:
            for choice in ('N/A', 'NONE', 'ANY'):
                try:
                    i = schema['enum'].index(choice)
                    break
                except ValueError:
                    i = 0
            val = schema['enum'][i]
        elif 'type' in schema:
            typ = schema['type']
            if typ == 'string':
                if keyword.startswith('meta.cal_step'):
                    val = 'COMPLETE'
                else:
                    val = ''
            elif typ == 'number':
                val = 0.0
            elif typ == 'integer':
                val = 1
            elif typ == 'boolean':
                val = True
            elif typ == 'object':
                for field in schema['properties']:
                    subkey = '.'.join((keyword, field))
                    set_value(im, defaults, subkey)
            else:
                val = None
        else:
            val = None

    if val is not None:
        im[keyword] = val

def set_wcs_defaults(im, defaults):
    """
    Set default values for the wcs header keywords
    """
    wcs_defaults = {
                    'crval1':5.3196,     'crval2':-72.98605,
                    'ctype1':'RA---TAN', 'ctype2':'DEC--TAN',
                    'cunit1':'deg',      'cunit2':'deg',
                    'cdelt1':5.5E-06,    'cdelt2':5.5E-06}

    shape = get_shape(im, defaults)
    for i in (1, 2):
        name = "crpix%d" % i
        wcs_defaults[name] = 0.5 * shape[i-1]

        for j in (1,2):
            name = "pc%d_%d" % (i, j)
            if i == j:
                wcs_defaults[name] = 1.0
            else:
                wcs_defaults[name] = 0.0

    for name in wcs_defaults:
        wcs_name = 'meta.wcsinfo.' + name
        defaults[wcs_name] = wcs_defaults[name]

def undefined_metadata(module, string):
    """
    Find the metadata keywords that are undefined in the pipeline
    """
    name_set = []
    for i in range(2):
        name_set.append(set())

    pat = re.compile(r'([\w\.]+' + string + r'[\w\.]+\(?)')
    seen = set()
    matched = search_source_lines(module,
                                   '.' + string + '.',
                                   seen)
    log_results("matched lines", matched)

    for line in matched:
        sides = assignment.split(line)
        if len(sides) == 1:
            sides.append(sides[0][:])
            sides[0] = ''

        j = 1
        for i in range(len(sides), 0, -1):
            pat_match = pat.search(sides[i-1])
            if pat_match is not None:
                for name in pat_match.groups():
                    if name is not None:
                        name = drop_suffix(name)
                        name_set[j].add(name)
            j = 0

    log_results("lhs metadata", name_set[0])
    log_results("rhs metadata", name_set[1])

    undefined = set()
    for name in name_set[1]:
        name = drop_prefix(string, name)
        if name:
            undefined.add(name)

    log_results("undefined metadata", undefined)
    return undefined

def unquote(line):
    """
    Remove quoted strings and comments from a line
    """
    line = triple_quote.sub('', line)
    line = double_quote.sub('""', line)
    line = single_quote.sub("''", line)
    line = comment.sub('', line)
    return line

def validate_value(im, name, value):
    """
    Test a value against the schema to see if it is valid
    """
    valid = True
    schema = get_subschema(im, name)
    if schema:
        try:
            validate._check_value(value, schema, im)
        except jsonschema.ValidationError:
            valid = False
    return valid
