#!/usr/bin/env python3
"""Search the extracted target rootfs for audio-like assets.

The scan is static: it reads files, checks names, file(1) output, magic
signatures, and ffprobe stream metadata. It does not execute target binaries.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import subprocess
from pathlib import Path


TARGET_ROOT_DIRS = [
    "bin",
    "boot",
    "connectivity",
    "data",
    "dev",
    "etc",
    "games",
    "home",
    "lib",
    "lib64",
    "linuxrc",
    "lost+found",
    "media",
    "media2",
    "mnt",
    "opt",
    "proc",
    "root",
    "run",
    "sbin",
    "sys",
    "tmp",
    "usr",
    "var",
]

NAME_RE = re.compile(
    r"(audio|sound|music|voice|vocal|speech|sfx|song|speaker|alsa|bluealsa|"
    r"\.(wav|wave|ogg|oga|mp3|flac|aac|m4a|opus|mid|midi|aif|aiff|aifc|"
    r"bnk|wem|fsb|bank|it|xm|mod|s3m)$)",
    re.IGNORECASE,
)

TEXT_RE = re.compile(r"(audio|sound|music|voice|vocal|speech|sfx|song|alsa|volume)", re.IGNORECASE)

MAGICS = [
    ("ogg", b"OggS"),
    ("flac", b"fLaC"),
    ("id3", b"ID3"),
    ("midi", b"MThd"),
    ("riff", b"RIFF"),
    ("aiff", b"FORM"),
    ("mp4_ftyp", b"ftyp"),
    ("wwise_bkhd", b"BKHD"),
    ("fsb5", b"FSB5"),
    ("opus_head", b"OpusHead"),
]


def target_files(root: Path) -> list[Path]:
    paths: list[Path] = []
    for name in TARGET_ROOT_DIRS:
        start = root / name
        if not start.exists():
            continue
        if start.is_file():
            paths.append(start)
            continue
        for dirpath, dirnames, filenames in os.walk(start):
            dirnames[:] = [d for d in dirnames if d != "__pycache__"]
            for filename in filenames:
                path = Path(dirpath) / filename
                if path.is_file() and not path.is_symlink():
                    paths.append(path)
    return sorted(paths)


def rel(path: Path, root: Path) -> str:
    return "/" + path.relative_to(root).as_posix()


def run_file(paths: list[Path]) -> dict[Path, str]:
    out: dict[Path, str] = {}
    batch_size = 128
    for i in range(0, len(paths), batch_size):
        batch = paths[i : i + batch_size]
        result = subprocess.run(["file", *map(str, batch)], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, check=False)
        for line in result.stdout.decode("utf-8", "replace").splitlines():
            left, sep, right = line.partition(":")
            if not sep:
                continue
            out[Path(left)] = right.strip()
    return out


def interesting_file_type(desc: str) -> bool:
    return bool(
        re.search(
            r"audio|mpeg.*layer|mp3|ogg|wave|flac|aac|midi|aiff|matroska|iso media|quicktime|video",
            desc,
            re.IGNORECASE,
        )
    )


def scan_magic(path: Path) -> list[dict[str, object]]:
    hits: list[dict[str, object]] = []
    size = path.stat().st_size
    if size == 0:
        return hits
    chunk_size = 1024 * 1024
    overlap = 64
    offset = 0
    prev_tail = b""
    with path.open("rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            data = prev_tail + chunk
            base = offset - len(prev_tail)
            for name, magic in MAGICS:
                start = 0
                while True:
                    idx = data.find(magic, start)
                    if idx < 0:
                        break
                    absolute = base + idx
                    if absolute >= offset and len([h for h in hits if h["magic"] == name]) < 32:
                        context = data[idx : idx + 16].hex()
                        detail = ""
                        if magic == b"RIFF" and idx + 12 <= len(data):
                            detail = data[idx + 8 : idx + 12].decode("ascii", "replace")
                        elif magic == b"FORM" and idx + 12 <= len(data):
                            detail = data[idx + 8 : idx + 12].decode("ascii", "replace")
                        elif magic == b"ftyp" and idx >= 4 and idx + 8 <= len(data):
                            detail = data[idx + 4 : idx + 8].decode("ascii", "replace")
                        hits.append({"magic": name, "offset": absolute, "detail": detail, "context": context})
                    start = idx + 1
            prev_tail = data[-overlap:]
            offset += len(chunk)
    return hits


def ffprobe(path: Path, ffprobe_bin: str) -> list[dict[str, object]]:
    result = subprocess.run(
        [
            ffprobe_bin,
            "-v",
            "error",
            "-show_entries",
            "stream=index,codec_type,codec_name",
            "-of",
            "json",
            str(path),
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        check=False,
    )
    if result.returncode != 0:
        return []
    try:
        return json.loads(result.stdout.decode("utf-8", "replace")).get("streams", [])
    except json.JSONDecodeError:
        return []


def first_strings(path: Path, limit: int = 20) -> list[str]:
    result = subprocess.run(["strings", "-a", "-n", "4", str(path)], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, check=False)
    strings: list[str] = []
    for line in result.stdout.decode("utf-8", "replace").splitlines():
        if TEXT_RE.search(line):
            strings.append(line[:240])
            if len(strings) >= limit:
                break
    return strings


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path("."))
    parser.add_argument("--output", type=Path, default=Path("analysis/audio-search"))
    parser.add_argument("--ffprobe", default="/opt/homebrew/bin/ffprobe")
    args = parser.parse_args()

    root = args.root.resolve()
    files = target_files(root)
    args.output.mkdir(parents=True, exist_ok=True)

    file_types = run_file(files)
    filename_hits = [p for p in files if NAME_RE.search(p.name)]
    file_type_hits = [p for p in files if interesting_file_type(file_types.get(p, ""))]

    magic_rows: list[dict[str, object]] = []
    ffprobe_candidates: set[Path] = set(file_type_hits)
    for p in files:
        hits = scan_magic(p)
        if hits:
            for hit in hits:
                magic_rows.append({"path": rel(p, root), **hit})
            if any(hit["magic"] in {"mp4_ftyp", "riff", "ogg", "flac", "id3", "midi", "aiff"} for hit in hits):
                ffprobe_candidates.add(p)

    ffprobe_rows: list[dict[str, object]] = []
    if Path(args.ffprobe).exists():
        for p in sorted(ffprobe_candidates):
            streams = ffprobe(p, args.ffprobe)
            if not streams:
                continue
            ffprobe_rows.append(
                {
                    "path": rel(p, root),
                    "size": p.stat().st_size,
                    "has_audio": any(s.get("codec_type") == "audio" for s in streams),
                    "has_video": any(s.get("codec_type") == "video" for s in streams),
                    "streams": ",".join(f"{s.get('index')}:{s.get('codec_type')}:{s.get('codec_name')}" for s in streams),
                }
            )

    string_rows = []
    for p in sorted(set(filename_hits + file_type_hits)):
        string_rows.append({"path": rel(p, root), "strings": first_strings(p)})

    with (args.output / "filename-hits.tsv").open("w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f, delimiter="\t")
        writer.writerow(["path", "size"])
        for p in filename_hits:
            writer.writerow([rel(p, root), p.stat().st_size])

    with (args.output / "file-type-hits.tsv").open("w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f, delimiter="\t")
        writer.writerow(["path", "size", "file_type"])
        for p in file_type_hits:
            writer.writerow([rel(p, root), p.stat().st_size, file_types.get(p, "")])

    with (args.output / "magic-hits.tsv").open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["path", "magic", "offset", "detail", "context"], delimiter="\t")
        writer.writeheader()
        writer.writerows(magic_rows)

    with (args.output / "ffprobe-streams.tsv").open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["path", "size", "has_audio", "has_video", "streams"], delimiter="\t")
        writer.writeheader()
        writer.writerows(ffprobe_rows)

    (args.output / "audio-related-strings.json").write_text(json.dumps(string_rows, indent=2), encoding="utf-8")

    summary = {
        "target_file_count": len(files),
        "filename_hits": len(filename_hits),
        "file_type_hits": len(file_type_hits),
        "magic_hits": len(magic_rows),
        "ffprobe_stream_files": len(ffprobe_rows),
        "ffprobe_audio_files": sum(1 for row in ffprobe_rows if row["has_audio"]),
        "outputs": {
            "filename_hits": str(args.output / "filename-hits.tsv"),
            "file_type_hits": str(args.output / "file-type-hits.tsv"),
            "magic_hits": str(args.output / "magic-hits.tsv"),
            "ffprobe_streams": str(args.output / "ffprobe-streams.tsv"),
            "audio_related_strings": str(args.output / "audio-related-strings.json"),
        },
    }
    (args.output / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))
    return 0


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