"""

Usage
-----
    python detect.py INPUT_FILE
    python detect.py INPUT_FILE --json
    python detect.py INPUT_FILE --audio-model retrained_models/audio/clf.joblib \
                                --video-model retrained_models/video/cnn.pt \
                                --lipsync-model retrained_models/lipsync/fusion.pt
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed


PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))

CORES = {
    "audio":   os.path.join(PROJECT_DIR, "audio-detection_core"),
    "video":   os.path.join(PROJECT_DIR, "video-deepfake-detection_core"),
    "lipsync": os.path.join(PROJECT_DIR, "avh-align_core"),
}


PY = {
    "audio":   "/opt/dfdetect-envs/audio/bin/python",
    "video":   "/opt/dfdetect-envs/video/bin/python",
    "lipsync": "/opt/conda/envs/avh/bin/python",
}

WEIGHTS = {"lipsync": 0.5, "audio": 0.4, "video": 0.1}

FFMPEG = "/usr/bin/ffmpeg"
FFPROBE = "/usr/bin/ffprobe"

# How each core is driven 

CORE_SPEC = {
    "audio":   {"file": "app.py",       "func": "process_audio",       "score_key": "score_audio",   "marker": None},
    "video":   {"file": "app.py",       "func": "process_video",       "score_key": "score_video",   "marker": "## local test"},
    "lipsync": {"file": "avh_align.py", "func": "process_audio_video", "score_key": "score_lipsync", "marker": "#### testing"},
}

RESULT_PREFIX = "__DETECT_RESULT__ "


def _apply_override(module, ns, model_path):
    import torch
    if module == "audio":
        import joblib
        ns["classifier"], ns["scaler"], ns["thresh"] = joblib.load(model_path)
    elif module == "video":
        sd = torch.load(model_path, map_location=ns.get("device", "cpu"))
        if isinstance(sd, dict) and "state_dict" in sd:
            sd = sd["state_dict"]
        ns["classifier"].load_state_dict(sd)
        ns["classifier"].eval()
    elif module == "lipsync":
        sd = torch.load(model_path, map_location="cpu", weights_only=False)
        if isinstance(sd, dict) and "state_dict" in sd:
            sd = sd["state_dict"]
        ns["fusion_model"].load_state_dict(sd)
        ns["fusion_model"].eval()


def run_worker(module, input_path, model_override):
    spec = CORE_SPEC[module]
    core_dir = CORES[module]
    out = {"ok": False, "module": module, "score": None, "detail": None, "error": None}
    try:
        os.chdir(core_dir)
        sys.path.insert(0, core_dir)

        with open(os.path.join(core_dir, spec["file"]), "r") as fh:
            src = fh.read()
        if spec["marker"]:
            src = src.split(spec["marker"])[0]

        ns = {"__name__": "__core_%s__" % module, "__file__": os.path.join(core_dir, spec["file"])}
        exec(compile(src, spec["file"], "exec"), ns)

        ns["load_models"]("model")

      
        try:
            import torch
            if module == "video" and torch.cuda.is_available():
                dev = torch.device("cuda")
                ns["device"] = dev
                if ns.get("model") is not None:
                    ns["model"].to(dev)
                if ns.get("classifier") is not None:
                    ns["classifier"].to(dev)
        except Exception:
            pass

        if model_override:
            _apply_override(module, ns, model_override)

        result = ns[spec["func"]](input_path)
        detail = json.loads(result) if isinstance(result, str) else result
        out["detail"] = detail
        score = detail.get(spec["score_key"]) if isinstance(detail, dict) else None
        if score is None:
            out["error"] = detail.get("error", "core produced no '%s'" % spec["score_key"]) if isinstance(detail, dict) else "no score"
        else:
            out["score"] = float(score)
            out["ok"] = True
    except Exception as exc:
        import traceback
        out["error"] = "%s: %s" % (type(exc).__name__, exc)
        out["trace"] = traceback.format_exc()
    print(RESULT_PREFIX + json.dumps(out))
    return 0


## orchestrator based of the input type
def detect_modality(path):
    try:
        raw = subprocess.run(
            [FFPROBE, "-v", "error", "-print_format", "json", "-show_streams", path],
            capture_output=True, text=True, check=True,
        ).stdout
        streams = json.loads(raw).get("streams", [])
    except Exception as exc:
        raise RuntimeError("ffprobe failed on %s: %s" % (path, exc))

    has_video = any(
        s.get("codec_type") == "video" and s.get("disposition", {}).get("attached_pic", 0) != 1
        for s in streams
    )
    has_audio = any(s.get("codec_type") == "audio" for s in streams)
    if has_video:
        return "video", has_audio
    if has_audio:
        return "audio", True
    raise RuntimeError("No audio or video stream found in %s" % path)


## audio as wav(since the classifier is trained on wav extentions at 16kHz) extrator from video using ffmpeg
def extract_audio(video_path, dest_dir):
    wav = os.path.join(dest_dir, "extracted_audio.wav")
    subprocess.run(
        [FFMPEG, "-y", "-i", video_path, "-vn", "-ac", "1", "-ar", "16000", "-f", "wav", wav,
         "-loglevel", "error"],
        check=True,
    )
    return wav


def _log(msg):
    sys.stderr.write(msg + "\n")
    sys.stderr.flush()

## check which GPUs are free for parallel processing and faster inference time
def _pick_gpus(n):

    try:
        out = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,memory.free,utilization.gpu",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True).stdout
        rows = []
        for ln in out.strip().splitlines():
            idx, free, util = (x.strip() for x in ln.split(","))
            rows.append((int(idx), int(free), int(util)))
        rows.sort(key=lambda r: (-r[1], r[2])) # sorts the GPU usage
        return [r[0] for r in rows[:n]]
    except Exception:
        return []


def run_core(module, input_path, model_override, verbose, gpu_id=None, online=False):
    interp = PY[module]
    if not os.path.exists(interp):
        return {"ok": False, "module": module, "score": None,
                "error": "interpreter not found: %s" % interp}

    cmd = [interp, os.path.abspath(__file__), "--worker", module, os.path.abspath(input_path)]
    if model_override:
        cmd += ["--model", os.path.abspath(model_override)]

    env = dict(os.environ)
    if gpu_id is not None:
        env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    if not online:
        # weights are local and the HF cache is populated; skipping HF's online
        # metadata checks cuts transformers load time from minutes to seconds.
        env["HF_HUB_OFFLINE"] = "1"
        env["TRANSFORMERS_OFFLINE"] = "1"

    _log("[detect] %-8s starting (loading models)%s ..."
         % (module, "" if gpu_id is None else " on GPU %s" % gpu_id))
    t0 = time.time()
    proc = subprocess.run(cmd, capture_output=True, text=True, env=env)
    dt = time.time() - t0
    if verbose and proc.stdout:
        sys.stderr.write("\n----- [%s] stdout -----\n%s\n" % (module, proc.stdout))
    if verbose and proc.stderr:
        sys.stderr.write("\n----- [%s] stderr -----\n%s\n" % (module, proc.stderr))

    result = None
    for line in proc.stdout.splitlines():
        if line.startswith(RESULT_PREFIX):
            result = json.loads(line[len(RESULT_PREFIX):])
    if result is None:
        tail = (proc.stderr or proc.stdout or "").strip().splitlines()[-8:]
        _log("[detect] %-8s FAILED after %.0fs (exit %d)" % (module, dt, proc.returncode))
        return {"ok": False, "module": module, "score": None,
                "error": "no result from core (exit %d)\n%s" % (proc.returncode, "\n".join(tail))}
    if result.get("ok"):
        _log("[detect] %-8s done in %.0fs  score=%.4f" % (module, dt, result.get("score") or 0.0))
    else:
        _log("[detect] %-8s done in %.0fs  (no score: %s)"
             % (module, dt, (result.get("error") or "").splitlines()[0] if result.get("error") else "?"))
    return result


def orchestrate(args):
    input_path = os.path.abspath(args.input)
    if not os.path.exists(input_path):
        sys.exit("input file not found: %s" % input_path)

    if args.modality:
        modality = args.modality
        has_audio = True
    else:
        modality, has_audio = detect_modality(input_path)

    overrides = {"audio": args.audio_model, "video": args.video_model, "lipsync": args.lipsync_model}

    tmpdir = tempfile.mkdtemp(prefix="detect_")
    targets = {}  
    try:
        if modality == "audio":
            targets["audio"] = input_path
        else:
            if has_audio:
                targets["audio"] = extract_audio(input_path, tmpdir)
            else:
                sys.stderr.write("[warn] video has no audio stream; skipping the audio core\n")
            targets["video"] = input_path
            targets["lipsync"] = input_path

        
        for m in list(targets):
            if getattr(args, "no_" + m):
                targets.pop(m)

        
        ordered = list(targets.items())
        free_gpus = _pick_gpus(len(ordered))
        gpu_for = {m: (free_gpus[i] if i < len(free_gpus) else None)
                   for i, (m, _) in enumerate(ordered)}

        _log("[detect] modality=%s; running cores: %s%s"
             % (modality, ", ".join(targets),
                "" if args.sequential or len(targets) < 2 else " (parallel)"))

        per_module = {}
        if args.sequential or len(targets) < 2:
            for module, fpath in ordered:
                per_module[module] = run_core(module, fpath, overrides[module], args.verbose,
                                              gpu_for[module], args.online)
        else:
            with ThreadPoolExecutor(max_workers=len(ordered)) as ex:
                futs = {ex.submit(run_core, m, f, overrides[m], args.verbose, gpu_for[m], args.online): m
                        for m, f in ordered}
                for fut in as_completed(futs):
                    per_module[futs[fut]] = fut.result()
    finally:
        try:
            import shutil
            shutil.rmtree(tmpdir, ignore_errors=True)
        except Exception:
            pass


    ok = {m: r for m, r in per_module.items() if r.get("ok") and r.get("score") is not None}
    failed = {m: r.get("error") for m, r in per_module.items() if m not in ok}

    final_score = None
    used_weights = {}
    if ok:
        wsum = sum(WEIGHTS[m] for m in ok)
        used_weights = {m: WEIGHTS[m] / wsum for m in ok}
        final_score = sum(used_weights[m] * ok[m]["score"] for m in ok)

    summary = {
        "input": input_path,
        "modality": modality,
        "scores": {m: round(r["score"], 4) for m, r in ok.items()},
        "weights_used": {m: round(w, 4) for m, w in used_weights.items()},
        "final_fakeness_score": round(final_score, 4) if final_score is not None else None,
        "final_label": (None if final_score is None else ("fake" if final_score >= 0.5 else "real")),
        "failed": failed,
    }

    if args.json:
        print(json.dumps(summary, indent=2))
    else:
        print(_pretty(summary, per_module))

    return 0 if final_score is not None else 2


## function for text summary result of the module, it can be json format as well
def _pretty(summary, per_module):
    lines = []
    lines.append("=" * 56)
    lines.append("Deepfake detection report")
    lines.append("=" * 56)
    lines.append("input    : %s" % summary["input"])
    lines.append("modality : %s" % summary["modality"])
    lines.append("-" * 56)
    lines.append("%-10s %-10s %-10s %s" % ("module", "fakeness", "weight", "status"))
    for m in ("lipsync", "audio", "video"):
        if m not in per_module:
            continue
        r = per_module[m]
        if r.get("ok"):
            lines.append("%-10s %-10.4f %-10.4f %s" % (
                m, r["score"], summary["weights_used"].get(m, 0.0), "ok"))
        else:
            err = (r.get("error") or "").splitlines()[0] if r.get("error") else "failed"
            lines.append("%-10s %-10s %-10s FAILED: %s" % (m, "-", "-", err))
    lines.append("-" * 56)
    if summary["final_fakeness_score"] is None:
        lines.append("FINAL    : no score (all cores failed)")
    else:
        lines.append("FINAL    : %.4f  ->  %s" % (
            summary["final_fakeness_score"], summary["final_label"].upper()))
    lines.append("=" * 56)
    return "\n".join(lines)


def build_parser():
    p = argparse.ArgumentParser(description="Unified deepfake detection (audio / video / lipsync).")
    p.add_argument("input", nargs="?", help="path to an audio or video file")
    p.add_argument("--json", action="store_true", help="emit machine-readable JSON only")
    p.add_argument("--verbose", action="store_true", help="show each core's stdout/stderr")
    p.add_argument("--modality", choices=["audio", "video"], help="force input modality (skip ffprobe)")
    p.add_argument("--audio-model", help="override path to retrained audio joblib")
    p.add_argument("--video-model", help="override path to retrained video CNN .pt")
    p.add_argument("--lipsync-model", help="override path to retrained lipsync FusionModel .pt")
    p.add_argument("--sequential", action="store_true", help="run cores one at a time instead of in parallel")
    p.add_argument("--online", action="store_true",
                   help="allow Hugging Face online checks (needed only to fetch a model the first time)")
    p.add_argument("--no-audio", action="store_true", help="disable the audio core")
    p.add_argument("--no-video", action="store_true", help="disable the video core")
    p.add_argument("--no-lipsync", action="store_true", help="disable the lipsync core")
    p.add_argument("--worker", choices=list(CORE_SPEC), help=argparse.SUPPRESS)
    p.add_argument("--model", dest="worker_model", help=argparse.SUPPRESS)
    return p


def main():
    parser = build_parser()
    args, _ = parser.parse_known_args()

    if args.worker:
        return run_worker(args.worker, args.input, args.worker_model)

    if not args.input:
        parser.error("an input file is required")
    return orchestrate(args)


if __name__ == "__main__":
    sys.exit(main())