#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import pyfits
import numpy as np
from os import path
from astropy.io import votable

import argparse

COLUMN_FORMATS_FITS = {
    'J': 'integer',
    'I': 'integer',
    'L': 'boolean',
    'E': 'real',
    'D': 'double precision',
    'A': 'character varying(%s)'
}

COLUMN_FORMATS_VOTABLE = {
    'double': 'double precision',
    'int': 'integer',
    'long': 'integer',
    'float': 'float',
    'char': 'character varying%s'
}

parser = argparse.ArgumentParser(description="""
A tool to convert fits/votable file into SQL database data.
""", formatter_class=argparse.RawDescriptionHelpFormatter)

parser.add_argument('-i', '--input', type=str, default=None,
                    help='FITS file to use')

parser.add_argument('-db', '--database', type=str, default=None,
                    help='Database to use (unused so far)')

parser.add_argument('-t', '--table', type=str, default=None,
                    help='Table to use')

parser.add_argument('-n', '--null', type=str, default='',
                    help='Null value')

parser.add_argument('-e', '--execute', action="store_true",
                    default=False,
                    help='Execute SQL on the server (TODO!!!)')

parser.add_argument('-d', '--data', action="store_true",
                    default=False,
                    help='Insert data')

args = parser.parse_args()


ext = path.splitext(args.input)[1][1:]
input_format = 'Unknown'
if ext in ['fits', 'fit']:
    fitstbl = pyfits.open(args.input)[1]
    input_format = 'fits'
    column_format = COLUMN_FORMATS_FITS
    table_columns = fitstbl.columns
elif ext in ['vo', 'vot', 'votable']:
    fitstbl = votable.parse_single_table(args.input)
    input_format = 'votable'
    column_format = COLUMN_FORMATS_VOTABLE
    table_columns = fitstbl.fields


print 'CREATE TABLE %s (' % args.table
columns = []
column_quotes_needed = []
insert_header = 'insert into %s values ' % args.table
for column in table_columns:
    if input_format == 'fits':
        form = column.format
        data_type = column_format[form[-1:]]
        if form[-1:] == 'A':
            data_type = data_type % form[:-1]
            column_quotes_needed.append(True)
        else:
            column_quotes_needed.append(False)
    elif input_format == 'votable':
        data_type = column_format[column.datatype]
        if column.datatype == 'char':
            #import ipdb; ipdb.set_trace()
            column_quotes_needed.append(True)
            if column.arraysize == '*':
                data_type = data_type % '' # Arbitrary char length
            else:
                data_type = data_type % ('(%s)' % column.arraysize)
        else:
            column_quotes_needed.append(False)

    columns.append("%s %s not null" % (column.name.replace('.', '_'), data_type))
print ",\n".join(columns)
print ');'

def to_str(item):
    if type(item) == str:
        if item == args.null:
            return 'null'
        else:
            return '"%s"' % item
    elif np.isnan(item):
        return 'Null'
    else:
        return str(item)

if args.data:
    first_line = True
    print insert_header
    if input_format == 'fits':
        data = fitstbl.data
    else:
        data = fitstbl.array.data
    for row in data:
        if first_line:
            print '(',
            first_line = False
        else:
            print ', (',
        #for iitem, item in enumerate(row):
            #if column_quotes_needed[iitem]:
            #    row[iitem] = '"%s"' % row[iitem]
        print '%s)' % ','.join(map(to_str, row))
    print ';'