from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass
class ProtoField:
    number: int
    wire: int
    value: Any


def encode_varint(value: int) -> bytes:
    out = bytearray()
    value = int(value)
    while True:
        b = value & 0x7F
        value >>= 7
        if value:
            out.append(b | 0x80)
        else:
            out.append(b)
            break
    return bytes(out)


def read_varint(data: bytes, pos: int) -> tuple[int, int]:
    value = 0
    shift = 0
    while pos < len(data):
        b = data[pos]
        pos += 1
        value |= (b & 0x7F) << shift
        if not (b & 0x80):
            return value, pos
        shift += 7
        if shift > 70:
            raise ValueError("varint too long")
    raise EOFError("unexpected EOF while reading varint")


def parse_fields(data: bytes) -> list[ProtoField]:
    pos = 0
    fields: list[ProtoField] = []
    while pos < len(data):
        key, pos = read_varint(data, pos)
        number = key >> 3
        wire = key & 7
        if wire == 0:
            value, pos = read_varint(data, pos)
        elif wire == 1:
            value = data[pos:pos + 8]
            pos += 8
        elif wire == 2:
            length, pos = read_varint(data, pos)
            value = data[pos:pos + length]
            pos += length
        elif wire == 5:
            value = data[pos:pos + 4]
            pos += 4
        else:
            raise ValueError(f"unsupported wire type {wire}")
        fields.append(ProtoField(number, wire, value))
    return fields


def build_fields(fields: list[ProtoField]) -> bytes:
    out = bytearray()
    for field in fields:
        key = (field.number << 3) | field.wire
        out.extend(encode_varint(key))
        if field.wire == 0:
            out.extend(encode_varint(int(field.value)))
        elif field.wire == 1:
            raw = bytes(field.value)
            if len(raw) != 8:
                raise ValueError(f"fixed64 field {field.number} has wrong length")
            out.extend(raw)
        elif field.wire == 2:
            raw = bytes(field.value)
            out.extend(encode_varint(len(raw)))
            out.extend(raw)
        elif field.wire == 5:
            raw = bytes(field.value)
            if len(raw) != 4:
                raise ValueError(f"fixed32 field {field.number} has wrong length")
            out.extend(raw)
        else:
            raise ValueError(f"unsupported wire type {field.wire}")
    return bytes(out)


def is_printable_utf8(raw: bytes) -> str | None:
    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError:
        return None
    if "\x00" in text:
        return None
    if all(ch.isprintable() or ch in "\r\n\t" for ch in text):
        return text
    return None


def patch_strings_recursive(raw: bytes, replacements: dict[str, str]) -> bytes:
    try:
        fields = parse_fields(raw)
    except Exception:
        return raw
    changed = False
    new_fields: list[ProtoField] = []
    for field in fields:
        if field.wire != 2 or not isinstance(field.value, (bytes, bytearray)):
            new_fields.append(field)
            continue
        value = bytes(field.value)
        try:
            nested_fields = parse_fields(value)
        except Exception:
            nested_fields = []
        try:
            nested_roundtrip = build_fields(nested_fields) if nested_fields else None
        except Exception:
            nested_roundtrip = None
        if nested_fields and nested_roundtrip == value:
            nested = patch_strings_recursive(value, replacements)
            if nested != value:
                new_fields.append(ProtoField(field.number, field.wire, nested))
                changed = True
                continue
        text = is_printable_utf8(value)
        if text is not None:
            new_text = text
            for old, new in replacements.items():
                new_text = new_text.replace(old, new)
            if new_text != text:
                new_fields.append(ProtoField(field.number, field.wire, new_text.encode("utf-8")))
                changed = True
            else:
                new_fields.append(field)
        else:
            new_fields.append(field)
    return build_fields(new_fields) if changed else raw


def _normalize_exact_text(value: str) -> str:
    return value.replace("/", "\\") if ("\\" in value or "/" in value) else value


def patch_strings_recursive_exact(raw: bytes, replacements: dict[str, str]) -> bytes:
    try:
        fields = parse_fields(raw)
    except Exception:
        return raw
    changed = False
    new_fields: list[ProtoField] = []
    for field in fields:
        if field.wire != 2 or not isinstance(field.value, (bytes, bytearray)):
            new_fields.append(field)
            continue
        value = bytes(field.value)
        try:
            nested_fields = parse_fields(value)
        except Exception:
            nested_fields = []
        try:
            nested_roundtrip = build_fields(nested_fields) if nested_fields else None
        except Exception:
            nested_roundtrip = None
        if nested_fields and nested_roundtrip == value:
            nested = patch_strings_recursive_exact(value, replacements)
            if nested != value:
                new_fields.append(ProtoField(field.number, field.wire, nested))
                changed = True
                continue
        text = is_printable_utf8(value)
        if text is not None:
            key = _normalize_exact_text(text)
            new_text = replacements.get(key, text)
            if new_text != text:
                new_fields.append(ProtoField(field.number, field.wire, new_text.encode("utf-8")))
                changed = True
            else:
                new_fields.append(field)
        else:
            new_fields.append(field)
    return build_fields(new_fields) if changed else raw

