import tree_sitter_cpp as ts_cpp
from tree_sitter import Language, Node, Parser

from my_mcp.umlaut_config import UmlautConfig

CPP = Language(ts_cpp.language())

UMLAUTS_STRING = {
    "ä": b"AE_UMLT",
    "Ä": b"CAE_UMLT",
    "ö": b"OE_UMLT",
    "Ö": b"COE_UMLT",
    "ü": b"UE_UMLT",
    "Ü": b"CUE_UMLT",
    "ß": b"SS_UMLT",
}

UMLAUTS_CHAR = {
    "ä": "AE_UMLT_CHAR",
    "Ä": "CAE_UMLT_CHAR",
    "ö": "OE_UMLT_CHAR",
    "Ö": "COE_UMLT_CHAR",
    "ü": "UE_UMLT_CHAR",
    "Ü": "CUE_UMLT_CHAR",
    "ß": "SS_UMLT_CHAR",
}

UMLAUTS_COMMENT = {
    "ä": "ae",
    "Ä": "Ae",
    "ö": "oe",
    "Ö": "Oe",
    "ü": "ue",
    "Ü": "Ue",
    "ß": "ss",
}


def make_string_literal(content: list[str] | str, quotes='"') -> str:
    if isinstance(content, list):
        content = "".join(content)
    return quotes + "".join(content) + quotes


def transform_string(string: str, mapping: dict[str, str]) -> str | None:
    if not string.startswith('"') or not string.endswith('"'):
        # raw, special-encoding or user-defined literals -> let's not mess with this
        return None

    content = string[1:-1]
    substring: list[str] = []
    result = []
    change_needed = False
    for c in content:
        replacement = mapping.get(c)
        if replacement is not None:
            change_needed = True
            if substring:
                result.append(make_string_literal(substring))
                substring = []
            result.append(replacement)
        else:
            substring.append(c)

    if not change_needed:
        return None

    if substring:
        result.append(make_string_literal(substring))

    return " ".join(result)


def transform_comment(comment: str, mapping: dict[str, str]) -> str | None:
    result = []
    change_needed = False
    for c in comment:
        replacement = mapping.get(c)
        if replacement is not None:
            change_needed = True
            result.append(replacement)
        else:
            result.append(c)

    if not change_needed:
        return None
    return "".join(result)


def transform_char(char: str, mapping: dict[str, str]) -> str | None:
    if len(char) != 3:
        # some wide string literal or multi-character unicode -> not our problem
        return None

    c = char.strip("'")
    replacement = mapping.get(c)
    if replacement:
        return replacement
    return None


def process_node(
    node: Node, code_bytes: bytes, umlaut_config: UmlautConfig
) -> str | None:
    start = node.start_byte
    end = node.end_byte
    text = code_bytes[start:end].decode("utf-8")

    if node.type == "string_literal":
        return transform_string(text, umlaut_config.string_mapping)
    elif node.type in ("comment", "line_comment", "block_comment"):
        return transform_comment(text, umlaut_config.comment_mapping)
    elif node.type in ("char_literal"):
        return transform_char(text, umlaut_config.comment_mapping)
    return None


def fix_document(code: str, umlaut_config: UmlautConfig) -> str:
    parser = Parser(CPP)

    code_bytes = code.encode("utf-8")
    tree = parser.parse(code_bytes)

    root_node = tree.root_node

    edits = []

    def traverse(node: Node):
        replacement = process_node(node, code_bytes, umlaut_config)
        if replacement:
            edits.append((node.start_byte, node.end_byte, replacement))

        for child in node.children:
            traverse(child)

    traverse(root_node)

    edits.sort(key=lambda x: x[0], reverse=True)

    result_bytes = code_bytes
    for start, end, replacement in edits:
        result_bytes = (
            result_bytes[:start] + replacement.encode("utf-8") + result_bytes[end:]
        )

    return result_bytes.decode("utf-8")
