"""
BrainGPT pipeline core (demo build).
Adds: selectable voice, debug log file, saved transcript, end-of-run report.
"""

import subprocess, json, os, time, base64, tempfile, requests, threading
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import anthropic
from faster_whisper import WhisperModel

# Parallelism (scheduling only — does not change output, just speed)
CORRECT_WORKERS = 4   # concurrent Claude correction batches
TTS_WORKERS     = 3   # concurrent MiniMax voice requests
FRAME_WORKERS   = 4   # concurrent screenshot extractions

os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")

WHISPER_REPOS = {
    "base":            ("Systran/faster-whisper-base",                 "faster-whisper-base"),
    "small":           ("Systran/faster-whisper-small",                "faster-whisper-small"),
    "large-v3-turbo":  ("mobiuslabsgmbh/faster-whisper-large-v3-turbo","faster-whisper-large-v3-turbo"),
}


def _model_dir(name: str) -> Path:
    base = Path(os.environ.get("LOCALAPPDATA", str(Path.home())))
    return base / "BrainGPT" / "models" / name

def _existing_model(folder: str):
    """Reuse a model already downloaded by any BrainGPT app, to avoid re-downloading."""
    base = Path(os.environ.get("LOCALAPPDATA", str(Path.home())))
    for root in ("BrainGPT", "YouTubePipeline"):
        cand = base / root / "models" / folder
        if (cand / "model.bin").exists():
            return cand
    return None

def ensure_model(model_key="base", log=None) -> str:
    repo, folder = WHISPER_REPOS.get(model_key, WHISPER_REPOS["base"])

    found = _existing_model(folder)
    if found:
        return str(found)

    md = _model_dir(folder)
    if (md / "model.bin").exists():
        return str(md)
    if log:
        log("Downloading AI model", f"Downloading {model_key} model (one time)…")
    md.parent.mkdir(parents=True, exist_ok=True)

    # Be patient with slow/unstable connections
    os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
    os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "30")
    from huggingface_hub import snapshot_download

    # Try the official endpoint first, then a public mirror if it keeps timing out.
    # snapshot_download resumes partial files, so retries don't restart from zero.
    endpoints = [None, None, "https://hf-mirror.com", "https://hf-mirror.com"]
    last_err = None
    for attempt, endpoint in enumerate(endpoints):
        try:
            if endpoint:
                os.environ["HF_ENDPOINT"] = endpoint
                if log:
                    log("Downloading AI model", "Primary server slow — trying mirror…")
            snapshot_download(repo_id=repo, local_dir=str(md), max_workers=2)
            os.environ.pop("HF_ENDPOINT", None)
            return str(md)
        except Exception as e:
            last_err = e
            if log:
                log("Downloading AI model",
                    f"Network issue — retrying ({attempt + 1}/{len(endpoints)})…")
            time.sleep(5 * (attempt + 1))
    os.environ.pop("HF_ENDPOINT", None)

    raise RuntimeError(
        "Couldn't download the AI voice-to-text model — the servers could not be "
        "reached.\n\n"
        "• Check your internet connection.\n"
        "• If you use a VPN, firewall, or are on a restricted network, that may be "
        "blocking the download — try another network.\n"
        "• Then run it again (it resumes where it left off).\n\n"
        f"Details: {last_err}")


def _run(cmd, **kw):
    return subprocess.run(cmd, check=True, capture_output=True, **kw)

def get_duration(path: str) -> float:
    r = subprocess.run(["ffprobe", "-v", "quiet", "-print_format", "json",
                        "-show_streams", path], capture_output=True, text=True, check=True)
    for s in json.loads(r.stdout)["streams"]:
        if "duration" in s:
            return float(s["duration"])
    return 0.0

def make_thumbnail(video_path: str, out_path: str, at: float = 3.0):
    try:
        _run(["ffmpeg", "-ss", str(at), "-i", video_path, "-frames:v", "1",
              "-vf", "scale=320:-1", out_path, "-y"])
        return out_path
    except Exception:
        return None


# ── Step 1 ────────────────────────────────────────────────────────────────────

def extract_audio(video_path, audio_path):
    _run(["ffmpeg", "-i", video_path, "-vn", "-acodec", "pcm_s16le",
          "-ar", "16000", "-ac", "1", audio_path, "-y"])


# ── Step 2 ────────────────────────────────────────────────────────────────────

def _fmt(sec):
    m, s = divmod(int(sec), 60)
    return f"{m}:{s:02d}"

def _pick_device():
    """Use the GPU if an NVIDIA card is available (same model, same accuracy, much faster)."""
    try:
        import ctranslate2
        if ctranslate2.get_cuda_device_count() > 0:
            return "cuda", "float16"
    except Exception:
        pass
    return "cpu", "int8"

def _encode_video(video_path, vf_filter, out_path, log=None):
    """Render the speed-adjusted video. Uses the GPU's NVENC encoder when available
    (much faster on long videos), and falls back to CPU x264 anywhere else."""
    base = ["ffmpeg", "-i", video_path, "-filter_complex", vf_filter, "-map", "[vout]"]
    device, _ = _pick_device()
    if device == "cuda":
        try:
            if log:
                log("Building video", "Matching video speed to voice (GPU encode)…")
            subprocess.run(
                base + ["-c:v", "h264_nvenc", "-preset", "p5",
                        "-rc", "vbr", "-cq", "23", "-pix_fmt", "yuv420p",
                        out_path, "-y"],
                check=True, capture_output=True)
            return
        except Exception:
            if log:
                log("Building video", "GPU encoder unavailable — using CPU…")
    subprocess.run(
        base + ["-c:v", "libx264", "-preset", "fast", out_path, "-y"],
        check=True)

def transcribe(audio_path, log, model_key="base"):
    log("Transcribing", "Loading model…")
    path = ensure_model(model_key, log)
    device, compute = _pick_device()
    try:
        model = WhisperModel(path, device=device, compute_type=compute)
    except Exception:
        device, compute = "cpu", "int8"          # GPU present but unusable → fall back
        model = WhisperModel(path, device=device, compute_type=compute)
    log("Transcribing", f"Listening to audio (using {device.upper()})…")
    gen, info = model.transcribe(audio_path, beam_size=5)
    total = info.duration or 0
    segs = []
    for i, s in enumerate(gen):
        segs.append({"id": i, "start": s.start, "end": s.end, "text": s.text.strip()})
        if total:
            pct = min(100, int(s.end / total * 100))
            log("Transcribing", f"{pct}%  ({_fmt(s.end)} / {_fmt(total)})")
    return segs


# ── Step 3 ────────────────────────────────────────────────────────────────────

CORRECTION_RULES = (
    "You are correcting a screen-recording tutorial transcript. "
    "Use the screenshots to understand the topic before correcting.\n\n"
    "Rules:\n"
    "1. Fix spelling and grammar.\n"
    "2. Use the screenshots to correct misheard on-screen names, buttons, or terms.\n"
    "3. If a sentence makes no sense for the video, rewrite it using the screenshots as a guide.\n"
    "4. If a segment repeats the previous one, rewrite it to continue naturally.\n"
    "5. Keep the speaker's natural voice. Do not add new content.\n\n"
    "Return ONLY a JSON array with the SAME ids: [{\"id\":0,\"text\":\"...\"},...]. "
    "No markdown, no explanation."
)

def grab_frame(video_path, ts, out):
    _run(["ffmpeg", "-ss", str(ts), "-i", video_path, "-frames:v", "1",
          "-vf", "scale=768:-1", "-q:v", "5", out, "-y"])

def _extract_json_array(raw):
    raw = raw.strip()
    if raw.startswith("```"):
        parts = raw.split("```")
        if len(parts) >= 2:
            raw = parts[1]
        if raw.lstrip().lower().startswith("json"):
            raw = raw.lstrip()[4:]
        raw = raw.strip()
    a, b = raw.find("["), raw.rfind("]")
    if a != -1 and b != -1 and b > a:
        raw = raw[a:b + 1]
    return json.loads(raw)

def correct_transcript(segments, video_path, tmpdir, log, api_key, report):
    client = anthropic.Anthropic(api_key=api_key)
    n = len(segments)

    # 1) Extract every screenshot in parallel (same frames → same accuracy)
    log("Correcting transcript", "Capturing screenshots…")
    def _frame(seg):
        mid = (seg["start"] + seg["end"]) / 2
        fp = os.path.join(tmpdir, f"frame_{seg['id']}.jpg")
        try:
            grab_frame(video_path, mid, fp)
            return seg["id"], fp
        except Exception:
            return seg["id"], None
    frames = {}
    with ThreadPoolExecutor(max_workers=FRAME_WORKERS) as ex:
        for sid, fp in ex.map(_frame, segments):
            frames[sid] = fp

    # 2) Correct independent batches concurrently (same prompts/model → same accuracy)
    BATCH = 20
    batches = [segments[i:i + BATCH] for i in range(0, n, BATCH)]
    done = {"c": 0}
    lock = threading.Lock()

    def _do_batch(batch):
        content = [{"type": "text", "text": CORRECTION_RULES}]
        for seg in batch:
            fp = frames.get(seg["id"])
            if fp and os.path.exists(fp):
                with open(fp, "rb") as f:
                    b64 = base64.standard_b64encode(f.read()).decode()
                content.append({"type": "image", "source": {"type": "base64",
                                "media_type": "image/jpeg", "data": b64}})
            content.append({"type": "text", "text": f"Segment id {seg['id']}: {seg['text']}"})

        cmap = {}
        for attempt in range(2):
            try:
                msg = client.messages.create(model="claude-sonnet-4-6", max_tokens=4096,
                    messages=[{"role": "user", "content": content}])
                arr = _extract_json_array(msg.content[0].text)
                cmap = {it["id"]: it["text"] for it in arr
                        if isinstance(it, dict) and "id" in it and "text" in it}
                if cmap:
                    break
            except Exception:
                if attempt == 0:
                    time.sleep(2)
                    continue
        with lock:
            done["c"] += len(batch)
            log("Correcting transcript", f"{min(done['c'], n)} / {n} segments…")
        return batch, cmap

    failed = 0
    with ThreadPoolExecutor(max_workers=CORRECT_WORKERS) as ex:
        for batch, cmap in ex.map(_do_batch, batches):
            if not cmap:
                failed += len(batch)
            for seg in batch:
                seg["corrected"] = cmap.get(seg["id"], seg["text"])

    if failed:
        report["skipped"].append(f"Grammar correction skipped for {failed} segment(s) — kept original text")
    report["corrected_ok"] = n - failed
    return segments


# ── Steps 4-6 ─────────────────────────────────────────────────────────────────

def generate_tts(text, path, mm_key, mm_group, voice_id, model, speed):
    url = f"https://api.minimaxi.chat/v1/t2a_v2?GroupId={mm_group}"
    headers = {"Authorization": f"Bearer {mm_key}", "Content-Type": "application/json"}
    body = {"model": model, "text": text, "stream": False,
            "voice_setting": {"voice_id": voice_id, "speed": speed, "vol": 1.0, "pitch": 0},
            "audio_setting": {"sample_rate": 32000, "bitrate": 128000, "format": "mp3"}}
    last_err = None
    for attempt in range(8):
        try:
            data = requests.post(url, headers=headers, json=body, timeout=60).json()
        except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
            last_err = e
            time.sleep(8 * (attempt + 1))
            continue
        hex_audio = data.get("data", {}).get("audio")
        if hex_audio:
            with open(path, "wb") as f:
                f.write(bytes.fromhex(hex_audio))
            return True
        if data.get("base_resp", {}).get("status_code") == 1002:
            time.sleep(10 * (attempt + 1))
        else:
            raise ValueError(f"MiniMax error: {data}")
    raise ValueError(f"MiniMax failed. Last network error: {last_err}")


def build_output(segments, video_path, video_duration, tmpdir, output_path,
                 log, mm_key, mm_group, voice, report):

    before = len(segments)
    segments = [s for s in segments if s.get("corrected", "").strip()]
    empty_skipped = before - len(segments)
    if empty_skipped:
        report["skipped"].append(f"Removed {empty_skipped} silent/empty segment(s)")

    total = len(segments)
    done = {"c": 0}
    lock = threading.Lock()

    def _voice(item):
        i, seg = item
        tts_path = os.path.join(tmpdir, f"tts_{i}.mp3")
        ok = False
        try:
            generate_tts(seg["corrected"], tts_path, mm_key, mm_group,
                         voice["voice_id"], voice.get("model", "speech-02-hd"),
                         float(voice.get("speed", 1.0)))
            dur = get_duration(tts_path)
            if dur >= 0.05:
                seg["tts_path"] = tts_path
                seg["tts_dur"] = dur
                ok = True
        except Exception:
            ok = False
        if not ok:
            seg["skip"] = True
        with lock:
            done["c"] += 1
            log("Generating voice", f"{min(done['c'], total)} / {total} clips…")
        return ok

    # Generate voice clips concurrently (no fixed delay; generate_tts self-throttles
    # on rate limits). Output audio is identical — only the scheduling changes.
    with ThreadPoolExecutor(max_workers=TTS_WORKERS) as ex:
        results = list(ex.map(_voice, list(enumerate(segments))))
    tts_failed = sum(1 for r in results if not r)

    segments = [s for s in segments if not s.get("skip")]
    if tts_failed:
        report["skipped"].append(f"{tts_failed} voice clip(s) failed and were skipped")

    # Timeline
    chunks, orig_pos, new_pos = [], 0.0, 0.0
    for seg in segments:
        gap = seg["start"] - orig_pos
        if gap > 0.01:
            chunks.append({"os": orig_pos, "od": gap, "ns": new_pos, "nd": gap, "tts": None})
            new_pos += gap
        od = seg["end"] - seg["start"]
        chunks.append({"os": seg["start"], "od": od, "ns": new_pos, "nd": seg["tts_dur"],
                       "tts": seg["tts_path"]})
        new_pos += seg["tts_dur"]
        orig_pos = seg["end"]
    tail = video_duration - orig_pos
    if tail > 0.01:
        chunks.append({"os": orig_pos, "od": tail, "ns": new_pos, "nd": tail, "tts": None})
        new_pos += tail
    total_duration = new_pos

    # Speed-adjust video
    log("Building video", "Matching video speed to voice…")
    vf, vl = [], []
    for i, c in enumerate(chunks):
        factor = c["nd"] / c["od"]
        vf.append(f"[0:v]trim=start={c['os']:.4f}:duration={c['od']:.4f},"
                  f"setpts={factor:.6f}*(PTS-STARTPTS)[v{i}]")
        vl.append(f"[v{i}]")
    vf.append(f"{''.join(vl)}concat=n={len(chunks)}:v=1:a=0[vout]")
    processed = os.path.join(tmpdir, "video.mp4")
    _encode_video(video_path, ";".join(vf), processed, log)

    # Audio
    log("Building audio", "Mixing voice track…")
    tts_chunks = [(c["tts"], c["ns"]) for c in chunks if c["tts"]]
    ai = ["-f", "lavfi", "-i", "anullsrc=r=32000:cl=mono"]
    af = []
    for idx, (tp, st) in enumerate(tts_chunks):
        ai += ["-i", tp]
        ms = int(st * 1000)
        af.append(f"[{idx+1}]adelay={ms}|{ms}[a{idx}]")
    all_a = "[0]" + "".join(f"[a{i}]" for i in range(len(tts_chunks)))
    af.append(f"{all_a}amix=inputs={1+len(tts_chunks)}:normalize=0[aout]")
    final_audio = os.path.join(tmpdir, "audio.mp3")
    subprocess.run(["ffmpeg"] + ai + ["-filter_complex", ";".join(af), "-map", "[aout]",
                    "-t", str(total_duration), final_audio, "-y"], check=True)

    # Merge
    log("Finalizing", "Merging video and audio…")
    subprocess.run(["ffmpeg", "-i", processed, "-i", final_audio, "-c:v", "copy",
                    "-map", "0:v:0", "-map", "1:a:0", "-shortest", output_path, "-y"],
                   check=True)

    report["voiced_segments"] = total - tts_failed
    return segments


# ── Public entry ──────────────────────────────────────────────────────────────

def process_video(video_path, output_path, keys, voice, settings,
                  progress=None, logfile=None) -> dict:
    """Returns a report dict. progress(step, detail) for UI updates."""
    logf = open(logfile, "w", encoding="utf-8") if logfile else None

    def log(step, detail=""):
        line = f"[{time.strftime('%H:%M:%S')}] {step}: {detail}"
        if logf:
            logf.write(line + "\n"); logf.flush()
        if progress:
            progress(step, detail)

    report = {"skipped": [], "stages": {}, "total_segments": 0,
              "corrected_ok": 0, "voiced_segments": 0,
              "transcript_path": None, "output": output_path}

    api_key  = keys.get("ANTHROPIC_API_KEY", "")
    mm_key   = keys.get("MINIMAX_API_KEY", "")
    mm_group = keys.get("MINIMAX_GROUP_ID", "")
    for nm, val in [("ANTHROPIC_API_KEY", api_key), ("MINIMAX_API_KEY", mm_key),
                    ("MINIMAX_GROUP_ID", mm_group)]:
        if not val:
            raise ValueError(f"Missing API key: {nm}")
    if not voice or not voice.get("voice_id"):
        raise ValueError("No voice selected")

    try:
        with tempfile.TemporaryDirectory() as tmp:
            t0 = time.time()
            log("Extracting audio", "")
            audio = os.path.join(tmp, "audio.wav")
            extract_audio(video_path, audio)
            video_dur = get_duration(video_path)
            report["stages"]["extract"] = round(time.time() - t0, 1)

            t1 = time.time()
            segs = transcribe(audio, log, settings.get("whisper_model", "base"))
            report["total_segments"] = len(segs)
            report["stages"]["transcribe"] = round(time.time() - t1, 1)
            if not segs:
                report["skipped"].append("No speech detected — nothing to voice")
                raise ValueError("No speech detected in this video")

            t2 = time.time()
            log("Correcting transcript", "")
            segs = correct_transcript(segs, video_path, tmp, log, api_key, report)
            report["stages"]["correct"] = round(time.time() - t2, 1)

            # Save corrected transcript next to output
            t_path = str(Path(output_path).with_suffix("")) + "_transcript.txt"
            try:
                with open(t_path, "w", encoding="utf-8") as tf:
                    tf.write("\n".join(s.get("corrected", s["text"]) for s in segs))
                report["transcript_path"] = t_path
            except Exception:
                pass

            t3 = time.time()
            build_output(segs, video_path, video_dur, tmp, output_path, log,
                         mm_key, mm_group, voice, report)
            report["stages"]["render"] = round(time.time() - t3, 1)

        log("Done", "Complete")
    finally:
        if logf:
            logf.close()
    return report