# agent/sql/process_sql.py
import re, json, requests
from config import *
from agent.sql.exec import exec_sql_s

def replace_date_with_day(sql):
    """
    This function replaces instances of exact date conditions in a SQL 
    statement from a format like:
        TradingDate = 'YYYY-MM-DD'
    to:
        date(TradingDate) = 'YYYY-MM-DD'
    
    Parameters:
        sql (str): The original SQL statement.
        
    Returns:
        str: The modified SQL statement, or the original if no match is found.
    """
    # Regex pattern to match patterns like: ColumnName = 'YYYY-MM-DD'
    pattern = r"([.\w]+)\s*=\s*'(\d{4}-\d{2}-\d{2})'"

    def replace_func(match):
        column_name = match.group(1)
        date_value = match.group(2)
        return f"date({column_name}) = '{date_value}'"

    new_sql = re.sub(pattern, replace_func, sql)

    # If no change was made, return the original SQL
    return new_sql if new_sql != sql else sql


def extract_sql(text):
    """
    从输入文本中提取被包裹在三重反引号（sql ... ）中的 SQL 语句
    Extracts an SQL statement from a block of text enclosed in triple backticks:
        ```sql
        SELECT ...
        ```
    
    Parameters:
        text (str): The full text containing an SQL statement.
        
    Returns:
        str: The extracted SQL statement, or a message if not found.
    """
    sql_pattern = re.compile(r'```sql(.*?)```', re.DOTALL)
    match = sql_pattern.search(text)
    if match:
        # Strip leading and trailing whitespace from the matched SQL
        return match.group(1).strip()
    else:
        # print(f"--------------------extract_sql else:{text}")
        return f"No SQL statement found :{text}."


def select_data(sql_text):
    """
    将指定的 SQL 查询 发送到某个 API 端点（通过 POST 请求），并返回该 API 的响应结果
    Sends the given SQL query to a specified endpoint and returns the JSON response.
    
    Parameters:
        sql_text (str): The SQL query to be executed.
        
    Returns:
        str: The JSON response from the API, formatted with indentation.
    """
    url = "https://comm.chatglm.cn/finglm2/api/query"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f'Bearer {Access_Token}'
    }
    data = {
        "sql": sql_text,  # e.g. SELECT * FROM constantdb.secumain LIMIT 10
        "limit": 50
    }
    response = requests.post(url, headers=headers, json=data)
    try:
        return json.dumps(response.json(), indent=2, ensure_ascii=False)
    except:
        return str(response.json())


def clean_sql_statement(sql):
    """
    清理SQL
    
    Parameters:
        sql (str): 输入的 SQL 语句。
    
    Returns:
        str: 清理后的 SQL 语句。
    """
    #     清理 SQL 语句中第一个 SELECT 前面的内容。
    match = re.search(r'\bSELECT\b.*', sql, re.IGNORECASE | re.DOTALL)
    if match:
        return match.group(0)
    else:
        return sql


def wrap_date_in_sql_with_conditions(sql_statement):
    # 正则匹配符合条件的时间字段
    pattern = r"(?<!\(\()\b(\w+)\b\s*=\s*'(2\d{3}-\d{2}-\d{2})'"
    # 替换，将字段用 date() 包裹
    modified_sql = re.sub(pattern, r"date(\1) = '\2'", sql_statement)
    return modified_sql


def ensure_date_in_between(sql_query):
    # 匹配 BETWEEN 左侧的字段，确保字段被 DATE() 包裹
    pattern = r"(\bAND\s+)(\w+)(\s+BETWEEN\s+)"
    replacement = r"\1date(\2)\3"
    updated_query = re.sub(pattern, replacement, sql_query)
    return updated_query


def clean_sql_query(sql):
    # 去除首尾空白字符
    cleaned = sql.strip()
    
    # 如果包含分号，去除分号后面的内容
    if ';' in cleaned:
        cleaned = cleaned.split(';')[0]
    
    # 检查并去除末尾的点
    if cleaned.endswith('.'):
        cleaned = cleaned[:-1]
    
    return cleaned
    
    
def validate_and_fix_sql_tables(sql_statement):
    """
    检查并修复 SQL 语句中的表名和库名是否匹配。
    
    参数：
        sql_statement (str): SQL 查询语句
        
    返回：
        str: 修复后的 SQL 查询语句
    """
    database_L_zh = database_L_zh_all
    if len(QUESTION_TYPE_LIST) == 0:
        database_L_zh = database_L_zh_all
    elif len(QUESTION_TYPE_LIST) >= 2:
        database_L_zh = database_L_zh_all
    elif "港股" in QUESTION_TYPE_LIST:
        database_L_zh = database_L_zh_hk
    elif "美股" in QUESTION_TYPE_LIST:
        database_L_zh = database_L_zh_us
    elif "A股" in QUESTION_TYPE_LIST:
        database_L_zh = database_L_zh_cn
    else:
        database_L_zh = database_L_zh_all

    # 构造表名到完整库表名的映射
    table_to_full_name = {item['数据表名'].split('.')[1]: item['数据表名'] for item in database_L_zh}

    # 提取 SQL 中的表名（包括库名）
    matches = re.findall(r"FROM\s+([a-zA-Z0-9_\.]+)", sql_statement, re.IGNORECASE)

    if not matches:
        if DEBUG_VER == 3:
            print("未在 SQL 语句中找到表名")
        return sql_statement

    fixed_sql = sql_statement

    for full_table_name in matches:
        if '.' in full_table_name:
            db_name, table_name = full_table_name.split('.')
        else:
            db_name = None
            table_name = full_table_name

        # 检查表名是否在 database_L_zh 中
        if table_name in table_to_full_name:
            correct_full_name = table_to_full_name[table_name]
            correct_db_name = correct_full_name.split('.')[0]

            # 如果库名不匹配或缺失，替换为正确的库表名
            if not db_name or db_name != correct_db_name:
                fixed_sql = re.sub(
                    rf"\b{re.escape(full_table_name)}\b",  # 精确匹配表名
                    correct_full_name, 
                    fixed_sql
                )
                if DEBUG_VER == 3:
                    print(f"修正库表名: {full_table_name} -> {correct_full_name}")
        else:
            if DEBUG_VER == 3:
                print(f"表名未找到: {table_name}")
            pass
    return fixed_sql
    

def to_select(text):
    """
    High-level function that:
      1. Extracts SQL from the given text.
      2. Optimizes the extracted SQL by converting date columns to 'date(...)'.
      3. Executes the optimized SQL through select_data and returns the result.
    
    Parameters:
        text (str): The input text containing an SQL statement.
        
    Returns:
        str: The JSON response from the SQL query.
    """
    global prev_tables_name_list
    global QUESTION_TYPE_LIST
    sql_statement = extract_sql(text)
    if DEBUG_VER == 3:
        print('***********Extracted SQL****************')
    sql_statement = clean_sql_statement(sql_statement)
    sql_statement = wrap_date_in_sql_with_conditions(sql_statement)
    sql_statement = validate_and_fix_sql_tables(sql_statement)
    sql_statement = ensure_date_in_between(sql_statement)
    sql_statement = clean_sql_query(sql_statement)
    if DEBUG_VER == 3:
        print(f"---------------sql_statement:{sql_statement}")
        print('***********Extracted SQL****************')
        
    optimized_sql = replace_date_with_day(sql_statement)
        
    result = select_data(optimized_sql)
    if 'count' in result:
        if '"count": 0' not in result:
            prev_tables_name_list += [i.get('数据表名') for i in table_maps if i.get('数据表名') in sql_statement]
        if ('"count": 0' in result) and ('AS' in sql_statement or 'as' in sql_statement):
            result = f"查询异常。SQL语句：{sql_statement}没有找到数据，请判断使用的字段是否正确，可尝试其它字段或库表查询，并把当次0也作为结果返回（也可能是真实结果），表的结构如下：{table_maps_LL}"
    prev_tables_name_list = list(set(prev_tables_name_list))
    if DEBUG_VER == 3:
        print(f"----------prev_tables_name_list:{prev_tables_name_list}")
    for table_name in prev_tables_name_list:
        if table_name in content_CN and table_name not in content_US and table_name not in content_HK:
            QUESTION_TYPE_LIST.append("A股")
            QUESTION_TYPE_LIST = list(set(QUESTION_TYPE_LIST))

    if "查询执行失败" in result:
        LL = [i for i in table_maps if i.get('数据表名') in sql_statement]
        result = result + f"表的结构如下：{LL}"
    if "Unknown column" in result:
        LL = [i for i in table_maps if i.get('数据表名') in sql_statement]
        result = result + f"表的结构如下：{LL}"
    if '"data": []' in result:
        LL = [i for i in table_maps if i.get('数据表名') in sql_statement]
        result = f"查询异常。SQL语句：{sql_statement}没有找到数据，结果如下：{result}。请用其它相关字段或库表查询，表的结构如下：{LL}"
    if 'No database selected' in result:
        LL = [i for i in table_maps if i.get('数据表名') in sql_statement]
        result = f"查询异常。SQL语句：{sql_statement}没有找到数据，结果如下：{result}。请用其它相关字段或库表查询，表的结构如下：{LL}"
    try:
        data_dict = json.loads(result)
        data = data_dict.get("data", [])
        if len(data) >= 50:
            result = result + " 数据库最多返回50条数据，如果需要的数据超过50条，请用多sql嵌套组合或者聚合函数count、sum、avg、max、min等来查询。如果回复的内容数量跟预期不符，思考SQL语句是否存在问题。"
        for item in data:
            if all(value is None for value in item.values()):
                result = f"查询异常。SQL语句：{sql_statement}可能没有找到数据，结果如下：{result}。请判断使用的字段是否正确，可尝试其它类似字段或库表查询（如时间字段更换成EndDate或InfoPublDate），并把当次null也作为结果返回（也可能是真实结果），表的结构如下：{table_maps_LL}"
    except Exception as e:
        pass
    # print(f"--------------sql_result:{result}")
    # print(f"--------------type sql_result:{type(result)}")
    return result

def extract_table_names(sql):
    # 正则匹配 FROM 关键字后的库表名，库表名后会有空格
    matches = re.findall(r'FROM\s+([\w\.]+)\s', sql, re.IGNORECASE)
    return matches

def all_tables_in_prompt(tables_name_list, main_sql_prompts):
    lower_prompts = main_sql_prompts.lower()
    return all(table.lower() in lower_prompts for table in tables_name_list)