#!/usr/bin/python2.7
from __future__ import unicode_literals

import argparse
import collections
import os.path
import polib
import sqlparse
import sys

import rad.sql

class Collector:
    def __init__(self):
        self._path = ""
        self._current_table = None
        self._msgids = collections.defaultdict(set)

    def set_path(self, path):
        self._path = path

    def get_msgids(self):
        return self._msgids

    def table(self, name):
        self._current_table = name
        self._push("table:{}".format(name))

    def column(self, name):
        assert self._current_table is not None
        self._push("column:{}.{}".format(self._current_table, name))

    def enum(self, value):
        self._push(value)

    def _push(self, disp):
        self._msgids[disp].add((self._path, ""))

def parse_statement(stmt, collect):
    def match_token(tok, ttype, tvalue=None):
        if tok.ttype not in ttype:
            return False
        if tvalue is not None and tok.value.upper() != tvalue.upper():
            return False
        return True

    T = sqlparse.tokens.Token
    skip_parentheses = 0

    state = "SEARCH_CREATE"
    for tok in stmt.flatten():
        if tok.ttype in T.Comment or tok.ttype in T.Text.Whitespace:
            continue

        if skip_parentheses > 0:
            if match_token(tok, T.Punctuation, "("):
                skip_parentheses += 1
            elif match_token(tok, T.Punctuation, ")"):
                skip_parentheses -= 1
        else:
            if state == "SEARCH_CREATE":
                if match_token(tok, T.Keyword.DDL, "CREATE"):
                    state = "CREATE_BEGIN"
            elif state == "CREATE_BEGIN":
                if match_token(tok, T.Keyword, "TABLE"):
                    state = "CREATE_AFTER_TABLE"
            elif state == "CREATE_AFTER_TABLE":
                if match_token(tok, T.Name):
                    collect.table(rad.sql.get_name(tok))
                    state = "CREATE_AFTER_NAME"
            elif state == "CREATE_AFTER_NAME":
                if match_token(tok, T.Punctuation, "("):
                    state = "CREATE_DEFINITION"

            elif state == "CREATE_DEFINITION":
                if match_token(tok, T.Name):
                    collect.column(rad.sql.get_name(tok))
                    state = "IN_DEFINITION"
                elif match_token(tok, T.Punctuation, "("):
                    # syntax error (unexpected open parentheses)
                    state = "ABORT_SYNTAX_ERROR"
                    break
                else:
                    state = "SKIP_DEFINITION"

            elif state == "SKIP_DEFINITION":
                if match_token(tok, T.Punctuation, "("):
                    skip_parentheses = 1
                elif match_token(tok, T.Punctuation, ","):
                    state = "CREATE_DEFINITION"
                elif match_token(tok, T.Punctuation, ")"):
                    state = "FINISH"
                    break

            elif state == "IN_DEFINITION":
                if match_token(tok, T.Name, "ENUM"):
                    state = "ENUM_BEGIN"
                elif match_token(tok, T.Punctuation, "("):
                    skip_parentheses = 1
                elif match_token(tok, T.Punctuation, ","):
                    state = "CREATE_DEFINITION"
                elif match_token(tok, T.Punctuation, ")"):
                    state = "FINISH"
                    break

            elif state == "ENUM_BEGIN":
                if match_token(tok, T.Punctuation, "("):
                    state = "IN_ENUM"
                else:
                    state = "SEARCH_ENUM"
            elif state == "IN_ENUM":
                if match_token(tok, T.Literal.String):
                    string_value = rad.sql.sql_parse_string(tok)
                    state = "ENUM_AFTER_STRING"
                else:
                    state = "IN_ENUM"
            elif state == "ENUM_AFTER_STRING":
                if match_token(tok, T.Punctuation, ":"):
                    state = "ENUM_ANNOT"
                elif match_token(tok, T.Punctuation, ")"):
                    collect.enum(string_value)
                    state = "SEARCH_ENUM"
                else:
                    collect.enum(string_value)
                    state = "IN_ENUM"
            elif state == "ENUM_ANNOT":
                if match_token(tok, T.Literal.String):
                    collect.enum(rad.sql.sql_parse_string(tok))
                    state = "ENUM_AFTER_ANNOT"
                elif match_token(tok, T.Punctuation, ")"):
                    state = "SEARCH_ENUM"
            elif state == "ENUM_AFTER_ANNOT":
                if match_token(tok, T.Punctuation, ")"):
                    state = "IN_DEFINITION"
                else:
                    state = "IN_ENUM"

    return state

parser = argparse.ArgumentParser()

parser.add_argument("-o", "--output", dest="outfile", metavar="PATH", required=True)
parser.add_argument("infiles", metavar="SQL_FILE", nargs="+")

args = parser.parse_args()

collect = Collector()

# extract:
#  - table and column names
#  - the values (or labels, where available) of ENUMs
for path in args.infiles:
    with open(path, "r") as f:
        collect.set_path(path)

        for stmt in sqlparse.parsestream(f, encoding="utf-8"):
            if stmt.get_type() != "CREATE":
                continue

            finish_state = parse_statement(stmt, collect)

# add the extracted values to an existing PO file
if os.path.exists(args.outfile):
    po = polib.pofile(args.outfile, encoding="utf-8")
    changed = False
else:
    po = polib.POFile()
    changed = True

for (msgid, occurrences) in collect.get_msgids().items():
    entry = polib.POEntry(msgid=msgid, occurrences=list(occurrences))
    existing = po.find(msgid)
    if existing is not None:
        entry.occurrences = list(occurrences | set(existing.occurrences))
        existing.merge(entry)
        if entry != existing:
            changed = True
    else:
        po.append(entry)
        changed = True

# generate output PO file
if changed:
    print("Adding schema related strings to '{}'.".format(args.outfile))
    po.save(args.outfile)
else:
    print("No changes to be made to '{}'.".format(args.outfile))