"""
patch_notebook_v2.py
====================
Correctly patches Video_Deepfake_Detection_Cloud.ipynb to fix the following:

  1. ZeroDivisionError in Cell 12 when DATASET_ROOT has no labeled videos.
  2. prob_real / prob_fake key-swap bug in the predict_single() return dict
     (Cell 8). The notebook displayed correct labels in the console table,
     but saved swapped values in the JSON result file.

Run from the video_detection/ folder:
    python patch_notebook_v2.py

After running, re-upload the patched .ipynb to Colab or Kaggle and re-run
Cell 8 onward to produce correct JSON exports and Cell 12 evaluation.
"""
import json
import os
import sys

NB_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "notebooks",
    "Video_Deepfake_Detection_Cloud.ipynb",
)


def load_notebook(path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_notebook(nb, path):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(nb, f, indent=2, ensure_ascii=False)
    print("[OK] Saved:", path)


def patch_notebook():
    if not os.path.exists(NB_PATH):
        print("[ERR] Notebook not found:", NB_PATH)
        sys.exit(1)

    nb = load_notebook(NB_PATH)
    total_patches = 0

    for cell in nb["cells"]:
        if cell["cell_type"] != "code":
            continue

        source_lines = cell["source"]
        joined = "".join(source_lines)

        # ------------------------------------------------------------------ #
        # Patch 1 – Fix prob_real/prob_fake swap in predict_single (Cell 8)  #
        # ------------------------------------------------------------------ #
        # Original (wrong):
        #   prob_fake = 1.0 - y_val if label == 'REAL' else y_val
        # Explanation:
        #   GenConViT class-0 = fake, class-1 = real.
        #   max_prediction_value returns y_val = P(winning class) for FAKE,
        #   but abs(1 - P(real)) for REAL, which also equals P(fake).
        #   So y_val is always P(fake) — we just need:
        #     prob_fake = y_val
        #     prob_real = 1 - y_val
        OLD_PROB = (
            "    prob_fake = 1.0 - y_val if label == 'REAL' else y_val\n"
            "    prob_real = 1.0 - prob_fake"
        )
        NEW_PROB = (
            "    # y_val is always P(fake) from max_prediction_value()\n"
            "    prob_fake = y_val\n"
            "    prob_real = 1.0 - prob_fake"
        )

        if OLD_PROB in joined:
            patched = joined.replace(OLD_PROB, NEW_PROB, 1)
            lines = patched.splitlines(keepends=True)
            # ipynb convention: last line has no trailing newline
            if lines and lines[-1].endswith("\n"):
                lines[-1] = lines[-1][:-1]
            cell["source"] = lines
            total_patches += 1
            # Reload so Patch 2 sees the updated source
            joined = "".join(cell["source"])
            print("[OK] Patch 1 applied: fixed prob_real/prob_fake swap in predict_single()")

        # ------------------------------------------------------------------ #
        # Patch 2 – Fix ZeroDivisionError in Cell 12                         #
        # ------------------------------------------------------------------ #
        OLD_ZERO_DIV = (
            "print(f'\\n\U0001f4ca Results:"
            " Acc={correct/total:.1%}"
            "  TPR={tp/(tp+fn):.1%}"
            "  FPR={fp/(fp+tn):.1%}')"
        )
        # Fall back to ASCII version in case the emoji was already stripped
        OLD_ZERO_DIV_ASCII = (
            "print(f'\\n Results:"
            " Acc={correct/total:.1%}"
            "  TPR={tp/(tp+fn):.1%}"
            "  FPR={fp/(fp+tn):.1%}')"
        )
        NEW_ZERO_DIV = (
            "if total == 0:\n"
            "        print('[WARN] No labeled videos processed. "
            "Check DATASET_ROOT structure.')\n"
            "        print('   Expected: DATASET_ROOT/real/*.mp4"
            " and DATASET_ROOT/fake/*.mp4')\n"
            "    else:\n"
            "        acc = correct / total\n"
            "        tpr = tp / (tp + fn) if (tp + fn) > 0 else 0.0\n"
            "        fpr = fp / (fp + tn) if (fp + tn) > 0 else 0.0\n"
            "        print(f'\\n[RESULTS]"
            " Acc={acc:.1%}  TPR={tpr:.1%}  FPR={fpr:.1%}')"
        )

        matched_old = OLD_ZERO_DIV if OLD_ZERO_DIV in joined else (
            OLD_ZERO_DIV_ASCII if OLD_ZERO_DIV_ASCII in joined else None
        )
        if matched_old:
            patched = joined.replace(matched_old, NEW_ZERO_DIV, 1)
            lines = patched.splitlines(keepends=True)
            if lines and lines[-1].endswith("\n"):
                lines[-1] = lines[-1][:-1]
            cell["source"] = lines
            total_patches += 1
            print("[OK] Patch 2 applied: ZeroDivisionError guard in Cell 12")

    if total_patches == 0:
        print("[INFO] No patches needed (already applied or pattern not found).")
    else:
        save_notebook(nb, NB_PATH)
        print()
        print("[DONE]", total_patches, "patch(es) applied successfully.")
        print("   Re-upload", os.path.basename(NB_PATH),
              "to Colab/Kaggle and re-run from Cell 8.")


if __name__ == "__main__":
    patch_notebook()