import re
from dataclasses import dataclass
from functools import lru_cache
from typing import List, Tuple

from .parser import Parser


@lru_cache(maxsize=1024)
def fix_statement(statement: str) -> str:
    # to fix https://github.com/tidbcloud/SimCluster/issues/5
    # Handle composite key patterns first, before the generic ( ... ) replacement
    def fix_composite_key_patterns(stmt):
        def fix_composite_in(match):
            columns_part = match.group(1)
            column_count = columns_part.count(",") + 1
            placeholder_values = ", ".join(["1"] * column_count)
            # print(
            #     f"Debug: Matched composite pattern with {column_count} columns: '{match.group(0)}'"
            # )
            return f"( {columns_part} ) in ( ( {placeholder_values} ) )"

        # original_stmt = stmt
        # Pattern for composite key IN clauses: ( col1 , col2 , ... ) in ( ( ... ) )
        stmt = re.sub(
            r"\(\s*([^)]+)\s*\)\s*in\s*\(\s*\(\s*\.\.\.\s*\)\s*\)",
            fix_composite_in,
            stmt,
        )

        # if stmt != original_stmt:
        #     print(f"Debug: Fixed composite pattern: '{original_stmt}' -> '{stmt}'")

        return stmt

    statement = fix_composite_key_patterns(statement)

    replace_map = {
        "count ( ? )": "count ( 1 )",
        "values ( ... )": "values ( )",
        "limit ?": "limit 1",
        "offset ?": "offset 0",
        "- ?": "- 0",
        "+ ?": "+ 0",
        ") , ? )": ") , 1 )",
        "in ( ... )": "in ( 1 )",
        "( ... )": "( 1 )",  # This catches remaining cases after composite keys are handled
    }
    for key, value in replace_map.items():
        statement = statement.replace(key, value)

    statement = re.sub(
        r"\b(count|sum|max|min|date_sub|now) \(",
        lambda match: match.group(1) + "(",
        statement,
    )
    statement = re.sub(
        # NULL/TRUE/FALSE/UNKNOWN are all valid values for is/is not.
        # We only use NULL because it's most common.
        r"\b(is|is not) \?",
        lambda match: match.group(1) + " NULL",
        statement,
    )

    return statement


@dataclass
class Statement:
    digest_text: str
    stmt_type: str
    schema_name: str
    table_names: List[str]
    summary_begin_time: str
    exec_count: int
    parser: Parser

    def correlations(self) -> List[Tuple[str, str]]:
        return self.parser.parse_correlations(fix_statement(self.digest_text))

    def placeholders(self):
        return self.parser.parse_placeholders(
            fix_statement(self.digest_text), self.stmt_type
        )

    def fixed_statement(self) -> str:
        return fix_statement(self.digest_text).replace("?", "%s")