#!/usr/bin/env python3
"""
Lemenzo Voice v4.5 - hardened cellular Call-TX integration.

BUILD-TIME TOOL ONLY.
This script is NOT installed on the phone.

It patches:

1. frameworks/av/services/audiopolicy/managerdefault/AudioPolicyManager.cpp

   Legacy HAL 2.x (bluejay / lynx):
   - Create the real cellular Call-TX patch when Lemenzo Voice is enabled.

   Modern HAL / AIDL (akita / tokay):
   - Force only the real isCallTx() path through the AudioFlinger
     software bridge when Lemenzo Voice cellular processing is enabled.

2. frameworks/av/services/audioflinger/PatchPanel.cpp

   - Prevent PassthruPatchRecord only for an explicitly enabled
     Lemenzo cellular Telephony-TX software patch.

   - This guarantees the software Call-TX path passes RecordThread,
     where the existing Lemenzo Voice DSP hook is already installed.

Voice OFF preserves the original Google/AOSP routing behavior.

The script performs a full preflight before modifying either source file.
If writing fails, both framework files are automatically restored.
"""

from pathlib import Path
import argparse
import os
import shutil
import sys
import tempfile


MARKER_APM = "LEMENZO_VOICE_CALLTX_V45"
MARKER_PATCHPANEL = "LEMENZO_VOICE_CALLTX_NO_PASSTHRU_V45"
MARKER_PATCHPANEL_STATE = "LEMENZO_VOICE_CALLTX_PATCHPANEL_STATE_V45"

BACKUP_SUFFIX = ".pre-lemenzo-calltx-v45"

OLD_MARKERS = (
    "LEMENZO_VOICE_CALLTX_V43",
    "LEMENZO_VOICE_CALLTX_V44",
    "LEMENZO_VOICE_CALLTX_NO_PASSTHRU_V44",
)


def die(message: str) -> None:
    print(f"ERROR: {message}", file=sys.stderr)
    raise SystemExit(1)


def read_text(path: Path) -> str:
    try:
        return path.read_text(encoding="utf-8")
    except Exception as exc:
        die(f"could not read {path}: {exc}")


def backup_path(path: Path) -> Path:
    return Path(str(path) + BACKUP_SUFFIX)


def create_fresh_backup(path: Path) -> Path:
    """
    Always snapshot the CURRENT clean source before applying v4.5.

    This deliberately overwrites an old v4.5 backup because an old backup
    may belong to an earlier repo revision/QPR and must not be used for
    rollback of the current source tree.
    """
    bak = backup_path(path)

    try:
        shutil.copy2(path, bak)
    except Exception as exc:
        die(f"could not create backup {bak}: {exc}")

    print(f"Backup: {bak}")
    return bak


def restore_from_backup(path: Path) -> None:
    bak = backup_path(path)

    if not bak.exists():
        die(f"backup missing: {bak}")

    try:
        shutil.copy2(bak, path)
    except Exception as exc:
        die(f"could not restore {path}: {exc}")

    print(f"Restored: {path}")


def atomic_write(path: Path, content: str) -> None:
    """
    Replace one source file atomically.

    The temporary file is created in the same directory so os.replace()
    stays on the same filesystem.
    """
    fd = None
    tmp_name = None

    try:
        fd, tmp_name = tempfile.mkstemp(
            prefix=f".{path.name}.lemenzo-",
            suffix=".tmp",
            dir=str(path.parent),
            text=True,
        )

        with os.fdopen(fd, "w", encoding="utf-8", newline="") as handle:
            fd = None
            handle.write(content)
            handle.flush()
            os.fsync(handle.fileno())

        # Preserve mode/ownership metadata as far as normal build-user
        # permissions allow.
        shutil.copymode(path, tmp_name)

        os.replace(tmp_name, path)
        tmp_name = None

    except Exception:
        if fd is not None:
            try:
                os.close(fd)
            except Exception:
                pass

        if tmp_name is not None:
            try:
                os.unlink(tmp_name)
            except Exception:
                pass

        raise


def reject_old_or_mixed_markers(apm_text: str, panel_text: str) -> None:
    for marker in OLD_MARKERS:
        if marker in apm_text or marker in panel_text:
            die(
                f"old Lemenzo Call-TX marker detected: {marker}. "
                "Restore the original framework files before applying v4.5."
            )

    apm_v45 = MARKER_APM in apm_text
    panel_v45 = MARKER_PATCHPANEL in panel_text

    if apm_v45 and panel_v45:
        print("Lemenzo Voice v4.5 Call-TX integration is already fully applied.")
        raise SystemExit(0)

    if apm_v45 != panel_v45:
        die(
            "partial v4.5 integration detected. "
            "One framework file is patched while the other is not. "
            "Use --rollback or restore clean framework sources before retrying."
        )


def build_audio_policy_patch(original: str) -> str:
    text = original

    if MARKER_APM in text:
        die("AudioPolicyManager.cpp is already marked as v4.5 patched")

    # ------------------------------------------------------------
    # system property include
    # ------------------------------------------------------------

    if "#include <sys/system_properties.h>" not in text:
        pivot = '#include "AudioPolicyManager.h"'

        if pivot in text:
            text = text.replace(
                pivot,
                '#include <sys/system_properties.h>\n\n' + pivot,
                1,
            )
        else:
            pivot = "namespace android {"

            if pivot not in text:
                die(
                    "AudioPolicyManager.cpp: could not find include "
                    "or namespace insertion point"
                )

            text = text.replace(
                pivot,
                "#include <sys/system_properties.h>\n\n" + pivot,
                1,
            )

    # ------------------------------------------------------------
    # Lemenzo state helper
    # ------------------------------------------------------------

    ns = "namespace android {"

    if ns not in text:
        die("AudioPolicyManager.cpp: namespace android not found")

    helper = r'''namespace android {

// LEMENZO_VOICE_CALLTX_V45
//
// Missing properties are intentionally treated as OFF.
// Cellular Call-TX is active only when BOTH properties are explicitly "1".
static bool lemenzoVoicePropertyEnabled(const char* key) {
    char value[PROP_VALUE_MAX] = {};
    const int len = __system_property_get(key, value);

    return len > 0
            && value[0] == '1'
            && value[1] == '\0';
}

static bool lemenzoVoiceCellularTxEnabled() {
    return lemenzoVoicePropertyEnabled(
                    "persist.audio.secureline.enabled")
            && lemenzoVoicePropertyEnabled(
                    "persist.audio.secureline.cellular_enabled");
}
'''

    text = text.replace(ns, helper, 1)

    # ------------------------------------------------------------
    # Legacy HAL 2.x cellular TX patch creation
    #
    # Needed for bluejay / lynx where the normal legacy logic
    # does not create a TX patch for the primary built-in mic.
    # ------------------------------------------------------------

    old_legacy = '''createTxPatch = !(availablePrimaryModuleInputDevices().contains(txSourceDevice)) &&
                (txSinkDevice != 0);'''

    new_legacy = '''const bool originalCreateTxPatch =
                !(availablePrimaryModuleInputDevices().contains(txSourceDevice)) &&
                (txSinkDevice != 0);

        const bool lemenzoCreateTxPatch =
                lemenzoVoiceCellularTxEnabled() &&
                (txSinkDevice != 0);

        createTxPatch =
                originalCreateTxPatch ||
                lemenzoCreateTxPatch;

        ALOGI_IF(lemenzoCreateTxPatch && !originalCreateTxPatch,
                "Lemenzo Voice: enabling cellular TX audio patch for legacy HAL");'''

    count = text.count(old_legacy)

    if count != 1:
        die(
            "AudioPolicyManager.cpp: expected exactly one legacy "
            f"createTxPatch anchor, found {count}. "
            "Source differs from the audited Android 16 tree."
        )

    text = text.replace(old_legacy, new_legacy, 1)

    # ------------------------------------------------------------
    # Force the real Call-TX descriptor through SW bridge
    #
    # isCallTx() is Android's explicit telephony TX discriminator.
    # ------------------------------------------------------------

    old_bridge = '''if (!srcDevice->hasSameHwModuleAs(sinkDevice) ||
                        (srcDevice->getModuleVersionMajor() < 3) ||
                        !srcDevice->getModule()->supportsPatch(srcDevice, sinkDevice) ||
                        (!sourceDesc->isInternal() &&
                         srcDevice->getAudioPort()->getGains().size() == 0)) {'''

    new_bridge = '''const bool lemenzoForceCallTxSwBridge =
                        sourceDesc != nullptr &&
                        sourceDesc->isCallTx() &&
                        lemenzoVoiceCellularTxEnabled();

                if (lemenzoForceCallTxSwBridge ||
                        !srcDevice->hasSameHwModuleAs(sinkDevice) ||
                        (srcDevice->getModuleVersionMajor() < 3) ||
                        !srcDevice->getModule()->supportsPatch(srcDevice, sinkDevice) ||
                        (!sourceDesc->isInternal() &&
                         srcDevice->getAudioPort()->getGains().size() == 0)) {

                    ALOGI_IF(lemenzoForceCallTxSwBridge,
                            "Lemenzo Voice: forcing cellular Call-TX through "
                            "AudioFlinger SW bridge");'''

    count = text.count(old_bridge)

    if count != 1:
        die(
            "AudioPolicyManager.cpp: expected exactly one SW bridge "
            f"selection anchor, found {count}. "
            "Source differs from the audited Android 16 tree."
        )

    text = text.replace(old_bridge, new_bridge, 1)

    # ------------------------------------------------------------
    # Final AudioPolicy sanity
    # ------------------------------------------------------------

    if text.count(MARKER_APM) != 1:
        die("AudioPolicyManager.cpp: v4.5 marker validation failed")

    if "sourceDesc->isCallTx()" not in text:
        die("AudioPolicyManager.cpp: isCallTx() validation failed")

    if "lemenzoForceCallTxSwBridge" not in text:
        die("AudioPolicyManager.cpp: SW bridge integration validation failed")

    if "lemenzoCreateTxPatch" not in text:
        die("AudioPolicyManager.cpp: legacy TX integration validation failed")

    return text


def build_patch_panel_patch(original: str) -> str:
    text = original

    if MARKER_PATCHPANEL in text:
        die("PatchPanel.cpp is already marked as v4.5 patched")

    # ------------------------------------------------------------
    # system property include
    # ------------------------------------------------------------

    if "#include <sys/system_properties.h>" not in text:
        pivot = "namespace android {"

        if pivot not in text:
            die("PatchPanel.cpp: namespace android not found")

        text = text.replace(
            pivot,
            "#include <sys/system_properties.h>\n\n" + pivot,
            1,
        )

    # ------------------------------------------------------------
    # Lemenzo state helper
    # ------------------------------------------------------------

    ns = "namespace android {"

    if ns not in text:
        die("PatchPanel.cpp: namespace android not found")

    helper = r'''namespace android {

// LEMENZO_VOICE_CALLTX_PATCHPANEL_STATE_V45
//
// Missing properties are intentionally OFF.
// Voice OFF preserves the original AOSP/Google PatchPanel behavior.
static bool lemenzoPatchPanelPropertyEnabled(const char* key) {
    char value[PROP_VALUE_MAX] = {};
    const int len = __system_property_get(key, value);

    return len > 0
            && value[0] == '1'
            && value[1] == '\0';
}

static bool lemenzoPatchPanelCellularTxEnabled() {
    return lemenzoPatchPanelPropertyEnabled(
                    "persist.audio.secureline.enabled")
            && lemenzoPatchPanelPropertyEnabled(
                    "persist.audio.secureline.cellular_enabled");
}
'''

    text = text.replace(ns, helper, 1)

    # ------------------------------------------------------------
    # Prevent DIRECT passthrough only for ACTIVE Lemenzo
    # cellular Telephony-TX.
    #
    # This guarantees RecordThread is used, where the existing
    # Lemenzo DSP hook already processes PCM.
    # ------------------------------------------------------------

    old_passthru = '''    const bool usePassthruPatchRecord =
            (inputFlags & AUDIO_INPUT_FLAG_DIRECT)
            && (outputFlags & AUDIO_OUTPUT_FLAG_DIRECT);'''

    new_passthru = '''    // LEMENZO_VOICE_CALLTX_NO_PASSTHRU_V45
    //
    // Only an explicitly enabled Lemenzo cellular Telephony-TX
    // software patch is forced through RecordThread.
    //
    // Voice OFF preserves the original AOSP/Google passthrough behavior.
    const bool isTelephonyTxPatch =
            streamType == AUDIO_STREAM_VOICE_CALL
            && mAudioPatch.num_sinks > 0
            && mAudioPatch.sinks[0].type == AUDIO_PORT_TYPE_DEVICE
            && mAudioPatch.sinks[0].ext.device.type ==
                    AUDIO_DEVICE_OUT_TELEPHONY_TX;

    const bool lemenzoCallTxNeedsRecordThread =
            isTelephonyTxPatch &&
            lemenzoPatchPanelCellularTxEnabled();

    const bool usePassthruPatchRecord =
            !lemenzoCallTxNeedsRecordThread
            && (inputFlags & AUDIO_INPUT_FLAG_DIRECT)
            && (outputFlags & AUDIO_OUTPUT_FLAG_DIRECT);'''

    count = text.count(old_passthru)

    if count != 1:
        die(
            "PatchPanel.cpp: expected exactly one passthru selection anchor, "
            f"found {count}. Source differs from the audited Android 16 tree."
        )

    text = text.replace(old_passthru, new_passthru, 1)

    # ------------------------------------------------------------
    # Final PatchPanel sanity
    # ------------------------------------------------------------

    if text.count(MARKER_PATCHPANEL) != 1:
        die("PatchPanel.cpp: no-passthru v4.5 marker validation failed")

    if text.count(MARKER_PATCHPANEL_STATE) != 1:
        die("PatchPanel.cpp: state helper marker validation failed")

    if "AUDIO_DEVICE_OUT_TELEPHONY_TX" not in text:
        die("PatchPanel.cpp: Telephony TX validation failed")

    if "lemenzoCallTxNeedsRecordThread" not in text:
        die("PatchPanel.cpp: RecordThread guard validation failed")

    return text


def rollback_both(apm: Path, panel: Path) -> None:
    apm_bak = backup_path(apm)
    panel_bak = backup_path(panel)

    missing = []

    if not apm_bak.exists():
        missing.append(str(apm_bak))

    if not panel_bak.exists():
        missing.append(str(panel_bak))

    if missing:
        die(
            "cannot rollback because backup file(s) are missing:\n  "
            + "\n  ".join(missing)
        )

    try:
        shutil.copy2(apm_bak, apm)
        shutil.copy2(panel_bak, panel)
    except Exception as exc:
        die(f"rollback failed: {exc}")

    print(f"Restored: {apm}")
    print(f"Restored: {panel}")
    print("Lemenzo Voice v4.5 Call-TX integration rolled back.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Apply hardened Lemenzo Voice v4.5 Call-TX integration"
    )

    parser.add_argument(
        "--tree",
        default=".",
        help="Android source tree root",
    )

    parser.add_argument(
        "--rollback",
        action="store_true",
        help="restore both framework files from v4.5 backups",
    )

    args = parser.parse_args()

    tree = Path(args.tree).resolve()

    apm = (
        tree
        / "frameworks/av/services/audiopolicy/managerdefault/"
          "AudioPolicyManager.cpp"
    )

    panel = (
        tree
        / "frameworks/av/services/audioflinger/"
          "PatchPanel.cpp"
    )

    for path in (apm, panel):
        if not path.is_file():
            die(f"missing required framework file: {path}")

    if args.rollback:
        rollback_both(apm, panel)
        return

    # ------------------------------------------------------------
    # READ CURRENT SOURCE
    # ------------------------------------------------------------

    apm_original = read_text(apm)
    panel_original = read_text(panel)

    # ------------------------------------------------------------
    # REFUSE OLD / MIXED / HALF-PATCHED SOURCE
    # ------------------------------------------------------------

    reject_old_or_mixed_markers(
        apm_original,
        panel_original,
    )

    print("Preflight: validating AudioPolicyManager.cpp...")
    apm_patched = build_audio_policy_patch(apm_original)

    print("Preflight: validating PatchPanel.cpp...")
    panel_patched = build_patch_panel_patch(panel_original)

    # ------------------------------------------------------------
    # IMPORTANT:
    #
    # At this point BOTH patches have been successfully generated
    # in memory.
    #
    # Nothing in the Android tree has been modified yet.
    # ------------------------------------------------------------

    print("Preflight passed for both framework files.")

    # ------------------------------------------------------------
    # CREATE FRESH BACKUPS OF THE EXACT CURRENT SOURCE
    # ------------------------------------------------------------

    create_fresh_backup(apm)
    create_fresh_backup(panel)

    # ------------------------------------------------------------
    # WRITE TRANSACTION
    # ------------------------------------------------------------

    try:
        atomic_write(apm, apm_patched)
        print(f"Patched: {apm}")

        atomic_write(panel, panel_patched)
        print(f"Patched: {panel}")

    except Exception as exc:
        print(
            f"ERROR: framework write failed: {exc}",
            file=sys.stderr,
        )

        print(
            "Attempting automatic rollback of both framework files...",
            file=sys.stderr,
        )

        rollback_errors = []

        try:
            shutil.copy2(backup_path(apm), apm)
        except Exception as rollback_exc:
            rollback_errors.append(
                f"AudioPolicyManager rollback failed: {rollback_exc}"
            )

        try:
            shutil.copy2(backup_path(panel), panel)
        except Exception as rollback_exc:
            rollback_errors.append(
                f"PatchPanel rollback failed: {rollback_exc}"
            )

        if rollback_errors:
            for message in rollback_errors:
                print(f"CRITICAL: {message}", file=sys.stderr)

            die(
                "automatic rollback was not completely successful. "
                "Inspect both framework files manually before building."
            )

        die(
            "patch operation failed, but both framework files "
            "were restored successfully"
        )

    # ------------------------------------------------------------
    # VERIFY FILES AFTER WRITE
    # ------------------------------------------------------------

    apm_after = read_text(apm)
    panel_after = read_text(panel)

    if MARKER_APM not in apm_after:
        print(
            "CRITICAL: AudioPolicy post-write verification failed. "
            "Rolling back.",
            file=sys.stderr,
        )
        rollback_both(apm, panel)
        raise SystemExit(1)

    if MARKER_PATCHPANEL not in panel_after:
        print(
            "CRITICAL: PatchPanel post-write verification failed. "
            "Rolling back.",
            file=sys.stderr,
        )
        rollback_both(apm, panel)
        raise SystemExit(1)

    print()
    print("======================================================")
    print("Lemenzo Voice v4.5 Call-TX integration SUCCESS")
    print("======================================================")
    print()
    print("AudioPolicy:")
    print("  Legacy HAL 2.x:")
    print("    Lemenzo Voice ON -> create cellular TX patch")
    print()
    print("  Modern/AIDL HAL:")
    print("    isCallTx() + Voice ON -> force AudioFlinger SW bridge")
    print()
    print("PatchPanel:")
    print("  Voice ON + real Telephony TX:")
    print("    disable PassthruPatchRecord")
    print("    -> RecordThread")
    print("    -> existing Lemenzo DSP")
    print()
    print("Voice OFF:")
    print("  original Google/AOSP call routing behavior")
    print()
    print("This Python script is BUILD-TIME ONLY.")
    print("It is not installed on the phone.")
    print()


if __name__ == "__main__":
    main()