from __future__ import annotations

import glob
import json
import os
import re
import shutil
import subprocess
import sys
import time
from typing import Any

from langchain_core.tools import tool

from persistence import CLI_VENVS_DIR, MUGI_LOGS_DIR
from tools._logging import _debug_log, _emit_input_status, should_stop_tool
from tools.cli.runner import cli_rich_path
from tools.job_stream_schema import (
    close_job_stream,
    error_job_stream,
    make_capabilities,
    metric_def,
    open_job_stream,
    update_job_stream,
)
from tools.live_view_visuals import emit_live_view_preview, image_to_preview_base64
from tools.plugin_stream_controls import clear_stream_state, get_action_count, trigger_action
from utils.path_safety import get_user_output_dir, is_plugins_dir_write_blocked

_CTX = None
_PLUGIN_ID = "yt_dlp"
_PLUGIN_NAME = "YT-DLP"

# Seconds to wait for yt-dlp to exit after stdout closes (see _run_ytdlp_streaming).
_YTDLP_PROC_JOIN_TIMEOUT_S = 5
_MUGI_HOMEBREW = os.path.join(os.path.expanduser("~"), ".mugi", "homebrew")
_MUGI_VENV = os.path.join(os.path.expanduser("~"), ".mugi", "venv")
_YTDLP_PROGRESS_RE = re.compile(r"(\d+(?:\.\d+)?)%")
_YTDLP_SPEED_RE = re.compile(r"\bat\s+([0-9.]+\s*[KMG]?i?B/s)\b", re.IGNORECASE)
_YTDLP_ETA_RE = re.compile(r"\bETA\s+([0-9:]+)\b", re.IGNORECASE)


def _ytdlp_via_python(python_exe: str) -> list[str] | None:
    if not python_exe or not os.path.isfile(python_exe):
        return None
    try:
        env = {**os.environ, "PATH": cli_rich_path()}
        env.pop("PYTHONHOME", None)
        env.pop("PYTHONPATH", None)
        r = subprocess.run(
            [python_exe, "-m", "yt_dlp", "--version"],
            capture_output=True,
            text=True,
            timeout=10,
            env=env,
        )
        if r.returncode == 0:
            return [python_exe, "-m", "yt_dlp"]
    except Exception:
        pass
    return None


def _find_ytdlp() -> tuple[str | list[str] | None, str]:
    mugi_bin = os.path.join(_MUGI_HOMEBREW, "bin", "yt-dlp")
    if os.path.isfile(mugi_bin):
        return mugi_bin, ""
    cellar_glob = os.path.join(_MUGI_HOMEBREW, "Cellar", "yt-dlp", "*", "bin", "yt-dlp")
    for path in sorted(glob.glob(cellar_glob), reverse=True):
        if os.path.isfile(path):
            return path, ""

    cli_venv_bin = os.path.join(CLI_VENVS_DIR, "yt-dlp", "bin", "yt-dlp")
    if os.path.isfile(cli_venv_bin):
        return cli_venv_bin, ""
    cli_venv_bin_alt = os.path.join(CLI_VENVS_DIR, "yt_dlp", "bin", "yt-dlp")
    if os.path.isfile(cli_venv_bin_alt):
        return cli_venv_bin_alt, ""

    path = shutil.which("yt-dlp", path=cli_rich_path())
    if path:
        return path, ""

    venv_python = os.path.join(_MUGI_VENV, "bin", "python") if os.name != "nt" else os.path.join(_MUGI_VENV, "Scripts", "python.exe")
    for candidate in [venv_python, sys.executable]:
        args = _ytdlp_via_python(candidate)
        if args:
            _debug_log("yt-dlp: using python -m yt_dlp (no brew/venv binary found)", "INFO")
            return args, ""

    return (
        None,
        "yt-dlp not found. Install via:\n"
        "  - Homebrew: ensure_mugi_homebrew then mugi_brew_install('yt-dlp')\n"
        "  - Pip (isolated venv): install_cli_package('yt-dlp')\n"
        "Then run this tool again.",
    )


def _ytdlp_format_cli_args(format_preference: str) -> list[str]:
    """Map model-facing format hints to yt-dlp CLI args.

    Passing ``-f mp4`` on YouTube often selects a fragile HLS-only format (HTTP 403 on fragments).
    For MP4 container output, use merge/remux instead and keep the default best video+audio selector.
    """
    raw = (format_preference or "").strip()
    if not raw:
        return []
    low = raw.lower().replace(".", "").replace("/", "")
    if low in ("mp4", "m4v", "muxedmp4", "videomp4"):
        return ["--merge-output-format", "mp4"]
    return ["-f", raw]


def _parse_ytdlp_progress(line: str) -> dict[str, Any]:
    pct = None
    m = _YTDLP_PROGRESS_RE.search(line or "")
    if m:
        try:
            pct = float(m.group(1))
        except Exception:
            pct = None
    speed = ""
    s = _YTDLP_SPEED_RE.search(line or "")
    if s:
        speed = s.group(1).strip()
    eta = ""
    e = _YTDLP_ETA_RE.search(line or "")
    if e:
        eta = e.group(1).strip()
    return {"progress_pct": pct, "speed": speed, "eta": eta}


def _run_ytdlp_streaming(*, args: list[str], env: dict[str, str], stream_id: str, phase: str) -> tuple[int, str]:
    proc = subprocess.Popen(
        args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
        env=env,
        bufsize=1,
    )
    lines: list[str] = []
    last_pct = -1.0
    last_emit = 0.0
    try:
        while True:
            if should_stop_tool() or get_action_count(stream_id, "stop_download") > 0:
                try:
                    proc.terminate()
                except Exception:
                    pass
                update_job_stream(
                    plugin_id=_PLUGIN_ID,
                    plugin_name=_PLUGIN_NAME,
                    stream_id=stream_id,
                    status="stopped",
                    phase=phase,
                    message="Stop requested by user, terminating yt-dlp...",
                )
                break
            line = proc.stdout.readline() if proc.stdout else ""
            if not line:
                if proc.poll() is not None:
                    break
                time.sleep(0.05)
                continue
            txt = line.rstrip("\n")
            lines.append(txt)
            parsed = _parse_ytdlp_progress(txt)
            pct = parsed.get("progress_pct")
            now = time.time()
            should_emit = False
            if isinstance(pct, (int, float)):
                if pct >= last_pct + 1.0:
                    should_emit = True
                    last_pct = float(pct)
            if now - last_emit >= 1.2:
                should_emit = True
            if should_emit:
                metrics = {}
                if isinstance(pct, (int, float)):
                    metrics["download_progress_pct"] = float(pct)
                extra = {"log_line": txt[-400:]}
                if parsed.get("speed"):
                    extra["speed"] = parsed["speed"]
                if parsed.get("eta"):
                    extra["eta"] = parsed["eta"]
                update_job_stream(
                    plugin_id=_PLUGIN_ID,
                    plugin_name=_PLUGIN_NAME,
                    stream_id=stream_id,
                    status="running",
                    phase=phase,
                    progress_pct=float(pct) if isinstance(pct, (int, float)) else None,
                    message=txt[:160] or "Downloading...",
                    metrics=metrics,
                    extra=extra,
                )
                last_emit = now
        rc = proc.wait(timeout=_YTDLP_PROC_JOIN_TIMEOUT_S)
        return rc, "\n".join(lines).strip()
    finally:
        if proc.stdout:
            try:
                proc.stdout.close()
            except Exception:
                pass


@tool
def download_video_yt_dlp(
    url: str,
    output_path: str = "",
    audio_only: bool = False,
    format_preference: str = "",
) -> str:
    """Download a video (or audio only) using yt-dlp with progress and stop controls.

    Leave format_preference empty for best quality. For an MP4 container, pass format_preference="mp4"
    (uses remux; do not pass raw yt-dlp selectors unless you know them). To re-encode or convert an
    existing download, use convert_video instead of re-downloading.
    """
    global _CTX
    if _CTX is None:
        return "Plugin context is not initialized."
    url = re.sub(r"\s+", "", url or "")
    stream_id = f"{_PLUGIN_ID}:{int(time.time() * 1000)}"
    _emit_input_status("Downloading video…")
    _debug_log(f"yt-dlp: downloading from {url[:80]}{'...' if len(url) > 80 else ''}", "INFO")
    if not url:
        return "No URL provided."

    controls = [
        {
            "id": "stop_download",
            "type": "button",
            "label": "Stop Download",
            "action_id": "stop_download",
            "style": "destructive",
        }
    ]
    open_job_stream(
        plugin_id=_PLUGIN_ID,
        plugin_name=_PLUGIN_NAME,
        title="yt-dlp Download",
        job_type="ytdlp_download",
        stream_id=stream_id,
        controls=controls,
        capabilities=make_capabilities(
            cancel=True,
            progress=True,
            metrics=True,
            retry=True,
            phases=True,
            stream_logs=True,
            eta=True,
            artifacts=True,
        ),
        metric_schema=[
            metric_def("download_progress_pct", "Download Progress", unit="%", format="percent", priority=1),
        ],
        metadata={"url": url, "audio_only": bool(audio_only), "format_preference": format_preference or ""},
    )
    _CTX.register_stream_controls(stream_id, controls)

    ytdlp_path, err = _find_ytdlp()
    if err:
        error_job_stream(
            plugin_id=_PLUGIN_ID,
            plugin_name=_PLUGIN_NAME,
            stream_id=stream_id,
            phase="resolve_binary",
            message=err,
        )
        return err

    if output_path and output_path.strip():
        out_abs, out_err = _CTX.resolve_safe_path(output_path.strip(), allow_write=True)
        if out_err:
            error_job_stream(
                plugin_id=_PLUGIN_ID,
                plugin_name=_PLUGIN_NAME,
                stream_id=stream_id,
                phase="resolve_output",
                message=out_err,
            )
            return out_err
        if not out_abs:
            error_job_stream(
                plugin_id=_PLUGIN_ID,
                plugin_name=_PLUGIN_NAME,
                stream_id=stream_id,
                phase="resolve_output",
                message="Invalid output_path.",
            )
            return "Invalid output_path."
        if is_plugins_dir_write_blocked(out_abs):
            error_job_stream(
                plugin_id=_PLUGIN_ID,
                plugin_name=_PLUGIN_NAME,
                stream_id=stream_id,
                phase="resolve_output",
                message="Blocked: cannot write to plugins directory.",
            )
            return "Blocked: cannot write to the plugins directory. Use create_plugin (evolution mode) to add plugins."
        out_template = os.path.join(out_abs, "%(title)s.%(ext)s") if os.path.isdir(out_abs) else out_abs
    else:
        out_dir = get_user_output_dir()
        os.makedirs(out_dir, exist_ok=True)
        out_template = os.path.join(out_dir, "%(title)s.%(ext)s")

    env = os.environ.copy()
    env["PATH"] = cli_rich_path()
    env.pop("PYTHONHOME", None)
    env.pop("PYTHONPATH", None)
    args = (list(ytdlp_path) if isinstance(ytdlp_path, list) else [ytdlp_path]) + ["-o", out_template]
    if audio_only:
        args.extend(["-x", "--audio-format", "mp3"])
    args.extend(_ytdlp_format_cli_args(format_preference))
    args.append(url)

    emit_live_view_preview("Downloading…", tool_name="download_video_yt_dlp", progress_pct=0)
    try:
        thumb_args = (list(ytdlp_path) if isinstance(ytdlp_path, list) else [ytdlp_path]) + ["--get-thumbnail", url]
        thumb_proc = subprocess.run(thumb_args, capture_output=True, text=True, timeout=15, env=env)
        thumb_url = (thumb_proc.stdout or "").strip().splitlines()[0].strip() if thumb_proc.returncode == 0 else ""
        if thumb_url and thumb_url.startswith("http"):
            import requests as _req

            thumb_resp = _req.get(thumb_url, timeout=5)
            if thumb_resp.status_code == 200 and len(thumb_resp.content) > 500:
                b64 = image_to_preview_base64(thumb_resp.content, max_width=400)
                if b64:
                    emit_live_view_preview(
                        f"Downloading: {url[:60]}",
                        tool_name="download_video_yt_dlp",
                        progress_pct=0,
                        preview_image_base64=b64,
                        url=url,
                    )
    except Exception:
        pass

    try:
        update_job_stream(
            plugin_id=_PLUGIN_ID,
            plugin_name=_PLUGIN_NAME,
            stream_id=stream_id,
            phase="download",
            status="running",
            progress_pct=0,
            message="Starting yt-dlp download...",
        )
        rc, merged_out = _run_ytdlp_streaming(args=args, env=env, stream_id=stream_id, phase="download")
        out = (merged_out or "").strip()
        if rc != 0:
            msg = out or f"yt-dlp exited with code {rc}"
            stderr_lower = out.lower()
            _want_cookie_retry = (
                "403" in stderr_lower
                or "sign in to confirm" in stderr_lower
                or "not a bot" in stderr_lower
                or "cookies" in stderr_lower
                or "operation not permitted" in stderr_lower
                or "cookies.binarycookies" in stderr_lower
            )
            if _want_cookie_retry:
                update_job_stream(
                    plugin_id=_PLUGIN_ID,
                    plugin_name=_PLUGIN_NAME,
                    stream_id=stream_id,
                    phase="retry_decision",
                    status="running",
                    message="Auth/cookies issue detected. Waiting for user retry choice...",
                )
                choice = _CTX.request_user_decision(
                    "browser_cookie_source",
                    context={
                        "title": "Retry with browser cookies?",
                        "message": (
                            "Download failed due to cookie/auth restrictions. Pick the browser where you are "
                            "signed in to the site. Prefer Chrome or Brave if you use them—Safari’s cookie file "
                            "is often blocked unless Mugi has Full Disk Access (System Settings → Privacy & "
                            "Security → Full Disk Access)."
                        ),
                        "url": url,
                        "out_template": out_template,
                        "format_preference": format_preference or "",
                        "audio_only": bool(audio_only),
                    },
                    timeout_s=300.0,
                )
                if choice and choice != "cancel":
                    args_retry = (list(ytdlp_path) if isinstance(ytdlp_path, list) else [ytdlp_path]) + [
                        "-o",
                        out_template,
                        "--cookies-from-browser",
                        choice,
                    ]
                    if audio_only:
                        args_retry.extend(["-x", "--audio-format", "mp3"])
                    args_retry.extend(_ytdlp_format_cli_args(format_preference))
                    args_retry.append(url)
                    update_job_stream(
                        plugin_id=_PLUGIN_ID,
                        plugin_name=_PLUGIN_NAME,
                        stream_id=stream_id,
                        phase="retry_download",
                        status="running",
                        progress_pct=0,
                        message=f"Retrying with {choice} browser cookies...",
                    )
                    rc2, out2 = _run_ytdlp_streaming(
                        args=args_retry,
                        env=env,
                        stream_id=stream_id,
                        phase="retry_download",
                    )
                    out2 = (out2 or "").strip()
                    try:
                        os.makedirs(MUGI_LOGS_DIR, exist_ok=True)
                        with open(os.path.join(MUGI_LOGS_DIR, "ytdl_debug.log"), "a", encoding="utf-8") as f:
                            f.write(
                                json.dumps(
                                    {
                                        "id": "ytdl_result",
                                        "timestamp": int(time.time() * 1000),
                                        "location": "templates/bundled_plugins/yt_dlp/plugin.py",
                                        "message": "yt-dlp retry with cookies",
                                        "data": {
                                            "url_preview": (url or "")[:200],
                                            "returncode": rc2,
                                            "browser": choice,
                                            "stderr_preview": (out2 or "")[:800],
                                        },
                                    }
                                )
                                + "\n"
                            )
                    except Exception:
                        pass
                    if rc2 == 0:
                        emit_live_view_preview("Download complete.", tool_name="download_video_yt_dlp", progress_pct=100)
                        out_dir = os.path.dirname(out_template)
                        close_job_stream(
                            plugin_id=_PLUGIN_ID,
                            plugin_name=_PLUGIN_NAME,
                            stream_id=stream_id,
                            status="completed",
                            message=f"Download complete (retry using {choice} cookies).",
                            metrics={"download_progress_pct": 100.0},
                            extra={"output_directory": out_dir, "browser_cookies": choice},
                        )
                        return f"Download complete (using {choice} cookies). Output directory: {out_dir}\n{out2 or ''}".strip()
                    msg = out2 or f"yt-dlp exited with code {rc2}"
                    msg += "\n\nRetry with browser cookies still failed."
                    if "cookies.binarycookies" in (out2 or "").lower() or (
                        "operation not permitted" in (out2 or "").lower() and "library/cookies" in (out2 or "").lower()
                    ):
                        msg += (
                            " Safari’s cookie database is protected by macOS: choose Chrome or Brave if you use "
                            "those browsers, or enable Full Disk Access for Mugi, then try again."
                        )
                    else:
                        msg += " The user can run yt-dlp manually with --cookies-from-browser or try another browser."
                    error_job_stream(
                        plugin_id=_PLUGIN_ID,
                        plugin_name=_PLUGIN_NAME,
                        stream_id=stream_id,
                        phase="retry_download",
                        message=msg,
                    )
                    return msg
                msg += "\n\nSuggestion: YouTube may be blocking the download. You can retry and choose a browser when prompted."
            error_job_stream(
                plugin_id=_PLUGIN_ID,
                plugin_name=_PLUGIN_NAME,
                stream_id=stream_id,
                phase="download",
                message=msg,
            )
            return msg

        emit_live_view_preview("Download complete.", tool_name="download_video_yt_dlp", progress_pct=100)
        out_dir = os.path.dirname(out_template)
        close_job_stream(
            plugin_id=_PLUGIN_ID,
            plugin_name=_PLUGIN_NAME,
            stream_id=stream_id,
            status="completed",
            message="Download complete.",
            metrics={"download_progress_pct": 100.0},
            extra={"output_directory": out_dir},
        )
        return f"Download complete. Output directory: {out_dir}\n{out or ''}".strip()
    except subprocess.TimeoutExpired as e:
        waited = getattr(e, "timeout", None)
        if waited is None:
            waited = _YTDLP_PROC_JOIN_TIMEOUT_S
        error_job_stream(
            plugin_id=_PLUGIN_ID,
            plugin_name=_PLUGIN_NAME,
            stream_id=stream_id,
            phase="download",
            message=f"yt-dlp did not exit within {waited} seconds (subprocess wait).",
        )
        return f"yt-dlp did not exit within {waited} seconds (subprocess wait)."
    except Exception as e:
        error_job_stream(
            plugin_id=_PLUGIN_ID,
            plugin_name=_PLUGIN_NAME,
            stream_id=stream_id,
            phase="download",
            message=f"download_video_yt_dlp failed: {e}",
        )
        return f"download_video_yt_dlp failed: {e}"
    finally:
        clear_stream_state(stream_id)


def _handle_control_event(event: dict, _ctx) -> dict:
    action_id = str(event.get("action_id") or "").strip()
    stream_id = str(event.get("stream_id") or "").strip()
    if action_id == "stop_download" and stream_id:
        trigger_action(stream_id, "stop_download")
        return {"status": "applied", "message": "Stopping download..."}
    return {"status": "applied", "message": "Control applied."}


def register(ctx):
    global _CTX, _PLUGIN_ID, _PLUGIN_NAME
    _CTX = ctx
    _PLUGIN_ID = "yt_dlp"
    _PLUGIN_NAME = "YT-DLP"
    ctx.register_control_handler(_handle_control_event)
    ctx.register_tools([download_video_yt_dlp])
