import struct
import numpy as np
from astropy import units as u
from astropy.coordinates import Angle
import os
import six

try:
    import h5py
    HAS_HDF5 = True
except ImportError:
    HAS_HDF5 = False

###
# Header parsing
###

# Dictionary of allowed keywords and their types
# Here are the keywordss that a filter bank file may
# contain.  Items marked with "[*]" are not yet # supported.  See docs for
# indivisuabl attribtues for more detailed info.
#
#   * telescope_id (int): 0=fake data; 1=Arecibo; 2=Ooty... others to be added
#   * machine_id (int): 0=FAKE; 1=PSPM; 2=WAPP; 3=OOTY... others to be added
#   * data_type (int): 1=blimpy; 2=time series... others to be added
#   * rawdatafile (string): the name of the original data file
#   * source_name (string): the name of the source being observed by the telescope
#   * barycentric (int): equals 1 if data are barycentric or 0 otherwise
#   * pulsarcentric (int): equals 1 if data are pulsarcentric or 0 otherwise
#   * az_start (double): telescope azimuth at start of scan (degrees)
#   * za_start (double): telescope zenith angle at start of scan (degrees)
#   * src_raj (double): right ascension (J2000) of source (hours, converted from hhmmss.s)
#   * src_dej (double): declination (J2000) of source (degrees, converted from ddmmss.s)
#   * tstart (double): time stamp (MJD) of first sample
#   * tsamp (double): time interval between samples (s)
#   * nbits (int): number of bits per time sample
#   * nsamples (int): number of time samples in the data file (rarely used any more)
#   * fch1 (double): centre frequency (MHz) of first blimpy channel
#   * foff (double): blimpy channel bandwidth (MHz)
#   * FREQUENCY_START [*] (character): start of frequency table (see below for explanation)
#   * fchannel [*] (double): frequency channel value (MHz)
#   * FREQUENCY_END [*] (character): end of frequency table (see below for explanation)
#   * nchans (int): number of blimpy channels
#   * nifs (int): number of seperate IF channels
#   * refdm (double): reference dispersion measure (pc/cm**3)
#   * period (double): folding period (s)
#   * nbeams (int):total number of beams (?)
#   * ibeam (int): number of the beam in this file (?)

header_keyword_types = {
    b'telescope_id' : b'<l',
    b'machine_id'   : b'<l',
    b'data_type'    : b'<l',
    b'barycentric'  : b'<l',
    b'pulsarcentric': b'<l',
    b'nbits'        : b'<l',
    b'nsamples'     : b'<l',
    b'nchans'       : b'<l',
    b'nifs'         : b'<l',
    b'nbeams'       : b'<l',
    b'ibeam'        : b'<l',
    b'rawdatafile'  : b'str',
    b'source_name'  : b'str',
    b'az_start'     : b'<d',
    b'za_start'     : b'<d',
    b'tstart'       : b'<d',
    b'tsamp'        : b'<d',
    b'fch1'         : b'<d',
    b'foff'         : b'<d',
    b'refdm'        : b'<d',
    b'period'       : b'<d',
    b'src_raj'      : b'angle',
    b'src_dej'      : b'angle',
    }

def len_header(filename):
    """ Return the length of the blimpy header, in bytes

    Args:
        filename (str): name of file to open

    Returns:
        idx_end (int): length of header, in bytes
    """
    with  open(filename, 'rb') as f:
        header_sub_count = 0
        eoh_found = False
        while not eoh_found:
            header_sub = f.read(512)
            header_sub_count += 1
            if b'HEADER_END' in header_sub:
                idx_end = header_sub.index(b'HEADER_END') + len(b'HEADER_END')
                eoh_found = True
                break

        idx_end = (header_sub_count -1) * 512 + idx_end
    return idx_end

def read_next_header_keyword(fh):
    """

    Args:
        fh (file): file handler

    Returns:
    """
    n_bytes = np.fromstring(fh.read(4), dtype='uint32')[0]
    #print n_bytes

    if n_bytes > 255:
        n_bytes = 16

    keyword = fh.read(n_bytes)

    #print keyword

    if keyword == b'HEADER_START' or keyword == b'HEADER_END':
        return keyword, 0, fh.tell()
    else:
        dtype = header_keyword_types[keyword]
        #print dtype
        idx = fh.tell()
        if dtype == b'<l':
            val = struct.unpack(dtype, fh.read(4))[0]
        if dtype == b'<d':
            val = struct.unpack(dtype, fh.read(8))[0]
        if dtype == b'str':
            str_len = np.fromstring(fh.read(4), dtype='uint32')[0]
            val = fh.read(str_len)
        if dtype == b'angle':
            val = struct.unpack('<d', fh.read(8))[0]
            val = fil_double_to_angle(val)
            if keyword == b'src_raj':
                val = Angle(val, unit=u.hour)
            else:
                val = Angle(val, unit=u.deg)
        return keyword, val, idx

def is_filterbank(filename):
    """ Open file and confirm if it is a filterbank file or not. """
    with open(filename, 'rb') as fh:
        is_fil = True

        # Check this is a blimpy file
        try:
            keyword, value, idx = read_next_header_keyword(fh)
            try:
                assert keyword == b'HEADER_START'
            except AssertionError:
                is_fil = False
        except KeyError:
            is_fil = False
        return is_fil

def read_header(filename, return_idxs=False):
    """ Read blimpy header and return a Python dictionary of key:value pairs

    Args:
        filename (str): name of file to open

    Optional args:
        return_idxs (bool): Default False. If true, returns the file offset indexes
                            for values

    returns

    """
    with open(filename, 'rb') as fh:
        header_dict = {}
        header_idxs = {}

        # Check this is a blimpy file
        keyword, value, idx = read_next_header_keyword(fh)

        try:
            assert keyword == b'HEADER_START'
        except AssertionError:
            raise RuntimeError("Not a valid blimpy file.")

        while True:
            keyword, value, idx = read_next_header_keyword(fh)
            if keyword == b'HEADER_END':
                break
            else:
                header_dict[keyword] = value
                header_idxs[keyword] = idx

    if return_idxs:
        return header_idxs
    else:
        return header_dict

def fix_header(filename, keyword, new_value):
    """ Apply a quick patch-up to a Filterbank header by overwriting a header value

    Args:
        filename (str): name of file to open and fix. WILL BE MODIFIED.
        keyword (stt):  header keyword to update
        new_value (long, double, angle or string): New value to write.

    Notes:
        This will overwrite the current value of the blimpy with a desired
        'fixed' version. Note that this has limited support for patching
        string-type values - if the length of the string changes, all hell will
        break loose.

    """

    # Read header data and return indexes of data offsets in file
    hd = read_header(filename)
    hi = read_header(filename, return_idxs=True)
    idx = hi[keyword]

    # Find out the datatype for the given keyword
    dtype = header_keyword_types[keyword]
    dtype_to_type = {b'<l'  : np.int32,
                     b'str' : bytes,
                     b'<d'  : np.float64,
                     b'angle' : to_sigproc_angle}
    value_dtype = dtype_to_type[dtype]

    # Generate the new string
    if isinstance(value_dtype, bytes):
        if len(hd[keyword]) == len(new_value):
            val_str = np.int32(len(new_value)).tostring() + new_value
        else:
            raise RuntimeError("String size mismatch. Cannot update without rewriting entire file.")
    else:
        val_str = value_dtype(new_value).tostring()

    # Write the new string to file
    with open(filename, 'rb+') as fh:
        fh.seek(idx)
        fh.write(val_str)

def fil_double_to_angle(angle):
      """ Reads a little-endian double in ddmmss.s (or hhmmss.s) format and then
      converts to Float degrees (or hours).  This is primarily used to read
      src_raj and src_dej header values. """

      negative = (angle < 0.0)
      angle = np.abs(angle)

      dd = np.floor((angle / 10000))
      angle -= 10000 * dd
      mm = np.floor((angle / 100))
      ss = angle - 100 * mm
      dd += mm/60.0 + ss/3600.0

      if negative:
          dd *= -1

      return dd

###
# sigproc writing functions
###

def to_sigproc_keyword(keyword, value=None):
    """ Generate a serialized string for a sigproc keyword:value pair

    If value=None, just the keyword will be written with no payload.
    Data type is inferred by keyword name (via a lookup table)

    Args:
        keyword (str): Keyword to write
        value (None, float, str, double or angle): value to write to file

    Returns:
        value_str (str): serialized string to write to file.
    """

    keyword = bytes(keyword)

    if value is None:
        return np.int32(len(keyword)).tostring() + keyword
    else:
        dtype = header_keyword_types[keyword]

        dtype_to_type = {b'<l'  : np.int32,
                         b'str' : str,
                         b'<d'  : np.float64,
                         b'angle' : to_sigproc_angle}

        value_dtype = dtype_to_type[dtype]

        if value_dtype is str:
            return np.int32(len(keyword)).tostring() + keyword + np.int32(len(value)).tostring() + value
        else:
            return np.int32(len(keyword)).tostring() + keyword + value_dtype(value).tostring()

def generate_sigproc_header(f):
    """ Generate a serialzed sigproc header which can be written to disk.

    Args:
        f (Filterbank object): Filterbank object for which to generate header

    Returns:
        header_str (str): Serialized string corresponding to header
    """

    header_string = b''
    header_string += to_sigproc_keyword(b'HEADER_START')

    for keyword in f.header.keys():
        if keyword == b'src_raj':
            header_string += to_sigproc_keyword(b'src_raj')  + to_sigproc_angle(f.header[b'src_raj'])
        elif keyword == b'src_dej':
            header_string += to_sigproc_keyword(b'src_dej')  + to_sigproc_angle(f.header[b'src_dej'])
        elif keyword == b'az_start' or keyword == b'za_start':
            header_string += to_sigproc_keyword(keyword)  + np.float64(f.header[keyword]).tostring()
        elif keyword not in header_keyword_types.keys():
            pass
        else:
            header_string += to_sigproc_keyword(keyword, f.header[keyword])

    header_string += to_sigproc_keyword(b'HEADER_END')
    return header_string

def to_sigproc_angle(angle_val):
    """ Convert an astropy.Angle to the ridiculous sigproc angle format string. """
    x = str(angle_val)

    if '.' in x:
        if 'h' in x:
                d, m, s, ss = int(x[0:x.index('h')]), int(x[x.index('h')+1:x.index('m')]), \
                int(x[x.index('m')+1:x.index('.')]), float(x[x.index('.'):x.index('s')])
        if 'd' in x:
            d, m, s, ss = int(x[0:x.index('d')]), int(x[x.index('d')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('.')]), float(x[x.index('.'):x.index('s')])
    else:
        if 'h' in x:
            d, m, s = int(x[0:x.index('h')]), int(x[x.index('h')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('s')])
        if 'd' in x:
            d, m, s = int(x[0:x.index('d')]), int(x[x.index('d')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('s')])
        ss = 0
    num = str(d).zfill(2) + str(m).zfill(2) + str(s).zfill(2)+ '.' + str(ss).split(".")[-1]
    return np.float64(num).tostring()

def calc_n_ints_in_file(filename):
    """ Calculate number of integrations in a given file """
    # Load binary data
    h = read_header(filename)
    n_bytes  = int(h[b'nbits'] / 8)

    n_chans = h[b'nchans']
    n_ifs   = h[b'nifs']
    idx_data = len_header(filename)
    f = open(filename, 'rb')
    f.seek(idx_data)
    filesize = os.path.getsize(filename)
    n_bytes_data = filesize - idx_data

    if h[b'nbits'] == 2:
        n_ints = int(4 * n_bytes_data / (n_chans * n_ifs))
    else:
        n_ints = int(n_bytes_data / (n_bytes * n_chans * n_ifs))

    return n_ints


#!/usr/bin/env python
'''
    Simple script for quicly making an .fil file from a .h5.

    ..author: Emilio Enriquez (jeenriquez@gmail.com)

    July 28th 2017
'''

try:
    from .waterfall import Waterfall
except:
    from waterfall import Waterfall

from optparse import OptionParser
import sys
import os

#------
# Logging set up
import logging
logger = logging.getLogger(__name__)

level_log = logging.INFO

if level_log == logging.INFO:
    stream = sys.stdout
    format = '%(name)-15s %(levelname)-8s %(message)s'
else:
    stream =  sys.stderr
    format = '%%(relativeCreated)5d (name)-15s %(levelname)-8s %(message)s'

logging.basicConfig(format=format,stream=stream,level = level_log)
#------

def make_fil_file(filename,out_dir='./', new_filename=None, max_load = None):
    ''' Converts file to Sigproc filterbank (.fil) format.  Default saves output in current dir.
    '''

    fil_file = Waterfall(filename, max_load = max_load)
    if not new_filename:
        new_filename = out_dir+filename.replace('.h5','.fil').split('/')[-1]

    if '.fil' not in new_filename:
        new_filename = new_filename+'.fil'

    fil_file.write_to_fil(new_filename)

def cmd_tool():
    '''
    '''

    p = OptionParser()
    p.set_usage('Command line utility for converting HDF5 (.h5) to Sigproc filterbank (.fil) format \n >>h52fil <FULL_PATH_TO_FIL_FILE> [options]')
    p.add_option('-o', '--out_dir', dest='out_dir', type='str', default='./', help='Location for output files. Default: local dir. ')
    p.add_option('-n', '--new_filename', dest='new_filename', type='str', default='', help='New filename. Default: replaces extention to .fil')
    p.add_option('-d', '--delete_input', dest='delete_input', action='store_true', default=False, help='This option deletes the input file after conversion.')
    p.add_option('-l', action='store', default=None, dest='max_load', type=float,help='Maximum data limit to load. Default:1GB')

    opts, args = p.parse_args(sys.argv[1:])

    if len(args)!=1:
        logger.info('Please specify a file name \nExiting.')
        sys.exit()
    else:
        filename = args[0]

    make_fil_file(filename, out_dir = opts.out_dir, new_filename=opts.new_filename, max_load = opts.max_load)

    if opts.delete_input:
        logger.info("'Deleting input file: %s"%(filename))
        os.remove(filename)

if __name__ == "__main__":

    cmd_tool()


import re
import sys
from types import CodeType
from types import TracebackType

try:
    from __pypy__ import tproxy
except ImportError:
    tproxy = None
try:
    from .cpython import tb_set_next
except ImportError:
    tb_set_next = None

if not tb_set_next and not tproxy:
    raise ImportError("Cannot use tblib. Runtime not supported.")

__version__ = '1.3.2'
__all__ = 'Traceback',

PY3 = sys.version_info[0] == 3
FRAME_RE = re.compile(r'^\s*File "(?P<co_filename>.+)", line (?P<tb_lineno>\d+)(, in (?P<co_name>.+))?$')

class _AttrDict(dict):
    __slots__ = ()
    __getattr__ = dict.__getitem__

# noinspection PyPep8Naming
class __traceback_maker(Exception):
    pass

class TracebackParseError(Exception):
    pass

class Code(object):
    def __init__(self, code):
        self.co_filename = code.co_filename
        self.co_name = code.co_name

class Frame(object):
    def __init__(self, frame):
        self.f_globals = {
            k: v
            for k, v in frame.f_globals.items()
            if k in ("__file__", "__name__")
        }
        self.f_code = Code(frame.f_code)

    def clear(self):
        # For compatibility with PyPy 3.5;
        # clear() was added to frame in Python 3.4
        # and is called by traceback.clear_frames(), which
        # in turn is called by unittest.TestCase.assertRaises
        pass

class Traceback(object):

    tb_next = None

    def __init__(self, tb):
        self.tb_frame = Frame(tb.tb_frame)
        # noinspection SpellCheckingInspection
        self.tb_lineno = int(tb.tb_lineno)

        # Build in place to avoid exceeding the recursion limit
        tb = tb.tb_next
        prev_traceback = self
        cls = type(self)
        while tb is not None:
            traceback = object.__new__(cls)
            traceback.tb_frame = Frame(tb.tb_frame)
            traceback.tb_lineno = int(tb.tb_lineno)
            prev_traceback.tb_next = traceback
            prev_traceback = traceback
            tb = tb.tb_next

    def as_traceback(self):
        if tproxy:
            return tproxy(TracebackType, self.__tproxy_handler)
        if not tb_set_next:
            raise RuntimeError("Unsupported Python interpreter!")

        current = self
        top_tb = None
        tb = None
        while current:
            f_code = current.tb_frame.f_code
            code = compile('\n' * (current.tb_lineno - 1) + 'raise __traceback_maker', current.tb_frame.f_code.co_filename, 'exec')
            if PY3:
                code = CodeType(
                    0, code.co_kwonlyargcount,
                    code.co_nlocals, code.co_stacksize, code.co_flags,
                    code.co_code, code.co_consts, code.co_names, code.co_varnames,
                    f_code.co_filename, f_code.co_name,
                    code.co_firstlineno, code.co_lnotab, (), ()
                )
            else:
                code = CodeType(
                    0,
                    code.co_nlocals, code.co_stacksize, code.co_flags,
                    code.co_code, code.co_consts, code.co_names, code.co_varnames,
                    f_code.co_filename.encode(), f_code.co_name.encode(),
                    code.co_firstlineno, code.co_lnotab, (), ()
                )

            # noinspection PyBroadException
            try:
                exec(code, current.tb_frame.f_globals, {})
            except Exception:
                next_tb = sys.exc_info()[2].tb_next
                if top_tb is None:
                    top_tb = next_tb
                if tb is not None:
                    tb_set_next(tb, next_tb)
                tb = next_tb
                del next_tb

            current = current.tb_next
        try:
            return top_tb
        finally:
            del top_tb
            del tb

    # noinspection SpellCheckingInspection
    def __tproxy_handler(self, operation, *args, **kwargs):
        if operation in ('__getattribute__', '__getattr__'):
            if args[0] == 'tb_next':
                return self.tb_next and self.tb_next.as_traceback()
            else:
                return getattr(self, args[0])
        else:
            return getattr(self, operation)(*args, **kwargs)

    def to_dict(self):
        """Convert a Traceback into a dictionary representation"""
        if self.tb_next is None:
            tb_next = None
        else:
            tb_next = self.tb_next.to_dict()

        code = {
            'co_filename': self.tb_frame.f_code.co_filename,
            'co_name': self.tb_frame.f_code.co_name,
        }
        frame = {
            'f_globals': self.tb_frame.f_globals,
            'f_code': code,
        }
        return {
            'tb_frame': frame,
            'tb_lineno': self.tb_lineno,
            'tb_next': tb_next,
        }

    @classmethod
    def from_dict(cls, dct):
        if dct['tb_next']:
            tb_next = cls.from_dict(dct['tb_next'])
        else:
            tb_next = None

        code = _AttrDict(
            co_filename=dct['tb_frame']['f_code']['co_filename'],
            co_name=dct['tb_frame']['f_code']['co_name'],
        )
        frame = _AttrDict(
            f_globals=dct['tb_frame']['f_globals'],
            f_code=code,
        )
        tb = _AttrDict(
            tb_frame=frame,
            tb_lineno=dct['tb_lineno'],
            tb_next=tb_next,
        )
        return cls(tb)

    @classmethod
    def from_string(cls, string, strict=True):
        frames = []
        header = strict

        for line in string.splitlines():
            line = line.rstrip()
            if header:
                if line == 'Traceback (most recent call last):':
                    header = False
                continue
            frame_match = FRAME_RE.match(line)
            if frame_match:
                frames.append(frame_match.groupdict())
            elif line.startswith('  '):
                pass
            elif strict:
                break  # traceback ended

        if frames:
            previous = None
            for frame in reversed(frames):
                previous = _AttrDict(
                    frame,
                    tb_frame=_AttrDict(
                        frame,
                        f_globals=_AttrDict(
                            __file__=frame['co_filename'],
                            __name__='?',
                        ),
                        f_code=_AttrDict(frame),
                    ),
                    tb_next=previous,
                )
            return cls(previous)
        else:
            raise TracebackParseError("Could not find any frames in %r." % string)


#!/usr/bin/python

"""
Known limitations:
    * only one $ORIGIN and one $TTL are supported
    * only the IN class is supported
    * PTR records must have a non-empty name
    * currently only supports the following:
    '$ORIGIN', '$TTL', 'SOA', 'NS', 'A', 'AAAA', 'CNAME', 'MX', 'PTR',
    'TXT', 'SRV', 'SPF', 'URI'
"""

import os
import copy
import datetime
import time
import argparse
from collections import defaultdict

from .configs import SUPPORTED_RECORDS, DEFAULT_TEMPLATE
from .exceptions import InvalidLineException

class ZonefileLineParser(argparse.ArgumentParser):
    def error(self, message):
        """
        Silent error message
        """
        raise InvalidLineException(message)

def make_rr_subparser(subparsers, rec_type, args_and_types):
    """
    Make a subparser for a given type of DNS record
    """
    sp = subparsers.add_parser(rec_type)

    sp.add_argument("name", type=str)
    sp.add_argument("ttl", type=int, nargs='?')
    sp.add_argument(rec_type, type=str)

    for my_spec in args_and_types:
        (argname, argtype) = my_spec[:2]
        if len(my_spec) > 2:
            nargs = my_spec[2]
            sp.add_argument(argname, type=argtype, nargs=nargs)
        else:
            sp.add_argument(argname, type=argtype)
    return sp

def make_txt_subparser(subparsers):
    sp = subparsers.add_parser("TXT")

    sp.add_argument("name", type=str)
    sp.add_argument("--ttl", type=int)
    sp.add_argument("TXT", type=str)
    sp.add_argument("txt", type=str, nargs='+')
    return sp

def make_parser():
    """
    Make an ArgumentParser that accepts DNS RRs
    """
    line_parser = ZonefileLineParser()
    subparsers = line_parser.add_subparsers()

    # parse $ORIGIN
    sp = subparsers.add_parser("$ORIGIN")
    sp.add_argument("$ORIGIN", type=str)

    # parse $TTL
    sp = subparsers.add_parser("$TTL")
    sp.add_argument("$TTL", type=int)

    # parse each RR
    args_and_types = [
        ("mname", str), ("rname", str), ("serial", int), ("refresh", int),
        ("retry", int), ("expire", int), ("minimum", int)
    ]
    make_rr_subparser(subparsers, "SOA", args_and_types)

    make_rr_subparser(subparsers, "NS", [("host", str)])
    make_rr_subparser(subparsers, "A", [("ip", str)])
    make_rr_subparser(subparsers, "AAAA", [("ip", str)])
    make_rr_subparser(subparsers, "CNAME", [("alias", str)])
    make_rr_subparser(subparsers, "ALIAS", [("host", str)])
    make_rr_subparser(subparsers, "MX", [("preference", str), ("host", str)])
    make_txt_subparser(subparsers)
    make_rr_subparser(subparsers, "PTR", [("host", str)])
    make_rr_subparser(subparsers, "SRV", [("priority", int), ("weight", int), ("port", int), ("target", str)])
    make_rr_subparser(subparsers, "SPF", [("data", str)])
    make_rr_subparser(subparsers, "URI", [("priority", int), ("weight", int), ("target", str)])

    return line_parser

def tokenize_line(line):
    """
    Tokenize a line:
    * split tokens on whitespace
    * treat quoted strings as a single token
    * drop comments
    * handle escaped spaces and comment delimiters
    """
    ret = []
    escape = False
    quote = False
    tokbuf = ""
    ll = list(line)
    while len(ll) > 0:
        c = ll.pop(0)
        if c.isspace():
            if not quote and not escape:
                # end of token
                if len(tokbuf) > 0:
                    ret.append(tokbuf)

                tokbuf = ""
            elif quote:
                # in quotes
                tokbuf += c
            elif escape:
                # escaped space
                tokbuf += c
                escape = False
            else:
                tokbuf = ""

            continue

        if c == '\\':
            escape = True
            continue
        elif c == '"':
            if not escape:
                if quote:
                    # end of quote
                    ret.append(tokbuf)
                    tokbuf = ""
                    quote = False
                    continue
                else:
                    # beginning of quote
                    quote = True
                    continue
        elif c == ';':
            if not escape:
                # comment
                ret.append(tokbuf)
                tokbuf = ""
                break

        # normal character
        tokbuf += c
        escape = False

    if len(tokbuf.strip(" ").strip("\n")) > 0:
        ret.append(tokbuf)

    return ret

def serialize(tokens):
    """
    Serialize tokens:
    * quote whitespace-containing tokens
    * escape semicolons
    """
    ret = []
    for tok in tokens:
        if " " in tok:
            tok = '"%s"' % tok

        if ";" in tok:
            tok = tok.replace(";", "\;")

        ret.append(tok)

    return " ".join(ret)

def remove_comments(text):
    """
    Remove comments from a zonefile
    """
    ret = []
    lines = text.split("\n")
    for line in lines:
        if len(line) == 0:
            continue

        line = serialize(tokenize_line(line))
        ret.append(line)

    return "\n".join(ret)

def flatten(text):
    """
    Flatten the text:
    * make sure each record is on one line.
    * remove parenthesis
    """
    lines = text.split("\n")

    # tokens: sequence of non-whitespace separated by '' where a newline was
    tokens = []
    for l in lines:
        if len(l) == 0:
            continue

        l = l.replace("\t", " ")
        tokens += filter(lambda x: len(x) > 0, l.split(" ")) + ['']

    # find (...) and turn it into a single line ("capture" it)
    capturing = False
    captured = []

    flattened = []
    while len(tokens) > 0:
        tok = tokens.pop(0)
        if not capturing and len(tok) == 0:
            # normal end-of-line
            if len(captured) > 0:
                flattened.append(" ".join(captured))
                captured = []
            continue

        if tok.startswith("("):
            # begin grouping
            tok = tok.lstrip("(")
            capturing = True

        if capturing and tok.endswith(")"):
            # end grouping.  next end-of-line will turn this sequence into a flat line
            tok = tok.rstrip(")")
            capturing = False

        captured.append(tok)

    return "\n".join(flattened)

def remove_class(text):
    """
    Remove the CLASS from each DNS record, if present.
    The only class that gets used today (for all intents
    and purposes) is 'IN'.
    """

    # see RFC 1035 for list of classes
    lines = text.split("\n")
    ret = []
    for line in lines:
        tokens = tokenize_line(line)
        tokens_upper = [t.upper() for t in tokens]

        if "IN" in tokens_upper:
            tokens.remove("IN")
        elif "CS" in tokens_upper:
            tokens.remove("CS")
        elif "CH" in tokens_upper:
            tokens.remove("CH")
        elif "HS" in tokens_upper:
            tokens.remove("HS")

        ret.append(serialize(tokens))

    return "\n".join(ret)

def add_default_name(text):
    """
    Go through each line of the text and ensure that
    a name is defined.  Use '@' if there is none.
    """
    global SUPPORTED_RECORDS

    lines = text.split("\n")
    ret = []
    for line in lines:
        tokens = tokenize_line(line)
        if len(tokens) == 0:
            continue

        if tokens[0] in SUPPORTED_RECORDS and not tokens[0].startswith("$"):
            # add back the name
            tokens = ['@'] + tokens

        ret.append(serialize(tokens))

    return "\n".join(ret)

def parse_line(parser, record_token, parsed_records):
    """
    Given the parser, capitalized list of a line's tokens, and the current set of records
    parsed so far, parse it into a dictionary.

    Return the new set of parsed records.
    Raise an exception on error.
    """

    global SUPPORTED_RECORDS

    line = " ".join(record_token)

    # match parser to record type
    if len(record_token) >= 2 and record_token[1] in SUPPORTED_RECORDS:
        # with no ttl
        record_token = [record_token[1]] + record_token
    elif len(record_token) >= 3 and record_token[2] in SUPPORTED_RECORDS:
        # with ttl
        record_token = [record_token[2]] + record_token
        if record_token[0] == "TXT":
            record_token = record_token[:2] + ["--ttl"] + record_token[2:]
    try:
        rr, unmatched = parser.parse_known_args(record_token)
        assert len(unmatched) == 0, "Unmatched fields: %s" % unmatched
    except (SystemExit, AssertionError, InvalidLineException):
        # invalid argument
        raise InvalidLineException(line)

    record_dict = rr.__dict__
    if record_token[0] == "TXT" and len(record_dict['txt']) == 1:
        record_dict['txt'] = record_dict['txt'][0]

    # what kind of record? including origin and ttl
    record_type = None
    for key in record_dict.keys():
        if key in SUPPORTED_RECORDS and (key.startswith("$") or record_dict[key] == key):
            record_type = key
            if record_dict[key] == key:
                del record_dict[key]
            break

    assert record_type is not None, "Unknown record type in %s" % rr

    # clean fields
    for field in record_dict.keys():
        if record_dict[field] is None:
            del record_dict[field]

    current_origin = record_dict.get('$ORIGIN', parsed_records.get('$ORIGIN', None))

    # special record-specific fix-ups
    if record_type == 'PTR':
        record_dict['fullname'] = record_dict['name'] + '.' + current_origin

    if len(record_dict) > 0:
        if record_type.startswith("$"):
            # put the value directly
            record_dict_key = record_type.lower()
            parsed_records[record_dict_key] = record_dict[record_type]
        else:
            record_dict_key = record_type.lower()
            parsed_records[record_dict_key].append(record_dict)

    return parsed_records

def parse_lines(text, ignore_invalid=False):
    """
    Parse a zonefile into a dict.
    @text must be flattened--each record must be on one line.
    Also, all comments must be removed.
    """
    json_zone_file = defaultdict(list)
    record_lines = text.split("\n")
    parser = make_parser()

    for record_line in record_lines:
        record_token = tokenize_line(record_line)
        try:
            json_zone_file = parse_line(parser, record_token, json_zone_file)
        except InvalidLineException:
            if ignore_invalid:
                continue
            else:
                raise

    return json_zone_file

def parse_zone_file(text, ignore_invalid=False):
    """
    Parse a zonefile into a dict
    """
    text = remove_comments(text)
    text = flatten(text)
    text = remove_class(text)
    text = add_default_name(text)
    json_zone_file = parse_lines(text, ignore_invalid=ignore_invalid)
    return json_zone_file


from .record_processors import (
    process_origin, process_ttl, process_soa, process_ns, process_a,
    process_aaaa, process_cname, process_alias, process_mx, process_ptr,
    process_txt, process_srv, process_spf, process_uri
)
from .configs import DEFAULT_TEMPLATE
import copy

def make_zone_file(json_zone_file_input, origin=None, ttl=None, template=None):
    """
    Generate the DNS zonefile, given a json-encoded description of the
    zone file (@json_zone_file) and the template to fill in (@template)

    json_zone_file = {
        "$origin": origin server,
        "$ttl":    default time-to-live,
        "soa":     [ soa records ],
        "ns":      [ ns records ],
        "a":       [ a records ],
        "aaaa":    [ aaaa records ]
        "cname":   [ cname records ]
        "alias":   [ alias records ]
        "mx":      [ mx records ]
        "ptr":     [ ptr records ]
        "txt":     [ txt records ]
        "srv":     [ srv records ]
        "spf":     [ spf records ]
        "uri":     [ uri records ]
    }
    """

    if template is None:
        template = DEFAULT_TEMPLATE[:]

    # careful...
    json_zone_file = copy.deepcopy(json_zone_file_input)
    if origin is not None:
        json_zone_file['$origin'] = origin

    if ttl is not None:
        json_zone_file['$ttl'] = ttl

    soa_records = [json_zone_file.get('soa')] if json_zone_file.get('soa') else None

    zone_file = template
    zone_file = process_origin(json_zone_file.get('$origin', None), zone_file)
    zone_file = process_ttl(json_zone_file.get('$ttl', None), zone_file)
    zone_file = process_soa(soa_records, zone_file)
    zone_file = process_ns(json_zone_file.get('ns', None), zone_file)
    zone_file = process_a(json_zone_file.get('a', None), zone_file)
    zone_file = process_aaaa(json_zone_file.get('aaaa', None), zone_file)
    zone_file = process_cname(json_zone_file.get('cname', None), zone_file)
    zone_file = process_alias(json_zone_file.get('alias', None), zone_file)
    zone_file = process_mx(json_zone_file.get('mx', None), zone_file)
    zone_file = process_ptr(json_zone_file.get('ptr', None), zone_file)
    zone_file = process_txt(json_zone_file.get('txt', None), zone_file)
    zone_file = process_srv(json_zone_file.get('srv', None), zone_file)
    zone_file = process_spf(json_zone_file.get('spf', None), zone_file)
    zone_file = process_uri(json_zone_file.get('uri', None), zone_file)

    # remove newlines, but terminate with one
    zone_file = "\n".join(
        filter(
            lambda l: len(l.strip()) > 0, [tl.strip() for tl in zone_file.split("\n")]
        )
    ) + "\n"

    return zone_file


import copy

def process_origin(data, template):
    """
    Replace {$origin} in template with a serialized $ORIGIN record
    """
    record = ""
    if data is not None:
        record += "$ORIGIN %s" % data

    return template.replace("{$origin}", record)

def process_ttl(data, template):
    """
    Replace {$ttl} in template with a serialized $TTL record
    """
    record = ""
    if data is not None:
        record += "$TTL %s" % data

    return template.replace("{$ttl}", record)

def process_soa(data, template):
    """
    Replace {SOA} in template with a set of serialized SOA records
    """
    record = template[:]

    if data is not None:

        assert len(data) == 1, "Only support one SOA RR at this time"
        data = data[0]

        soadat = []
        domain_fields = ['mname', 'rname']
        param_fields = ['serial', 'refresh', 'retry', 'expire', 'minimum']

        for f in domain_fields + param_fields:
            assert f in data.keys(), "Missing '%s' (%s)" % (f, data)

        data_name = str(data.get('name', '@'))
        soadat.append(data_name)

        if data.get('ttl') is not None:
            soadat.append( str(data['ttl']) )

        soadat.append("IN")
        soadat.append("SOA")

        for key in domain_fields:
            value = str(data[key])
            soadat.append(value)

        soadat.append("(")

        for key in param_fields:
            value = str(data[key])
            soadat.append(value)

        soadat.append(")")

        soa_txt = " ".join(soadat)
        record = record.replace("{soa}", soa_txt)

    else:
        # clear all SOA fields
        record = record.replace("{soa}", "")

    return record

def quote_field(data, field):
    """
    Quote a field in a list of DNS records.
    Return the new data records.
    """
    if data is None:
        return None

    data_dup = copy.deepcopy(data)
    for i in xrange(0, len(data_dup)):
        data_dup[i][field] = '"%s"' % data_dup[i][field]
        data_dup[i][field] = data_dup[i][field].replace(";", "\;")

    return data_dup

def process_rr(data, record_type, record_keys, field, template):
    """
    Meta method:
    Replace $field in template with the serialized $record_type records,
    using @record_key from each datum.
    """
    if data is None:
        return template.replace(field, "")

    if type(record_keys) == list:
        pass
    elif type(record_keys) == str:
        record_keys = [record_keys]
    else:
        raise ValueError("Invalid record keys")

    assert type(data) == list, "Data must be a list"

    record = ""
    for i in xrange(0, len(data)):

        for record_key in record_keys:
            assert record_key in data[i].keys(), "Missing '%s'" % record_key

        record_data = []
        record_data.append( str(data[i].get('name', '@')) )
        if data[i].get('ttl') is not None:
            record_data.append( str(data[i]['ttl']) )

        record_data.append(record_type)
        record_data += [str(data[i][record_key]) for record_key in record_keys]
        record += " ".join(record_data) + "\n"

    return template.replace(field, record)

def process_ns(data, template):
    """
    Replace {ns} in template with the serialized NS records
    """
    return process_rr(data, "NS", "host", "{ns}", template)

def process_a(data, template):
    """
    Replace {a} in template with the serialized A records
    """
    return process_rr(data, "A", "ip", "{a}", template)

def process_aaaa(data, template):
    """
    Replace {aaaa} in template with the serialized A records
    """
    return process_rr(data, "AAAA", "ip", "{aaaa}", template)

def process_cname(data, template):
    """
    Replace {cname} in template with the serialized CNAME records
    """
    return process_rr(data, "CNAME", "alias", "{cname}", template)

def process_alias(data, template):
    """
    Replace {alias} in template with the serialized ALIAS records
    """
    return process_rr(data, "ALIAS", "host", "{alias}", template)

def process_mx(data, template):
    """
    Replace {mx} in template with the serialized MX records
    """
    return process_rr(data, "MX", ["preference", "host"], "{mx}", template)

def process_ptr(data, template):
    """
    Replace {ptr} in template with the serialized PTR records
    """
    return process_rr(data, "PTR", "host", "{ptr}", template)

def process_txt(data, template):
    """
    Replace {txt} in template with the serialized TXT records
    """
    if data is None:
        to_process = None
    else:
        # quote txt
        to_process = copy.deepcopy(data)
        for datum in to_process:
            if isinstance(datum["txt"], list):
                datum["txt"] = " ".join(['"%s"' % entry.replace(";", "\;")
                                         for entry in datum["txt"]])
            else:
                datum["txt"] = '"%s"' % datum["txt"].replace(";", "\;")
    return process_rr(to_process, "TXT", "txt", "{txt}", template)

def process_srv(data, template):
    """
    Replace {srv} in template with the serialized SRV records
    """
    return process_rr(data, "SRV", ["priority", "weight", "port", "target"], "{srv}", template)

def process_spf(data, template):
    """
    Replace {spf} in template with the serialized SPF records
    """
    return process_rr(data, "SPF", "data", "{spf}", template)

def process_uri(data, template):
    """
    Replace {uri} in templtae with the serialized URI records
    """
    # quote target
    data_dup = quote_field(data, "target")
    return process_rr(data_dup, "URI", ["priority", "weight", "target"], "{uri}", template)


# Copyright (c) 2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
from __future__ import absolute_import

try:
    import simplejson as json
except ImportError:
    import json

import sys
import codecs
from functools import partial
import logging

import pyschema

from pyschema_extensions import avro  # import to get avro mixins for fields
from pyschema import source_generation
assert avro  # silence linter

logger = logging.getLogger(__name__)

SIMPLE_FIELD_MAP = {
    "string": pyschema.Text,
    "float": partial(pyschema.Float, size=4),
    "double": pyschema.Float,
    "int": partial(pyschema.Integer, size=4),
    "boolean": pyschema.Boolean,
    "long": pyschema.Integer,
    "bytes": pyschema.Bytes
}

class AVSCParseException(Exception):
    pass

def parse_schema_string(schema_string):
    """
    Load and return a PySchema class from an avsc string
    """
    if isinstance(schema_string, str):
        schema_string = schema_string.decode("utf8")
    schema_struct = json.loads(schema_string)

    return AvroSchemaParser().parse_schema_struct(schema_struct)

class AvroSchemaParser(object):

    def __init__(self):
        self.schema_store = pyschema.core.SchemaStore()

    def _parse_schema_or_enum_struct(self, struct):
        # experimental interface...
        field = self._get_field_builder(struct, None)()
        if isinstance(field, pyschema.SubRecord):
            return field._schema
        elif isinstance(field, pyschema.Enum):
            self.schema_store.add_enum(field)
            return field
        else:
            raise AVSCParseException(
                "pyschema's avro parser doesn't support base schemas other than record and enum"
            )

    def parse_schema_struct(self, schema_struct, enclosing_namespace=None):
        record_name = schema_struct["name"]
        field_dct = {}
        namespace = None

        if "namespace" in schema_struct:
            namespace = schema_struct["namespace"]
        else:
            namespace = enclosing_namespace

        if namespace is not None:
            field_dct["_namespace"] = namespace

        if "fields" not in schema_struct:
            logger.warning("No fields in schema {}, skipping.".format(schema_struct))
            return None  # raise Exception("no fields in {}".format(schema_struct))

        for field_def in schema_struct["fields"]:
            field_name = field_def["name"]
            field_builder = self._get_field_builder(
                field_def["type"],
                namespace
            )

            if "default" in field_def:
                if field_def["default"] is not None:
                    default_parser = field_builder(nullable=False).avro_load
                    default_value = default_parser(field_def["default"])
                else:
                    default_value = None
            else:
                default_value = pyschema.core.NO_DEFAULT

            field = field_builder(
                description=field_def.get("doc"),
                default=default_value
            )
            field_dct[field_name] = field

        field_dct["__module__"] = "__avro_parser_runtime__"  # not great, but better than "abc"

        if "doc" in schema_struct:
            field_dct["__doc__"] = schema_struct["doc"]

        schema = pyschema.core.PySchema(record_name.encode("ascii"), (pyschema.core.Record,), field_dct)
        self.schema_store.add_record(schema)
        return schema

    def _get_field_builder(self, type_def_struct, enclosing_namespace):
        if isinstance(type_def_struct, list):
            return self._parse_union(type_def_struct, enclosing_namespace)
        elif isinstance(type_def_struct, dict):
            return self._parse_complex(type_def_struct, enclosing_namespace)
        elif type_def_struct in SIMPLE_FIELD_MAP:
            field_builder = SIMPLE_FIELD_MAP.get(type_def_struct)
            return partial(field_builder, nullable=False)
        elif isinstance(type_def_struct, basestring):
            return self._parse_reference(type_def_struct, enclosing_namespace)
        else:
            raise AVSCParseException((
                "Unknown type, and no prior declaration of: {0!r}"
            ).format(type_def_struct))

    def _parse_union(self, union_struct, enclosing_namespace):
        filtered = [subtype for subtype in union_struct if subtype != "null"]
        if len(filtered) > 1:
            raise AVSCParseException(
                "Unions with multiple non-null as type are unsupported: {0}".format(union_struct)
            )
        nullable = "null" in union_struct
        actual_type = filtered[0]
        field_builder = self._get_field_builder(actual_type, enclosing_namespace)

        return partial(field_builder, nullable=nullable)

    def _parse_map(self, type_def_struct, enclosing_namespace):
        values_field_builder = self._get_field_builder(
            type_def_struct["values"],
            enclosing_namespace
        )
        return partial(
            pyschema.Map,
            # NOTE: values field builder is executed here, so the field index
            # will be created at this point, but that shouldn't matter for the map type
            value_type=values_field_builder(),
            nullable=False
        )

    def _parse_reference(self, type_def_struct, enclosing_namespace):
        reference_name = type_def_struct
        if "." not in reference_name:
            # try first with added enclosing namespace
            full_name = ".".join([enclosing_namespace, reference_name])
            if self.schema_store.has_enum(full_name):
                enum_field_values = self.schema_store.get_enum(full_name)
                return partial(
                    pyschema.Enum,
                    name=full_name,
                    values=enum_field_values
                )
        if self.schema_store.has_enum(reference_name):
            enum_field_values = self.schema_store.get_enum(reference_name)
            return partial(
                pyschema.Enum,
                name=reference_name,
                values=enum_field_values
            )
        elif self.schema_store.has_schema(reference_name):
            schema_class = self.schema_store.get(reference_name)
            return partial(
                pyschema.SubRecord,
                schema_class,
                nullable=False
            )
        else:
            raise AVSCParseException((
                "A schema type '{0!r}' was referenced"
                " without prior declaration as either a record or an enum."
            ).format(reference_name))

    def _parse_subrecord(self, type_def_struct, enclosing_namespace):
        if isinstance(type_def_struct, dict):
            if "fields" not in type_def_struct:
                raise AVSCParseException((
                    "No 'fields' definition found in subrecord"
                    " declaration: {0!r}"
                ).format(type_def_struct))
            schema_class = self.parse_schema_struct(type_def_struct, enclosing_namespace)

        return partial(
            pyschema.SubRecord,
            schema_class,
            nullable=False
        )

    def _parse_array(self, type_def_struct, enclosing_namespace):
        item_type = self._get_field_builder(type_def_struct["items"], enclosing_namespace)()
        return partial(
            pyschema.List,
            item_type,
            nullable=False
        )

    def _parse_enum(self, type_def_struct, enclosing_namespace):
        # copy and make a tuple just to ensure it isn't modified later
        values = tuple(type_def_struct["symbols"])
        if "namespace" in type_def_struct:
            name = ".".join([type_def_struct["namespace"], type_def_struct["name"]])
        else:
            name = type_def_struct["name"]

        builder = partial(
            pyschema.Enum,
            values,
            name=name,
            nullable=False
        )

        def build_and_add_to_enum_store(*args, **kwargs):
            field = builder(*args, **kwargs)
            return self.schema_store.add_enum(field)
        return build_and_add_to_enum_store

    COMPLEX_MAPPING = {
        "map": _parse_map,
        "record": _parse_subrecord,
        "array": _parse_array,
        "enum": _parse_enum
    }

    def _parse_complex(self, type_def_struct, enclosing_namespace):
        typename = type_def_struct["type"]
        parser_func = self.COMPLEX_MAPPING.get(typename)
        if parser_func:
            return parser_func(self, type_def_struct, enclosing_namespace)
        elif typename in SIMPLE_FIELD_MAP:
            # hack to support the weird case where you double wrap a simple type
            return SIMPLE_FIELD_MAP[typename]
        raise AVSCParseException("Unknown complex type: {0}".format(type_def_struct))

def to_python_source(s):
    """Return a Python syntax declaration of the schemas contained in `s`"""
    schema = parse_schema_string(s)
    return source_generation.to_python_source([schema])

if __name__ == "__main__":
    utf8_stdin = codecs.getreader("utf8")(sys.stdin)
    schema_struct = json.load(utf8_stdin)
    schema = AvroSchemaParser().parse_schema_struct(schema_struct)
    print source_generation.to_python_source([schema])


import copy
from pyschema import types
import pyschema
import os
import sys
from collections import defaultdict

DEFAULT_INDENT = " " * 4

class SourceGenerationError(Exception):
    pass

def to_python_source(classes, indent=DEFAULT_INDENT):
    """Convert a set of pyschemas to executable python source code

    Currently supports all built-in types for basic usage.

    Notably not supported:
    * Maintaining class hierarchy
    * Methods, properties and non-field attributes
    * SELF-references
    """
    return header_source() + "\n" + classes_source(classes, indent)

RESERVED_KEYWORDS = [
    "and", "del", "from", "not", "while", "as", "elif",
    "global", "or", "with", "assert", "else", "if",
    "pass", "yield", "break", "except", "import",
    "print", "class", "exec", "in", "raise", "continue",
    "finally", "is", "return", "def", "for", "lambda", "try"
]

def make_safe(package_name):
    parts = package_name.split(".")
    for kw in RESERVED_KEYWORDS:
        while kw in parts:
            i = parts.index(kw)
            parts[i] = kw + "_"
    return ".".join(parts)

class PackageBuilder(object):
    def __init__(self, target_folder, parent_package, indent=DEFAULT_INDENT):
        self.target_folder = target_folder
        self.parent_package = parent_package
        self.indent = indent

    def get_namespace(self, schema):
        try:
            namespace = make_safe(schema._namespace)
        except AttributeError:
            namespace = None
        return namespace

    def get_namespace_clusters(self, all_classes):
        namespace_cluster = defaultdict(set)
        for c in all_classes:
            namespace = self.get_namespace(c)
            namespace_cluster[namespace].add(c)
        return namespace_cluster

    def format_definitions(self, classes):
        return "\n\n".join([_class_source(c, self.indent) for c in classes])

    def write_namespace_file(self, namespace, module_code):
        if not namespace:
            key = ['__init__']
        else:
            key = namespace.split('.')
        output_file = os.path.join(self.target_folder, *key) + '.py'
        output_dir = os.path.dirname(output_file)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        with open(output_file, 'w') as out_fn:
            out_fn.write(module_code)

    def write_init_files(self):
        def touch_init_file_in_folder(folder):
            path = os.path.join(folder, '__init__.py')
            if not os.path.exists(path):
                open(path, 'w').close()

        touch_init_file_in_folder(self.target_folder)
        for root, dirs, _ in os.walk(self.target_folder):
            for d in dirs:
                touch_init_file_in_folder(os.path.join(root, d))

    def format_imports(self, imported_classes):
        if not imported_classes:
            return "\n"
        imported_namespaces = self.get_namespace_clusters(imported_classes)
        lines = []
        for namespace, schemas in imported_namespaces.iteritems():
            if self.parent_package:
                if namespace:
                    module = "{}.{}".format(self.parent_package, namespace)
                else:
                    module = self.parent_package
            else:
                if not namespace:
                    module = "."
                else:
                    module = namespace
            class_part = ", ".join(s.__name__ for s in schemas)
            lines.append("from {} import {}".format(module, class_part))
        return "\n".join(lines) + "\n\n"

    def _get_namespace_prefixes(self, namespaces):
        prefixes = []
        for n in namespaces:
            if n:
                parts = n.split(".")
                prefixes.append(".".join(parts[:-1]))
        return prefixes

    def from_classes_with_refs(self, classes):
        class_graph = CachedGraphTraverser()

        all_classes = set(classes)
        for c in classes:
            referenced_schemas = class_graph.find_descendants(c)
            all_classes |= set(referenced_schemas)

        namespace_cluster = self.get_namespace_clusters(all_classes)
        parent_namespaces = self._get_namespace_prefixes(namespace_cluster.keys())
        ordered_schemas = class_graph.get_reference_ordered_schemas(all_classes)

        # Since we don't want to use the previous cached results we create a new instance
        # This CachedGraphTraverser will only keep one-child depth for its find_descendants
        child_only_class_graph = CachedGraphTraverser()
        for namespace, classes in namespace_cluster.iteritems():
            inlined_classes = [c for c in ordered_schemas if c in classes]
            imported_classes = set()

            for inlined in inlined_classes:
                direct_references = child_only_class_graph.find_descendants(inlined, max_depth=1)
                imported_classes |= set([c for c in direct_references if c not in inlined_classes])

            module_code = (
                header_source() +
                self.format_imports(imported_classes) +
                self.format_definitions(inlined_classes)
            )
            if namespace not in parent_namespaces:
                filename = namespace
            else:
                filename = "{}.{}".format(namespace, "__init__")

            self.write_namespace_file(filename, module_code)
            self.write_init_files()

def to_python_package(classes, target_folder, parent_package=None, indent=DEFAULT_INDENT):
    '''
    This function can be used to build a python package representation of pyschema classes.
    One module is created per namespace in a package matching the namespace hierarchy.

    Args:
        classes: A collection of classes to build the package from
        target_folder: Root folder of the package
        parent_package: Prepended on all import statements in order to support absolute imports.
            parent_package is not used when building the package file structure
        indent: Indent level. Defaults to 4 spaces
    '''
    PackageBuilder(target_folder, parent_package, indent).from_classes_with_refs(classes)

def classes_source(classes, indent=DEFAULT_INDENT):
    all_classes = set(classes)
    class_graph = CachedGraphTraverser()
    for c in classes:
        referenced_schemas = class_graph.find_descendants(c)
        all_classes |= set(referenced_schemas)

    ordered = class_graph.get_reference_ordered_schemas(all_classes)
    return "\n\n".join([_class_source(c, indent) for c in ordered])

def header_source():
    """Get the required header for generated source"""
    return (
        "import pyschema\n"
        "from pyschema.types import *\n"
        "from pyschema.core import NO_DEFAULT\n"
    )

def _class_source(schema, indent):
    """Generate Python source code for one specific class

    Doesn't include or take into account any dependencies between record types
    """

    def_pattern = (
        "class {class_name}(pyschema.Record):\n"
        "{indent}# WARNING: This class was generated by pyschema.to_python_source\n"
        "{indent}# there is a risk that any modification made to this class will be overwritten\n"
        "{optional_namespace_def}"
        "{field_defs}\n"
    )
    if hasattr(schema, '_namespace'):
        optional_namespace_def = "{indent}_namespace = {namespace!r}\n".format(
            namespace=schema._namespace, indent=indent)
    else:
        optional_namespace_def = ""

    field_defs = [
        "{indent}{field_name} = {field!r}".format(field_name=field_name, field=field, indent=indent)
        for field_name, field in schema._fields.iteritems()
    ]
    if not field_defs:
        field_defs = ["{indent}pass".format(indent=indent)]

    return def_pattern.format(
        class_name=schema._schema_name,
        optional_namespace_def=optional_namespace_def,
        field_defs="\n".join(field_defs),
        indent=indent
    )

class CachedGraphTraverser(object):
    def __init__(self):
        self.descendants = {}
        self.started = set()

    def find_descendants(self, a, max_depth=sys.getrecursionlimit()):
        if a in self.descendants:
            # fetch from cache
            return self.descendants[a]
        self.started.add(a)
        subs = set()
        if max_depth > 0:
            if pyschema.ispyschema(a):
                for _, field in a._fields.iteritems():
                    subs |= self.find_descendants(field, max_depth)
                self.descendants[a] = subs
            elif isinstance(a, types.List):
                subs |= self.find_descendants(a.field_type, max_depth)
            elif isinstance(a, types.Map):
                subs |= self.find_descendants(a.value_type, max_depth)
            elif isinstance(a, types.SubRecord):
                subs.add(a._schema)
                if a._schema not in self.started:  # otherwise there is a circular reference
                    subs |= self.find_descendants(a._schema, max_depth-1)
            self.started.remove(a)
        return subs

    def get_reference_ordered_schemas(self, schema_set):
        for schema in schema_set:
            self.find_descendants(schema)
        descendants = copy.deepcopy(self.descendants)  # a working copy

        ordered_output = []
        while descendants:
            leaves = []
            for root, referenced in descendants.iteritems():
                if len(referenced) == 0:
                    leaves.append(root)

            if not leaves:
                raise SourceGenerationError("Circular reference in input schemas, aborting")
            ordered_output += leaves
            for leaf in leaves:
                # remove all leaves
                descendants.pop(leaf)
                for root, referenced in descendants.iteritems():
                    if leaf in referenced:
                        referenced.remove(leaf)
        return ordered_output


# Copyright (c) 2013 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

""" Schema definition toolkit using Python classes

Usage example:

>>> class Foo(Record):
... bin = Bytes()
...
... class MyRecord(Record):
...     a_string = Text()
...     a_float = Float()
...     record = List(SubRecord(Foo))
...
... rec = MyRecord(a_string="hej")
... rec.record = [Foo(bin="bar")]
...
... s = dumps(rec)
... print loads(s)

Internals:

A valid PySchema class contains the following class variables:

`_fields`
    An OrderedDict of `field_name` => `field_type`
    where `field_type` is an instance of a Field subclass

`_schema_name`
    The qualifying name for this schema. This is used for registering a record
    in a `SchemaStore` and for auto-identification of serialized records.
    Should be unique within a specific SchemaStore, so if auto registering is
    used it should be unique within the execution chain of the current program.

"""
from __future__ import absolute_import
from abc import ABCMeta, abstractmethod
from itertools import izip

try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

import warnings
import types   # absolute import, this is the python standard library types
try:
    import simplejson as json
except ImportError:
    import json

SCHEMA_FIELD_NAME = "$schema"

class ParseError(Exception):
    """ Generic exception type for Record parse errors """
    pass

class InvalidSchemaSpecification(object):
    """
    Utility class that can be used to raise an exception on schema usage.

    This is used in the schema store as a placeholder for invalid schemas. Instead of raising when a schema is
    registered in the store, something that happens on import, we use this class to raise on usage.
    """
    def __init__(self, exception_msg):
        self.exception_msg = exception_msg

    def __getattr__(self, item):
        raise ValueError(self.exception_msg)

class SchemaStore(object):
    def __init__(self):
        self._schema_map = {}
        self._enum_map = {}

    def __str__(self):
        return str(self._schema_map.keys())

    def add_record(self, schema, _bump_stack_level=False):
        """ Add record class to record store for retrieval at record load time.

            Can be used as a class decorator
        """
        full_name = get_full_name(schema)
        has_namespace = '.' in full_name
        self._force_add(full_name, schema, _bump_stack_level, _raise_on_existing=has_namespace)
        if has_namespace and schema.__name__ not in self._schema_map:
            self._force_add(schema.__name__, schema, _bump_stack_level)
        return schema

    def add_enum(self, enum_definition):
        new_values_set = set(enum_definition.values)
        old_values_set = self._enum_map.get(enum_definition.name)

        if old_values_set is not None and new_values_set != old_values_set:
            warnings.warn(
                "Enum {!r} overwritten! Was: {}, Overwritten by: {}".format(
                    enum_definition.name,
                    old_values_set,
                    new_values_set
                )
            )

        if enum_definition.name is not None:
            self._enum_map[enum_definition.name] = enum_definition.values
        # return the definition to allow the method to be used as a decorator
        return enum_definition

    def _force_add(self, used_name, schema, _bump_stack_level=False, _raise_on_existing=False):
        existing = self._schema_map.get(used_name, None)
        if existing and existing != schema:
            full_name = get_full_name(schema)
            explanation = "(actually {0})".format() if full_name != used_name else ""

            warnings.warn(
                "{used_name}{explanation}: old definition in {prev_module} replaced by definition in {new_module}"
                .format(used_name=used_name,
                        explanation=explanation,
                        prev_module=existing.__module__,
                        new_module=schema.__module__),
                stacklevel=4 if _bump_stack_level else 3)

            if _raise_on_existing:
                if not isinstance(existing, InvalidSchemaSpecification):
                    schema = InvalidSchemaSpecification(
                        'Attempted to access data from a dubious schema specification. '
                        'The schema for: {used_name} was provided by both {existing} and {new}'
                        .format(used_name=used_name, existing=existing, new=schema))
                else:
                    schema = existing

        self._schema_map[used_name] = schema

    def get(self, record_name):
        """
        Will return a matching record or raise KeyError is no record is found.

        If the record name is a full name we will first check for a record matching the full name.
        If no such record is found any record matching the last part of the full name (without the namespace) will
        be returned.
        """
        if record_name in self._schema_map:
            return self._schema_map[record_name]
        else:
            last_name = record_name.split('.')[-1]
            return self._schema_map[last_name]

    def get_enum(self, name):
        return self._enum_map[name]

    def clear(self):
        self._schema_map.clear()
        self._enum_map.clear()

    def clone(self):
        r = SchemaStore()
        r._schema_map = self._schema_map.copy()
        r._enum_map = self._enum_map.copy()

        return r

    def has_schema(self, name):
        if name in self._schema_map:
            return True
        if "." in name:
            basename = name.split(".")[-1]
            return basename in self._schema_map
        return False

    def has_enum(self, name):
        return name in self._enum_map

    def __contains__(self, schema):
        return schema in self._schema_map.values()

class RecordStore(SchemaStore):
    def __init__(self):
        warnings.warn("RecordStore is deprecated and has been renamed to SchemaStore", DeprecationWarning, stacklevel=2)
        super(RecordStore, self).__init__()

def get_full_name(schema):
    full_name = schema.__name__
    if hasattr(schema, '_namespace'):
        full_name = '.'.join([schema._namespace, schema.__name__])
    elif hasattr(schema, '_avro_namespace_'):
        warnings.warn("_avro_namespace is deprecated, use _namespace instead", DeprecationWarning, stacklevel=3)
        full_name = '.'.join([schema._avro_namespace_, schema.__name__])
    return full_name

class _NoDefault:
    def __repr__(self):
        return "NO_DEFAULT"

# NO_DEFAULT is a special value to signify that a field has no default value
# and should fail to serialize unless a value has been assigned
# it's the default default-value for all non-nullable fields

NO_DEFAULT = _NoDefault()

_UNTOUCHED = object()

class Field(object):
    __metaclass__ = ABCMeta
    _next_index = 0

    def __init__(self, description=None, nullable=True, default=_UNTOUCHED):
        self.description = description
        self._index = Field._next_index
        self.nullable = nullable
        if default is _UNTOUCHED:
            # if default isn't explicitly set
            # use None for Nullables, and NO_DEFAULT for others
            if nullable:
                default = None
            else:
                default = NO_DEFAULT

        self.default = default
        Field._next_index += 1  # used for arg order in initialization

    def repr_vars(self):
        """Return a dictionary the field definition

        Should contain all fields that are required for the definition of this field in a pyschema class"""
        d = OrderedDict()
        d["nullable"] = repr(self.nullable)
        d["default"] = repr(self.default)
        if self.description is not None:
            d["description"] = repr(self.description)
        return d

    def __repr__(self):
        strings = ('{0}={1}'.format(vname, val) for vname, val in self.repr_vars().iteritems())
        return self.__class__.__name__ + '(' + ', '.join(strings) + ')'

    def set_parent(self, schema):
        # no-op by default but can be overridden by types
        # that need parent references
        pass

    @abstractmethod
    def dump(self, obj):
        pass

    @abstractmethod
    def load(self, obj):
        pass

    @classmethod
    def mixin(cls, mixin_cls):
        """Decorator for mixing in additional functionality into field type

        Example:

        >>> @Integer.mixin
        ... class IntegerPostgresExtensions:
        ...     postgres_type = 'INT'
        ...
        ...     def postgres_dump(self, obj):
        ...         self.dump(obj) + "::integer"

        Is roughly equivalent to:

        >>> Integer.postgres_type = 'INT'
        ...
        ... def postgres_dump(self, obj):
        ...     self.dump(obj) + "::integer"
        ...
        ... Integer.postgres_dump = postgres_dump

        """
        for item_name in dir(mixin_cls):
            if item_name.startswith("__"):
                # don't copy magic properties
                continue
            item = getattr(mixin_cls, item_name)

            if isinstance(item, types.MethodType):
                # unbound method will cause problems
                # so get the underlying function instead
                item = item.im_func

            setattr(cls, item_name, item)
        return mixin_cls

    def default_value(self):
        return self.default

    def is_similar_to(self, other):
        return(
            type(self) == type(other) and
            self.default == other.default and
            self.nullable == other.nullable and
            self.description == other.description
        )

auto_store = SchemaStore()

class PySchema(ABCMeta):
    """Metaclass for Records

    Builds schema on Record declaration and remembers Record types
    for easy generic parsing
    """
    auto_register = True

    def __new__(metacls, name, bases, dct):
        schema_attrs = metacls._get_schema_attributes(
            name=name,
            bases=bases,
            dct=dct
        )
        dct.update(schema_attrs)
        cls = ABCMeta.__new__(metacls, name, bases, dct)

        # allow self-references etc.
        for field_name, field in cls._fields.iteritems():
            field.set_parent(cls)

        if metacls.auto_register:
            auto_store.add_record(cls, _bump_stack_level=True)
        return cls

    @classmethod
    def _field_dupe_warning(metacls, name, fields):
        warnings.warn(
            "{schema}: Duplicate field definition for field{plural} {field}"
                .format(
                    schema=name,
                    field=fields,
                    plural="s" if len(fields) > 1 else ""
                ),
            stacklevel=4
        )

    @classmethod
    def _get_schema_attributes(metacls, name, bases, dct):
        fields = OrderedDict()
        for b in bases:
            if not isinstance(b, metacls):
                continue

            field_intersection = set(fields) & set(b._fields)
            if field_intersection:
                metacls._field_dupe_warning(name, field_intersection)
            fields.update(b._fields)

        new_fields = []
        for field_name, field_def in dct.iteritems():
            if isinstance(field_def, Field):
                new_fields.append((field_name, field_def))

        new_fields.sort(key=lambda fd: fd[1]._index)
        for field_name, field_def in new_fields:
            if field_name in fields:
                metacls._field_dupe_warning(name, (field_name,))
            fields[field_name] = field_def

        return {
            "_fields": fields,
            "_schema_name": name,
        }

    @classmethod
    def from_class(metacls, cls, auto_store=True):
        """Create proper PySchema class from cls

        Any methods and attributes will be transferred to the
        new object
        """
        if auto_store:
            def wrap(cls):
                return cls
        else:
            wrap = no_auto_store()

        return wrap(metacls.__new__(
            metacls,
            cls.__name__,
            (Record,),
            dict(cls.__dict__)
        ))

def disable_auto_register():
    PySchema.auto_register = False

def enable_auto_register():
    PySchema.auto_register = True

def no_auto_store():
    """ Temporarily disable automatic registration of records in the auto_store

    Decorator factory. This is _NOT_ thread safe

    >>> @no_auto_store()
    ... class BarRecord(Record):
    ...     pass
    >>> BarRecord in auto_store
    False

    """
    original_auto_register_value = PySchema.auto_register
    disable_auto_register()

    def decorator(cls):
        PySchema.auto_register = original_auto_register_value
        return cls

    return decorator

@no_auto_store()
class Record(object):
    """Abstract base class for structured logging records"""
    __metaclass__ = PySchema

    def __init__(self, *args, **kwargs):
        if args:
            # The idea behind only allowing keyword arguments
            # is to prevent accidental misuse of a changed schema
            raise TypeError('Non-keyword arguments not allowed'
                            ' when initializing Records')

        for k, field_type in self._fields.items():
            if k in kwargs:
                value = kwargs.get(k)
            else:
                value = field_type.default_value()

            object.__setattr__(self, k, value)

    def __setattr__(self, name, value):
        if name not in self._fields:
            raise AttributeError(
                "No field %r in %s"
                % (name, self._schema_name)
            )

        super(Record, self).__setattr__(name, value)

    def __unicode__(self):
        return str(self).decode('ascii')

    def __str__(self):
        return repr(self)

    def __repr__(self):
        strings = ('%s=%r' % (fname, getattr(self, fname))
                   for fname, f in self._fields.iteritems())

        return self._schema_name + '(' + ', '.join(strings) + ')'

    def __cmp__(self, other):
        if not isinstance(other, Record):
            # return default implementation cmp value
            return cmp(id(self), other)
        if self._schema_name != other._schema_name:
            return cmp(self._schema_name, other._schema_name)
        fields = self._fields.keys()
        a = (getattr(self, key) for key in fields)
        b = (getattr(other, key) for key in fields)

        for _a, _b in izip(a, b):
            r = cmp(_a, _b)
            if r:
                return r
        return 0

    def __eq__(self, other):
        return self.__cmp__(other) == 0

    def __ne__(self, other):
        return self.__cmp__(other) != 0

def to_json_compatible(record):
    "Dump record in json-encodable object format"
    d = {}
    for fname, f in record._fields.iteritems():
        val = getattr(record, fname)
        if val is not None:
            d[fname] = f.dump(val)
    return d

def from_json_compatible(schema, dct):
    "Load from json-encodable"
    kwargs = {}

    for key in dct:
        field_type = schema._fields.get(key)
        if field_type is None:
            raise ParseError("Unexpected field encountered in line for record %s: %s" % (schema.__name__, key))
        kwargs[key] = field_type.load(dct[key])

    return schema(**kwargs)

def ispyschema(schema):
    """ Is object PySchema instance?

    Returns true for PySchema Record *classes*
    i.e. NOT when schema is a *Record* instance

    >>> class FooRecord(Record):
    ...     pass
    >>> ispyschema(FooRecord)
    True
    >>> ispyschema(FooRecord())
    False
    """
    return isinstance(schema, PySchema)

def load_json_dct(
        dct,
        record_store=None,
        schema=None,
        loader=from_json_compatible
):
    """ Create a Record instance from a json-compatible dictionary

    The dictionary values should have types that are json compatible,
    as if just loaded from a json serialized record string.

    :param dct:
        Python dictionary with key/value pairs for the record

    :param record_store:
        Record store to use for schema lookups (when $schema field is present)

    :param schema:
        PySchema Record class for the record to load.
        This will override any $schema fields specified in `dct`

    """
    if schema is None:
        if record_store is None:
            record_store = auto_store
        try:
            schema_name = dct.pop(SCHEMA_FIELD_NAME)
        except KeyError:
            raise ParseError((
                "Serialized record missing '{0}' "
                "record identifier and no schema supplied")
                .format(SCHEMA_FIELD_NAME)
            )
        try:
            schema = record_store.get(schema_name)
        except KeyError:
            raise ParseError(
                "Can't recognize record type %r"
                % (schema_name,), schema_name)

    # if schema is explicit, use that instead of SCHEMA_FIELD_NAME
    elif SCHEMA_FIELD_NAME in dct:
        dct.pop(SCHEMA_FIELD_NAME)

    record = loader(schema, dct)
    return record

def loads(
        s,
        record_store=None,
        schema=None,
        loader=from_json_compatible,
        record_class=None  # deprecated in favor of schema
):
    """ Create a Record instance from a json serialized dictionary

    :param s:
        String with a json-serialized dictionary

    :param record_store:
        Record store to use for schema lookups (when $schema field is present)

    :param loader:
        Function called to fetch attributes from json. Typically shouldn't be used by end users

    :param schema:
        PySchema Record class for the record to load.
        This will override any $schema fields specified in `s`

    :param record_class:
        DEPRECATED option, old name for the `schema` parameter

    """
    if record_class is not None:
        warnings.warn(
            "The record_class parameter is deprecated in favour of schema",
            DeprecationWarning,
            stacklevel=2
        )
        schema = record_class
    if not isinstance(s, unicode):
        s = s.decode('utf8')
    if s.startswith(u"{"):
        json_dct = json.loads(s)
        return load_json_dct(json_dct, record_store, schema, loader)
    else:
        raise ParseError("Not a json record")

def dumps(obj, attach_schema_name=True):
    json_dct = to_json_compatible(obj)
    if attach_schema_name:
        json_dct[SCHEMA_FIELD_NAME] = get_full_name(obj.__class__)

    json_string = json.dumps(json_dct)
    return json_string


# Copyright (c) 2014 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the 'License'); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
'''
Extension for generating JSON schema schemas from PySchema classes

JSON schema: http://json-schema.org/

When dumping to JSON schema, all fields in a record are mandatory, although
a list or map can be empty (but must be present). These records are still
dump-able, but they will not validate against the schema.

Usage:

>>> class MyRecord(pyschema.Record):
>>>     foo = Text()
>>>     bar = Integer()
>>>
>>> [pyschema_extensions.jsonschema.]get_root_schema_string(MyRecord)

'{"additionalProperties": false, "required": ["bar", "foo"], "type": "object", "id": "MyRecord", "properties": {"foo": {"t
ype": "string"}, "bar": {"type": "integer"}}}  '
'''

from pyschema import core
from pyschema.types import Field, Boolean, Integer, Float
from pyschema.types import Text, Enum, List, Map, SubRecord
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict
try:
    import simplejson as json
except ImportError:
    import json

# Bytes are not supported
Boolean.jsonschema_type_name = 'boolean'
Integer.jsonschema_type_name = 'integer'
Float.jsonschema_type_name = 'number'
Text.jsonschema_type_name = 'string'
Enum.jsonschema_type_name = 'string'
List.jsonschema_type_name = 'array'
Map.avro_type_name = 'object'

@Field.mixin
class FieldMixin:
    def jsonschema_type_schema(self, state):
        return {
            'type': self.jsonschema_type_name,
        }

@Enum.mixin
class EnumMixin:
    def jsonschema_type_schema(self, state):
        return {
            'type': self.jsonschema_type_name,
            'enum': sorted(list(self.values)),
        }

@List.mixin
class ListMixin:
    def jsonschema_type_schema(self, state):
        return {
            'type': self.jsonschema_type_name,
            'items': self.field_type.jsonschema_type_schema(state),
        }

@Map.mixin
class MapMixin:
    def jsonschema_type_schema(self, state):
        return {
            'type': 'object',
            'additionalProperties': True,
            'patternProperties': {
                '^.*$': self.value_type.jsonschema_type_schema(state),
            },
        }

@SubRecord.mixin
class SubRecordMixin:
    def jsonschema_type_schema(self, state):
        cls = self._schema
        state.record_schemas[cls._schema_name] = get_schema_dict(cls, state)
        return {
            '$ref': self.jsonschema_type_name,
        }

    @property
    def jsonschema_type_name(self):
        return '#/definitions/{0}'.format(self._schema._schema_name)

# Schema generation
class SchemaGeneratorState(object):
    def __init__(self):
        self.record_schemas = dict()

def get_schema_dict(record, state=None):
    """Return a python dict representing the jsonschema of a record

    Any references to sub-schemas will be URI fragments that won't be
    resolvable without a root schema, available from get_root_schema_dict.
    """
    state = state or SchemaGeneratorState()
    schema = OrderedDict([
        ('type', 'object'),
        ('id', record._schema_name),
    ])
    fields = dict()
    for field_name, field_type in record._fields.iteritems():
        fields[field_name] = field_type.jsonschema_type_schema(state)

    required = set(fields.keys())
    schema['properties'] = fields
    schema['required'] = sorted(list(required))
    schema['additionalProperties'] = False

    state.record_schemas[record._schema_name] = schema
    return schema

def get_root_schema_dict(record):
    """Return a root jsonschema for a given record

    A root schema includes the $schema attribute and all sub-record
    schemas and definitions.
    """
    state = SchemaGeneratorState()
    schema = get_schema_dict(record, state)
    del state.record_schemas[record._schema_name]
    if state.record_schemas:
        schema['definitions'] = dict()
        for name, sub_schema in state.record_schemas.iteritems():
            schema['definitions'][name] = sub_schema
    return schema

def get_root_schema_string(record):
    return json.dumps(get_root_schema_dict(record))

def dumps(record):
    return json.dumps(core.to_json_compatible(record))

def loads(s, record_store=None, schema=None):
    return core.loads(s, record_store, schema, core.from_json_compatible)


# Copyright (c) 2013 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""
Extension for generating Avro schemas from PySchema Record classes

Usage:

>>> class MyRecord(pyschema.Record):
>>>     foo = Text()
>>>     bar = Integer()
>>>
>>> [pyschema_extensions.avro.]get_schema_string(MyRecord)

'{"fields": [{"type": "string", "name": "foo"},
{"type": "long", "name": "bar"}],
"type": "record", "name": "MyRecord"}'

"""
import warnings
from pyschema import core
from pyschema.types import Field, Boolean, Integer, Float
from pyschema.types import Bytes, Text, Enum, List, Map, SubRecord
try:
    import simplejson as json
except ImportError:
    import json

Boolean.avro_type_name = "boolean"
Bytes.avro_type_name = "bytes"
Text.avro_type_name = "string"
# "ENUM" is the avro 'type name' of all enums generated by pyschema
# this is pyschema convention, not avro, so it might change if
# need b
Enum.avro_type_name = "ENUM"
List.avro_type_name = "array"
Map.avro_type_name = "map"

@Float.mixin
class FloatMixin:
    @property
    def avro_type_name(self):
        if self.size <= 4:
            return 'float'
        return 'double'

@Integer.mixin
class IntegerMixin:
    @property
    def avro_type_name(self):
        if self.size <= 4:
            return 'int'
        return 'long'

@Field.mixin
class FieldMixin:
    def avro_type_schema(self, state):
        """Full type specification for the field

        I.e. the same as would go into the "type" field.
        For most field, only simplified_avro_type_schema has
        to be implemented.
        """
        simple_type = self.simplified_avro_type_schema(state)
        if self.nullable:
            # first value in union needs to be same as default
            if self.default in (None, core.NO_DEFAULT):
                return ["null", simple_type]
            else:
                return [simple_type, "null"]
        else:
            return simple_type

    def simplified_avro_type_schema(self, state):
        """The basic avro type for this field

        Not including nullability.
        """
        return self.avro_type_name

    def avro_dump(self, o):
        if o is None:
            return None
        else:
            # relying on the reference json dump behavior
            # could be a bit dangerous
            if self.nullable:
                return {self.avro_type_name: self.dump(o)}
            else:
                return self.dump(o)

    def avro_load(self, o):
        if o is None:
            return None
        else:
            if self.nullable:
                return self.load(o[self.avro_type_name])
            else:
                return self.load(o)

    def avro_default_value(self):
        return self.default

@List.mixin
class ListMixin:
    def simplified_avro_type_schema(self, state):
        return {
            "type": "array",
            "items": self.field_type.avro_type_schema(state)
        }

    def avro_dump(self, obj):
        if obj is None:
            return None
        else:
            l = [self.field_type.avro_dump(o) for o in obj]
            if self.nullable:
                return {self.avro_type_name: l}
            else:
                return l

    def avro_load(self, obj):
        if obj is None:
            return None
        else:
            if self.nullable:
                obj = obj[self.avro_type_name]
            return [
                self.field_type.avro_load(o)
                for o in obj
            ]

@Enum.mixin
class EnumMixin:
    @property
    def avro_type_name(self):
        if self.name is None:
            return "ENUM"
        return self.name

    def simplified_avro_type_schema(self, state):
        return {
            "type": "enum",
            "name": self.avro_type_name,
            "symbols": list(self.values)
        }

@SubRecord.mixin
class SubRecordMixin:
    def simplified_avro_type_schema(self, state):
        return get_schema_dict(self._schema, state)

    @property
    def avro_type_name(self):
        return core.get_full_name(self._schema)

    def avro_dump(self, obj):
        if obj is None:
            return None
        if self.nullable:
            return {self.avro_type_name: to_json_compatible(obj)}
        else:
            return to_json_compatible(obj)

    def _get_record_data(self, obj):
        """
        Try to get the data for this subrecord using the full_name including namespace.
        If no data is given for the fullname, we will fallback to use the SubRecord name without namespace.

        Doing this allows namespace addition to a schema while being able to read legacy data missing a namespace.
        """
        if self.avro_type_name in obj:
            return obj[self.avro_type_name]
        else:
            base_name = self.avro_type_name.split('.')[-1]
            return obj[base_name]

    def avro_load(self, obj):
        if obj is None:
            return None
        if self.nullable:
            return from_json_compatible(
                self._schema,
                self._get_record_data(obj)
            )
        else:
            return from_json_compatible(
                self._schema,
                obj
            )

@Map.mixin
class MapMixin:
    def simplified_avro_type_schema(self, state):
        assert isinstance(self.key_type, Text)
        return {
            "type": "map",
            "values": self.value_type.avro_type_schema(state)
        }

    def avro_dump(self, obj):
        if obj is None:
            return None
        else:
            m = dict([(
                # using json loader for key is kind of a hack
                # since this isn't an actual type in avro (always text)
                self.key_type.dump(k),
                self.value_type.avro_dump(v)
            ) for k, v in obj.iteritems()])
            if self.nullable:
                return {self.avro_type_name: m}
            else:
                return m

    def avro_load(self, obj):
        if obj is None:
            return None
        else:
            if self.nullable:
                obj = obj[self.avro_type_name]
            m = dict([(
                # using json loader for key is kind of a hack
                # since this isn't an actual type in avro (always text)
                self.key_type.load(k),
                self.value_type.avro_load(v)
            ) for k, v in obj.iteritems()])
            return m

# Schema generation
class SchemaGeneratorState(object):
    def __init__(self):
        self.declared_records = set()

def get_schema_dict(record, state=None):
    state = state or SchemaGeneratorState()

    full_name = core.get_full_name(record)
    if full_name in state.declared_records:
        return full_name
    state.declared_records.add(full_name)

    if '.' in full_name:
        namespace, _ = full_name.rsplit('.', 1)
    else:
        namespace = None

    avro_record = {
        "type": "record",
        "name": record._schema_name,
    }
    if namespace:
        avro_record["namespace"] = namespace

    if record.__doc__ is not None:
        avro_record["doc"] = record.__doc__

    avro_fields = []
    for field_name, field_type in record._fields.iteritems():
        field_spec = {
            "name": field_name,
            "type": field_type.avro_type_schema(state)
        }
        if field_type.default is not core.NO_DEFAULT:
            field_spec["default"] = field_type.avro_default_value()
        if field_type.description is not None:
            field_spec["doc"] = field_type.description
        avro_fields.append(field_spec)

    avro_record["fields"] = avro_fields
    return avro_record

def get_schema_string(record):
    return json.dumps(get_schema_dict(record))

def dumps(record):
    return json.dumps(to_json_compatible(record))

def to_json_compatible(record):
    dct = {}
    for name, fieldtype in record._fields.iteritems():
        value = getattr(record, name)
        dct[name] = fieldtype.avro_dump(value)
    return dct

def from_json_compatible(schema, dct):
    "Load from json-encodable"
    kwargs = {}

    for key in dct:
        field_type = schema._fields.get(key)
        if field_type is None:
            warnings.warn("Unexpected field encountered in line for record %s: %r" % (schema.__name__, key))
            continue
        kwargs[key] = field_type.avro_load(dct[key])

    return schema(**kwargs)

def loads(
    s,
    record_store=None,
    schema=None,
    record_class=None  # deprecated - replaced by `schema`
):
    if record_class is not None:
        warnings.warn(
            "The record_class parameter is deprecated in favour of schema",
            DeprecationWarning,
            stacklevel=2
        )
        schema = record_class

    return core.loads(s, record_store, schema, from_json_compatible)


# Copyright (c) 2013 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

"""Basic utilities for using PySchema in Luigi's Python MR lib
"""
import sys
from pyschema import core

def mr_reader(job, input_stream, loads=core.loads):
    """ Converts a file object with json serialised pyschema records
        to a stream of pyschema objects

    Can be used as job.reader in luigi.hadoop.JobTask
    """
    for line in input_stream:
        yield loads(line),

def mr_writer(job, outputs, output_stream,
              stderr=sys.stderr, dumps=core.dumps):
    """ Writes a stream of json serialised pyschema Records to a file object

    Can be used as job.writer in luigi.hadoop.JobTask
    """
    for output in outputs:
        try:
            print >> output_stream, dumps(output)
        except core.ParseError, e:
            print >> stderr, e
            raise