#!/usr/bin/env python3
"""
Unknown binary file exploration helper.

Main ideas:
- search for known values (numeric / strings)
- inspect offsets in multiple types
- find printable ASCII strings
- dump bytes around candidate offsets
- optionally search with alignment-aware steps

Examples
--------
# Search a float32 value with little endian, typical aligned scan
python bin_parser_to_be_tested.py search test.raw --type float32 --value 120.0 --endian little --step 4 --tolerance 1e-3

# Same, but full brute force scan
python bin_parser_to_be_tested.py search test.raw --type float32 --value 120.0 --step 1

# Search a text fragment directly in raw bytes
python bin_parser_to_be_tested.py search test.raw --type str --value DEVICE

# List printable ASCII strings
python bin_parser_to_be_tested.py strings test.raw --min-length 4

# Read values at a given offset as multiple candidate types
python bin_parser_to_be_tested.py inspect test.raw --offset 0x120 --types float32,int16,uint16,str --sizes 4,2,2,32

# Read an array
python bin_parser_to_be_tested.py array test.raw --offset 0x200 --type float32 --count 16 --step 4
"""

from __future__ import annotations

import argparse
import math
import os
import re
import struct
import sys
from typing import Iterable, List, Optional, Sequence, Tuple

__all__ = [
    "load_bin_file",
    "find_ascii_strings",
    "search_data",
    "get_data",
    "read_array",
    "hex_dump",
]

TYPE_MAP = {
    "float32": ("f", 4),
    "float64": ("d", 8),
    "uint8":   ("B", 1),
    "int8":    ("b", 1),
    "uint16":  ("H", 2),
    "int16":   ("h", 2),
    "uint32":  ("I", 4),
    "int32":   ("i", 4),
    "uint64":  ("Q", 8),
    "int64":   ("q", 8),
}

PRINTABLE_ASCII_RE = re.compile(rb"[\x20-\x7E]+")


# -----------------------------------------------------------------------------
# Core helpers
# -----------------------------------------------------------------------------

def load_bin_file(file_path: str) -> bytes:
    with open(file_path, "rb") as f:
        return f.read()


def endian_prefix(endian: str) -> str:
    if endian == "little":
        return "<"
    if endian == "big":
        return ">"
    raise ValueError(f"Unsupported endian: {endian}")


def parse_offset(text: str) -> int:
    text = text.strip().lower()
    return int(text, 0)


def default_step_for_type(vartype: str) -> int:
    if vartype == "str":
        return 1
    _, size = TYPE_MAP[vartype]
    return size


def guess_scalar_category(text: str) -> str:
    try:
        int(text, 0)
        return "int"
    except Exception:
        pass
    try:
        float(text)
        return "float"
    except Exception:
        pass
    return "str"


def safe_slice(data: bytes, start: int, end: int) -> bytes:
    start = max(0, start)
    end = min(len(data), end)
    if end < start:
        end = start
    return data[start:end]


# -----------------------------------------------------------------------------
# ASCII utilities
# -----------------------------------------------------------------------------

def find_ascii_strings(bin_data: bytes, istart: int = 0, ilast: Optional[int] = None,
                       min_length: int = 4) -> List[Tuple[int, str]]:
    """Find printable ASCII substrings.

    Returns absolute offsets in the original file.
    """
    if ilast is None:
        ilast = len(bin_data) - 1
    if ilast < istart:
        return []

    region = bin_data[istart:ilast + 1]
    pattern = rb"[\x20-\x7E]{%d,}" % min_length
    results: List[Tuple[int, str]] = []

    for match in re.finditer(pattern, region):
        offset = istart + match.start()  # absolute offset
        value = match.group().decode("ascii", errors="ignore")
        results.append((offset, value))
    return results


def extract_ascii_prefix(bin_data: bytes, offset: int, size: int) -> str:
    chunk = bin_data[offset: offset + size]
    result = bytearray()
    for b in chunk:
        if 0x20 <= b <= 0x7E:
            result.append(b)
        else:
            break
    return result.decode("ascii", errors="ignore")


def extract_c_string(bin_data: bytes, offset: int, max_size: int = 256,
                     encoding: str = "ascii") -> str:
    chunk = bin_data[offset: offset + max_size]
    nul = chunk.find(b"\x00")
    if nul >= 0:
        chunk = chunk[:nul]
    return chunk.decode(encoding, errors="ignore")


def search_raw_string(bin_data: bytes, needle: str, itop: int = 0,
                      ilast: Optional[int] = None, encoding: str = "ascii") -> List[int]:
    if ilast is None:
        ilast = len(bin_data) - 1
    if ilast < itop:
        return []
    haystack = bin_data[itop: ilast + 1]
    target = needle.encode(encoding, errors="ignore")
    if not target:
        return []

    offsets: List[int] = []
    pos = 0
    while True:
        found = haystack.find(target, pos)
        if found < 0:
            break
        offsets.append(itop + found)
        pos = found + 1
    return offsets


# -----------------------------------------------------------------------------
# Binary decoding
# -----------------------------------------------------------------------------

def get_data(bin_data: bytes, offset: int, size: int, vartype: str,
             endian: str = "little", varname: Optional[str] = None,
             str_mode: str = "ascii_prefix"):
    prefix = endian_prefix(endian)

    if offset < 0 or offset + size > len(bin_data):
        raise ValueError(
            f"Offset/size out of range: offset={offset}, size={size}, file_size={len(bin_data)}"
        )

    if vartype == "str":
        if str_mode == "ascii_prefix":
            return extract_ascii_prefix(bin_data, offset, size)
        if str_mode == "c_string":
            return extract_c_string(bin_data, offset, max_size=size)
        if str_mode == "fixed":
            raw = bin_data[offset: offset + size]
            return raw.decode("ascii", errors="ignore")
        raise ValueError(f"Unsupported str_mode: {str_mode}")

    if vartype not in TYPE_MAP:
        raise ValueError(f"Unsupported type: {vartype}")

    fmt, expected_size = TYPE_MAP[vartype]
    if size != expected_size:
        raise ValueError(
            f"Invalid size {size} for type {vartype} (expected {expected_size})"
        )

    try:
        return struct.unpack_from(prefix + fmt, bin_data, offset)[0]
    except Exception as e:
        name = varname if varname is not None else "(unnamed)"
        print(
            f"Error in get_data(): failed to read [{name}] "
            f"(type {vartype}) from offset={offset} size={size}: {e}",
            file=sys.stderr,
        )
        return None


def read_array(bin_data: bytes, offset: int, vartype: str, count: int,
               endian: str = "little", step: Optional[int] = None) -> List[object]:
    if vartype == "str":
        raise ValueError("read_array() does not support vartype='str'")

    fmt, size = TYPE_MAP[vartype]
    if step is None:
        step = size
    out = []
    for i in range(count):
        off = offset + i * step
        if off + size > len(bin_data):
            break
        out.append(get_data(bin_data, off, size, vartype, endian=endian))
    return out


# -----------------------------------------------------------------------------
# Search
# -----------------------------------------------------------------------------

def numeric_match(val, target, atol: float, rtol: float) -> bool:
    if isinstance(val, float):
        if math.isnan(val) or math.isinf(val):
            return False
        return abs(val - target) <= (atol + rtol * abs(target))
    return val == target


def search_data(bin_data: bytes, itop: int, ilast: int, size: int, vartype: str,
                endian: str, value, tolerance: float = 1e-6, rtol: float = 0.0,
                step: int = 1, varname: str = "not given",
                str_encoding: str = "ascii") -> List[int]:
    """Search for matching values.

    For numeric types, scans in increments of `step` bytes.
    For strings, performs raw byte substring search.
    """
    if ilast < itop:
        return []
    if step <= 0:
        raise ValueError("step must be >= 1")

    prefix = endian_prefix(endian)
    matches: List[int] = []

    if vartype == "str":
        return search_raw_string(bin_data, str(value), itop, ilast, encoding=str_encoding)

    if vartype not in TYPE_MAP:
        raise ValueError(f"Unsupported type: {vartype}")

    fmt, expected_size = TYPE_MAP[vartype]
    if size != expected_size:
        raise ValueError(
            f"Invalid size {size} for type {vartype} (expected {expected_size})"
        )

    max_offset = ilast - size + 1
    for offset in range(itop, max_offset + 1, step):
        try:
            val = struct.unpack_from(prefix + fmt, bin_data, offset)[0]
        except struct.error:
            continue

        if numeric_match(val, value, atol=tolerance, rtol=rtol):
            matches.append(offset)

    return matches


def search_nearby_pattern(matches_a: Sequence[int], matches_b: Sequence[int],
                          delta: int, tolerance: int = 0) -> List[Tuple[int, int]]:
    """Return pairs (a, b) such that b ~= a + delta.

    Useful when you know two fields should appear at a fixed relative offset.
    """
    set_b = set(matches_b)
    results = []
    if tolerance <= 0:
        for a in matches_a:
            b = a + delta
            if b in set_b:
                results.append((a, b))
        return results

    for a in matches_a:
        center = a + delta
        for b in range(center - tolerance, center + tolerance + 1):
            if b in set_b:
                results.append((a, b))
    return results


# -----------------------------------------------------------------------------
# Display helpers
# -----------------------------------------------------------------------------

def hex_dump(bin_data: bytes, offset: int, length: int = 64, width: int = 16) -> str:
    start = max(0, offset)
    end = min(len(bin_data), start + length)
    lines = []
    for row_start in range(start, end, width):
        row = bin_data[row_start: row_start + width]
        hex_part = " ".join(f"{b:02X}" for b in row)
        ascii_part = "".join(chr(b) if 32 <= b <= 126 else "." for b in row)
        lines.append(f"{row_start:08X}  {hex_part:<{width * 3 - 1}}  |{ascii_part}|")
    return "\n".join(lines)


def format_value(value) -> str:
    if isinstance(value, float):
        return f"{value:.12g}"
    return str(value)


# -----------------------------------------------------------------------------
# CLI
# -----------------------------------------------------------------------------

def cmd_search(args: argparse.Namespace) -> int:
    data = load_bin_file(args.infile)
    ilast = len(data) - 1 if args.ilast is None else min(args.ilast, len(data) - 1)
    itop = max(0, args.itop)

    vartype = args.vartype
    if vartype != "str":
        category = guess_scalar_category(args.value)
        if category == "int":
            value = int(args.value, 0)
        elif category == "float":
            value = float(args.value)
        else:
            raise ValueError(f"Numeric search requires numeric --value, got: {args.value}")
        _, size = TYPE_MAP[vartype]
        step = args.step if args.step is not None else default_step_for_type(vartype)
    else:
        value = args.value
        size = args.size if args.size is not None else len(value)
        step = 1

    matches = search_data(
        data,
        itop=itop,
        ilast=ilast,
        size=size,
        vartype=vartype,
        endian=args.endian,
        value=value,
        tolerance=args.tolerance,
        rtol=args.rtol,
        step=step,
    )

    if not matches:
        print("No matches found.")
        return 1

    print(f"Found {len(matches)} match(es).")
    for off in matches[:args.max_hits]:
        print(f"offset 0x{off:08X} ({off})")
        if args.dump > 0:
            base = max(0, off - args.dump_before)
            print(hex_dump(data, base, length=args.dump_before + args.dump))
            print()
    if len(matches) > args.max_hits:
        print(f"... {len(matches) - args.max_hits} more match(es) omitted")
    return 0


def cmd_strings(args: argparse.Namespace) -> int:
    data = load_bin_file(args.infile)
    ilast = len(data) - 1 if args.ilast is None else min(args.ilast, len(data) - 1)
    items = find_ascii_strings(data, istart=max(0, args.itop), ilast=ilast,
                               min_length=args.min_length)
    if args.contains:
        items = [(off, s) for off, s in items if args.contains in s]

    if not items:
        print("No ASCII strings found.")
        return 1

    for off, s in items[:args.max_hits]:
        print(f"0x{off:08X} ({off}): {s}")
    if len(items) > args.max_hits:
        print(f"... {len(items) - args.max_hits} more string(s) omitted")
    return 0


def cmd_inspect(args: argparse.Namespace) -> int:
    data = load_bin_file(args.infile)
    offset = args.offset
    types = [x.strip() for x in args.types.split(",") if x.strip()]
    sizes = [int(x, 0) for x in args.sizes.split(",")] if args.sizes else []

    print(f"File: {args.infile}")
    print(f"Offset: 0x{offset:08X} ({offset})")

    for i, vartype in enumerate(types):
        if vartype == "str":
            size = sizes[i] if i < len(sizes) else args.str_size
        else:
            _, size = TYPE_MAP[vartype]
        try:
            value = get_data(data, offset, size, vartype, endian=args.endian,
                             str_mode=args.str_mode)
            print(f"- {vartype:<8} size={size:<3} value={format_value(value)}")
        except Exception as e:
            print(f"- {vartype:<8} error: {e}")

    if args.dump > 0:
        base = max(0, offset - args.dump_before)
        print()
        print(hex_dump(data, base, length=args.dump_before + args.dump))
    return 0


def cmd_array(args: argparse.Namespace) -> int:
    data = load_bin_file(args.infile)
    step = args.step if args.step is not None else default_step_for_type(args.vartype)
    values = read_array(data, args.offset, args.vartype, args.count,
                        endian=args.endian, step=step)
    for i, value in enumerate(values):
        off = args.offset + i * step
        print(f"[{i:4d}] offset=0x{off:08X} ({off:8d}) value={format_value(value)}")
    return 0


def cmd_dump(args: argparse.Namespace) -> int:
    data = load_bin_file(args.infile)
    print(hex_dump(data, args.offset, length=args.length, width=args.width))
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Explore unknown binary files by searching, inspecting, and dumping data."
    )
    sub = p.add_subparsers(dest="command", required=True)

    # search
    ps = sub.add_parser("search", help="Search for numeric values or strings")
    ps.add_argument("infile", help="Input binary file")
    ps.add_argument("--type", dest="vartype", required=True,
                    choices=list(TYPE_MAP.keys()) + ["str"])
    ps.add_argument("--value", required=True,
                    help="Target value. For integers, 0x... is accepted.")
    ps.add_argument("--endian", choices=["little", "big"], default="little")
    ps.add_argument("--itop", type=parse_offset, default=0,
                    help="Start offset (default: 0)")
    ps.add_argument("--ilast", type=parse_offset, default=None,
                    help="Last offset inclusive (default: EOF)")
    ps.add_argument("--size", type=lambda x: int(x, 0), default=None,
                    help="Used only for strings; otherwise inferred from type")
    ps.add_argument("--step", type=lambda x: int(x, 0), default=None,
                    help="Scan step in bytes (default: numeric type size, str=1)")
    ps.add_argument("--tolerance", type=float, default=1e-6,
                    help="Absolute tolerance for float comparison")
    ps.add_argument("--rtol", type=float, default=0.0,
                    help="Relative tolerance for float comparison")
    ps.add_argument("--dump", type=lambda x: int(x, 0), default=64,
                    help="Dump N bytes from the displayed window")
    ps.add_argument("--dump-before", type=lambda x: int(x, 0), default=16,
                    help="Show this many bytes before each hit")
    ps.add_argument("--max-hits", type=int, default=20)
    ps.set_defaults(func=cmd_search)

    # strings
    pst = sub.add_parser("strings", help="List printable ASCII strings")
    pst.add_argument("infile")
    pst.add_argument("--min-length", type=int, default=4)
    pst.add_argument("--contains", default=None,
                     help="Filter strings containing this substring")
    pst.add_argument("--itop", type=parse_offset, default=0)
    pst.add_argument("--ilast", type=parse_offset, default=None)
    pst.add_argument("--max-hits", type=int, default=200)
    pst.set_defaults(func=cmd_strings)

    # inspect
    pi = sub.add_parser("inspect", help="Interpret one offset as multiple types")
    pi.add_argument("infile")
    pi.add_argument("--offset", type=parse_offset, required=True)
    pi.add_argument("--types", default="float32,int16,uint16,str",
                    help="Comma separated list, e.g. float32,int16,str")
    pi.add_argument("--sizes", default=None,
                    help="Optional comma separated sizes aligned to --types")
    pi.add_argument("--str-size", type=lambda x: int(x, 0), default=32,
                    help="Default size for vartype=str")
    pi.add_argument("--str-mode", choices=["ascii_prefix", "c_string", "fixed"],
                    default="ascii_prefix")
    pi.add_argument("--endian", choices=["little", "big"], default="little")
    pi.add_argument("--dump", type=lambda x: int(x, 0), default=64)
    pi.add_argument("--dump-before", type=lambda x: int(x, 0), default=16)
    pi.set_defaults(func=cmd_inspect)

    # array
    pa = sub.add_parser("array", help="Read a numeric array from an offset")
    pa.add_argument("infile")
    pa.add_argument("--offset", type=parse_offset, required=True)
    pa.add_argument("--type", dest="vartype", required=True,
                    choices=list(TYPE_MAP.keys()))
    pa.add_argument("--count", type=int, required=True)
    pa.add_argument("--step", type=lambda x: int(x, 0), default=None,
                    help="Stride between elements in bytes")
    pa.add_argument("--endian", choices=["little", "big"], default="little")
    pa.set_defaults(func=cmd_array)

    # dump
    pd = sub.add_parser("dump", help="Hex dump from an offset")
    pd.add_argument("infile")
    pd.add_argument("--offset", type=parse_offset, required=True)
    pd.add_argument("--length", type=lambda x: int(x, 0), default=128)
    pd.add_argument("--width", type=int, default=16)
    pd.set_defaults(func=cmd_dump)

    return p


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())