#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

import json
import logging
import re
import sys
import traceback
from collections import defaultdict

from text2sql.dataproc.dusql_dataset_v2 import load_tables

logging.basicConfig(
    level=logging.DEBUG,
    format="%(levelname)s: %(asctime)s %(filename)s" " [%(funcName)s:%(lineno)d][%(process)d] %(message)s",
    datefmt="%m-%d %H:%M:%S",
    filename=None,
    filemode="a",
)

g_date_patt = re.compile(r"(([0-9]{2})[0-9]{2}年)?[0-9]{1,2}月[0-9]{2}日|([0-9]{2})[0-9]{2}年[0-9]{1,2}月")


def get_char_list(sentence):
    def is_ascii(s):
        """check if s is English album or number
        Args:
            s (str): NULL
        Returns: bool
        """
        return ord(s) < 128

    if len(sentence) == 0:
        return []

    lst_result = [sentence[0]]
    last_is_ascii = lst_result[-1].isalnum()
    for char in sentence[1:]:
        if char == " ":
            last_is_ascii = False
            continue
        elif char == "-":
            last_is_ascii = False
            lst_result.append(char)
            continue

        if is_ascii(char) and last_is_ascii:
            lst_result[-1] += char
            continue

        if is_ascii(char):
            last_is_ascii = True
        else:
            last_is_ascii = False

        lst_result.append(char)

    return tuple(lst_result)


def _format_date_cell(old_cell):
    new_cell = old_cell.rstrip("月日")
    new_cell = new_cell.replace("年", "-")
    new_cell = new_cell.replace("月", "-")
    return new_cell


def _build(cells):
    dct_index = defaultdict(set)
    for cell in set(cells):
        if type(cell) is not str:
            continue
        cell = cell.strip()
        if re.match(g_date_patt, cell):
            cell = _format_date_cell(cell)
        cell_chars = get_char_list(cell.lower())
        dct_index[cell.lower()].add((cell, len(cell_chars)))
        for pos in range(len(cell_chars) - 1):
            bigram = cell_chars[pos : pos + 2]
            # tri_gram = cell_chars[pos: pos + 3]
            # four_gram = cell_chars[pos: pos + 4]
            dct_index[bigram].add((cell, len(cell_chars) - 1))
            # dct_index[tri_gram].add((cell, len(cell_chars) - 2))
            # dct_index[four_gram].add(cell)
    return dct_index


def build_cell_index(db_dict):
    for db in db_dict.values():
        column_cells = []
        for column in db.columns:
            cell_index = _build(column.cells)
            column_cells.append(cell_index)
        db.column_cells_index = column_cells


def extract_value_from_sql(sql_json, sql_format="dusql"):
    dct_col_values = defaultdict(list)
    if sql_format == "nl2sql":
        for col, _, val in item["sql"]["conds"]:
            dct_col_values[col].append(val)
        return dct_col_values

    def _merge_dict(base_dict, extra_dict):
        for k, v in extra_dict.items():
            base_dict[k].extend(v)

    def _extract_value_from_sql_cond(cond, dct_col_values):
        if type(cond[3]) is dict:
            new_col_values = extract_value_from_sql(cond[3])
            _merge_dict(dct_col_values, new_col_values)
            return
        col_id = cond[2][1][1]
        dct_col_values[col_id].append(cond[3])
        if cond[4] is not None:
            dct_col_values[col_id].append(cond[4])

    for table_unit in sql_json["from"]["table_units"]:
        if type(table_unit[1]) is dict:
            new_col_values = extract_value_from_sql(table_unit[1])
            _merge_dict(dct_col_values, new_col_values)

    for cond in sql_json["where"][::2]:
        _extract_value_from_sql_cond(cond, dct_col_values)
    for cond in sql_json["having"][::2]:
        _extract_value_from_sql_cond(cond, dct_col_values)

    if sql_json["intersect"] is not None:
        new_col_values = extract_value_from_sql(sql_json["intersect"])
        _merge_dict(dct_col_values, new_col_values)
    if sql_json["union"] is not None:
        new_col_values = extract_value_from_sql(sql_json["union"])
        _merge_dict(dct_col_values, new_col_values)
    if sql_json["except"] is not None:
        new_col_values = extract_value_from_sql(sql_json["except"])
        _merge_dict(dct_col_values, new_col_values)

    return dct_col_values


def search_values(query, db, extra_values):
    lst_match_values = []
    for column, cell_index in zip(db.columns, db.column_cells_index):
        if column.id == 0:
            lst_match_values.append([])
            continue

        candi_cnt = defaultdict(float)
        query_chars = get_char_list(query.lower())
        appear_set = set()
        for pos in range(len(query_chars)):
            unigram = query_chars[pos]
            if len(unigram) > 2 and unigram not in appear_set and unigram in cell_index:
                for cell, base in cell_index[unigram]:
                    candi_cnt[cell] += 1.0 / base
            if pos == len(query_chars) - 1:
                break

            bigram = query_chars[pos : pos + 2]
            if bigram not in cell_index:
                continue
            if bigram in appear_set:
                continue
            appear_set.add(bigram)
            for cell, base in cell_index[bigram]:
                candi_cnt[cell] += 1.0 / base

        if extra_values is not None and column.id in extra_values:
            gold_values = extra_values[column.id]
            for gval in gold_values:
                candi_cnt[str(gval)] += 2.0

        lst_match_values.append(list(sorted(candi_cnt.items(), key=lambda x: x[1], reverse=True))[:10])

    return lst_match_values


if __name__ == "__main__":
    import argparse

    try:
        arg_parser = argparse.ArgumentParser(description="linking candidate values for each column")
        arg_parser.add_argument(
            "input", nargs="?", type=argparse.FileType("r"), default=sys.stdin, help="input file path"
        )
        arg_parser.add_argument("-s", "--db-schema", required=True, help="file path")
        arg_parser.add_argument("-c", "--db-content", required=True, help="file path")
        arg_parser.add_argument(
            "-o", "--output", type=argparse.FileType("w"), default=sys.stdout, help="output file path"
        )
        arg_parser.add_argument("-t", "--is-train", default=False, action="store_true")
        arg_parser.add_argument("-f", "--sql-format", default="dusql", choices=["dusql", "nl2sql", "cspider"])
        args = arg_parser.parse_args()

        sys.stderr.write(">>> loading databases...\n")
        dct_db, _ = load_tables(args.db_schema, args.db_content)
        build_cell_index(dct_db)

        sys.stderr.write(">>> extracting values...\n")
        lst_output = []
        for idx, item in enumerate(json.load(args.input)):
            question_id = item.get("question_id", f"qid{idx:06d}")
            question = item["question"]
            db_id = item["db_id"]
            db = dct_db[db_id]

            extra_values = None
            if args.is_train:
                extra_values = extract_value_from_sql(item["sql"], args.sql_format)

            match_values = search_values(question, db, extra_values)
            lst_output.append(
                {"question_id": question_id, "question": question, "db_id": db_id, "match_values": match_values}
            )

        json.dump(lst_output, args.output, indent=2, ensure_ascii=False)
    except Exception:
        traceback.print_exc()
        # logging.critical(traceback.format_exc())
        exit(-1)
