#!/usr/bin/env python3
"""Remote MVLog worker.

Scans the web machine's in-dir, renders changed jobs locally with movmaker,
and publishes generated MP4 files back to the web machine's out-dir.
"""

from __future__ import annotations

import argparse
import logging
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from shlex import quote

ROOT = Path(__file__).resolve().parent
MOVMAKER = ROOT / "movmaker.py"
INTRO_LOGO = ROOT / "img" / "moto_travel.png"
STATE_NAME = ".movmaker-state.json"
LOCK_NAME = ".movmaker-lock"
ENABLED_NAME = ".movmaker-enabled"
PREVIEW_NAME = ".movmaker-preview"
HIDDEN_NAME = ".mvlog-hidden"
PERMALINK_NAME = ".mvlog-permalink"
ARTICLE_ID_NAME = ".mvlog-id"
INTRO_LOGO_ALPHA = "0.5"
MVLOG_WORKER_RECOMMENDED_INTERVAL_MINUTES = 5


def now() -> str:
    return datetime.now(timezone.utc).isoformat()


def run(cmd: list[str], *, input_text: str | None = None, check: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(cmd, input=input_text, text=True, capture_output=True, check=check)


def ssh(host: str, command: str, *, check: bool = True, input_text: str | None = None) -> subprocess.CompletedProcess[str]:
    return run(["ssh", host, command], input_text=input_text, check=check)


def remote_q(path: str) -> str:
    return quote(path)


def safe_chmod(path: Path, mode: int) -> None:
    try:
        path.chmod(mode)
    except PermissionError:
        # The worker may have group write access to www-data-owned files but
        # still cannot chmod files it does not own. Do not fail rendering for that.
        pass
    except FileNotFoundError:
        pass


def local_dir_accessible(path: Path) -> bool:
    return path.is_dir() and os.access(path, os.R_OK | os.X_OK)


def list_jobs(host: str | None, remote_root: str) -> list[str]:
    if host is None:
        root = Path(remote_root) / "in-dir"
        if not root.is_dir():
            return []
        jobs = []
        for path in root.iterdir():
            if path.name.startswith('.') or not path.is_dir():
                continue
            if local_dir_accessible(path):
                jobs.append(path.name)
            else:
                print(f"skip inaccessible input dir: {path.name}")
        return sorted(jobs)
    cmd = f"cd {remote_q(remote_root)} && find in-dir -mindepth 1 -maxdepth 1 -type d ! -name '.*' -readable -executable -printf '%f\\n' | sort"
    out = ssh(host, cmd).stdout
    return [line.strip() for line in out.splitlines() if line.strip()]


def remote_file_exists(host: str | None, path: str) -> bool:
    if host is None:
        return Path(path).is_file()
    return ssh(host, f"test -f {remote_q(path)}", check=False).returncode == 0


def read_remote_text(host: str | None, path: str) -> str:
    if host is None:
        try:
            return Path(path).read_text(encoding="utf-8")
        except Exception:
            return ""
    res = ssh(host, f"cat {remote_q(path)} 2>/dev/null", check=False)
    return res.stdout if res.returncode == 0 else ""


def read_article_id(host: str | None, job_dir: str) -> str | None:
    raw = read_remote_text(host, f"{job_dir}/{ARTICLE_ID_NAME}").strip().lower()
    if re.fullmatch(r"\d{14}[a-f0-9]{16}", raw):
        return raw
    return None


def acquire_lock(host: str | None, job_dir: str) -> bool:
    lock = f"{job_dir}/{LOCK_NAME}"
    if host is None:
        try:
            Path(lock).mkdir()
            Path(lock, "started_at").write_text(now() + "\n", encoding="utf-8")
            return True
        except FileExistsError:
            return False
    cmd = f"mkdir {remote_q(lock)} 2>/dev/null && date -Iseconds > {remote_q(lock + '/started_at')}"
    return ssh(host, cmd, check=False).returncode == 0


def release_lock(host: str | None, job_dir: str) -> None:
    if host is None:
        shutil.rmtree(Path(job_dir) / LOCK_NAME, ignore_errors=True)
        return
    ssh(host, f"rm -rf {remote_q(job_dir + '/' + LOCK_NAME)}", check=False)


def remote_fingerprint(host: str | None, job_dir: str) -> str:
    if host is None:
        entries = []
        root = Path(job_dir)
        for dirpath, dirnames, filenames in os.walk(root):
            dirnames[:] = [d for d in dirnames if d != LOCK_NAME]
            for name in filenames:
                if name in {STATE_NAME, ENABLED_NAME, PREVIEW_NAME, HIDDEN_NAME, PERMALINK_NAME, ARTICLE_ID_NAME, '.mvlog-hide-map'}:
                    continue
                path = Path(dirpath) / name
                rel = path.relative_to(root).as_posix()
                st = path.stat()
                entries.append([rel, st.st_size, st.st_mtime_ns])
        entries.sort()
        return hashlib.sha256(json.dumps(entries, ensure_ascii=False, separators=(",", ":")).encode("utf-8")).hexdigest()
    script = r'''
import hashlib, json, os, sys
root = sys.argv[1]
entries = []
for dirpath, dirnames, filenames in os.walk(root):
    dirnames[:] = [d for d in dirnames if d != '.movmaker-lock']
    for name in filenames:
        if name in {'.movmaker-state.json', '.movmaker-enabled', '.movmaker-preview', '.mvlog-hidden', '.mvlog-permalink', '.mvlog-id', '.mvlog-hide-map'}:
            continue
        path = os.path.join(dirpath, name)
        rel = os.path.relpath(path, root)
        st = os.stat(path)
        entries.append([rel, st.st_size, st.st_mtime_ns])
entries.sort()
print(hashlib.sha256(json.dumps(entries, ensure_ascii=False, separators=(',', ':')).encode('utf-8')).hexdigest())
'''
    cmd = f"python3 -c {quote(script)} {remote_q(job_dir)}"
    return ssh(host, cmd).stdout.strip()


def read_remote_json(host: str | None, path: str) -> dict:
    if host is None:
        try:
            data = json.loads(Path(path).read_text(encoding="utf-8"))
            return data if isinstance(data, dict) else {}
        except Exception:
            return {}
    res = ssh(host, f"cat {remote_q(path)} 2>/dev/null", check=False)
    if res.returncode != 0 or not res.stdout.strip():
        return {}
    try:
        data = json.loads(res.stdout)
        return data if isinstance(data, dict) else {}
    except json.JSONDecodeError:
        return {}


def write_remote_json(host: str, path: str, data: dict) -> None:
    tmp_path = f"{path}.tmp.$$"
    payload = json.dumps(data, ensure_ascii=False, indent=2) + "\n"
    cmd = (
        f"mkdir -p {remote_q(str(Path(path).parent))} && "
        f"cat > {remote_q(tmp_path)} && "
        f"chmod 664 {remote_q(tmp_path)} && "
        f"mv {remote_q(tmp_path)} {remote_q(path)}"
    )
    ssh(host, cmd, input_text=payload)


def read_state(host: str | None, job_dir: str) -> dict:
    if host is None:
        path = Path(job_dir) / STATE_NAME
        if not path.exists():
            return {}
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            return {}
    return read_remote_json(host, f"{job_dir}/{STATE_NAME}")
 

def write_state(host: str | None, job_dir: str, state: dict) -> None:
    if host is None:
        path = Path(job_dir) / STATE_NAME
        path.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        safe_chmod(path, 0o664)
        return
    path = f"{job_dir}/{STATE_NAME}"
    tmp_path = f"{path}.tmp.$$"
    data = json.dumps(state, ensure_ascii=False, indent=2) + "\n"
    cmd = (
        f"cat > {remote_q(tmp_path)} && "
        f"chmod 664 {remote_q(tmp_path)} && "
        f"mv {remote_q(tmp_path)} {remote_q(path)}"
    )
    ssh(host, cmd, input_text=data)
 

def copy_job_from_remote(host: str | None, job_dir: str, local_input: Path) -> None:
    local_input.mkdir(parents=True, exist_ok=True)
    excludes = {LOCK_NAME, STATE_NAME, PERMALINK_NAME}
    if host is None:
        src = Path(job_dir)
        for path in src.rglob("*"):
            if any(part in excludes for part in path.relative_to(src).parts):
                continue
            dest = local_input / path.relative_to(src)
            if path.is_dir():
                dest.mkdir(parents=True, exist_ok=True)
            elif path.is_file():
                dest.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(path, dest)
        return
    remote_cmd = (
        f"tar --exclude={quote(LOCK_NAME)} --exclude={quote(STATE_NAME)} --exclude={quote(PERMALINK_NAME)} "
        f"-C {remote_q(job_dir)} -cf - ."
    )
    p1 = subprocess.Popen(["ssh", host, remote_cmd], stdout=subprocess.PIPE)
    try:
        subprocess.run(["tar", "-xf", "-", "-C", str(local_input)], stdin=p1.stdout, check=True)
        assert p1.stdout is not None
        p1.stdout.close()
        rc = p1.wait()
        if rc != 0:
            raise subprocess.CalledProcessError(rc, ["ssh", host, remote_cmd])
    except Exception:
        p1.kill()
        raise




def publish_output(host: str | None, remote_root: str, local_mp4: Path, old_output: str | None) -> str:
    out_name = local_mp4.name
    if host is None:
        out_dir = Path(remote_root) / "out-dir"
        out_dir.mkdir(parents=True, exist_ok=True)
        final = out_dir / out_name
        tmp = out_dir / f".{out_name}.tmp"
        shutil.copy2(local_mp4, tmp)
        safe_chmod(tmp, 0o664)
        tmp.replace(final)
        safe_chmod(final, 0o664)
        (Path(remote_root) / "cache" / "videos.json").unlink(missing_ok=True)
        if old_output and old_output != out_name:
            (out_dir / old_output).unlink(missing_ok=True)
        return out_name
    remote_tmp = f"{remote_root}/out-dir/.{out_name}.tmp"
    remote_final = f"{remote_root}/out-dir/{out_name}"
    subprocess.run(["scp", str(local_mp4), f"{host}:{remote_tmp}"], check=True)
    cmds = [
        f"mv {remote_q(remote_tmp)} {remote_q(remote_final)}",
        f"rm -f {remote_q(remote_root + '/cache/videos.json')}",
    ]
    if old_output and old_output != out_name:
        cmds.append(f"rm -f {remote_q(remote_root + '/out-dir/' + old_output)}")
    ssh(host, " && ".join(cmds))
    return out_name


def upload_audio_files(host: str | None, job_dir: str, local_input: Path) -> list[str]:
    """Upload audio files found inside local_input to the remote job_dir/in-dir if they are missing.

    Returns a list of uploaded relative paths.
    """
    AUDIO_EXTS = {".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus", ".wma", ".aiff", ".aif"}
    uploaded: list[str] = []
    for path in sorted(local_input.rglob("*")):
        if not path.is_file():
            continue
        if path.suffix.lower() not in AUDIO_EXTS:
            continue
        rel = path.relative_to(local_input).as_posix()
        remote_path = f"{job_dir}/{rel}"
        if remote_file_exists(host, remote_path):
            print(f"audio exists on server, skipping: {rel}")
            continue
        parent = str(Path(remote_path).parent)
        if host is None:
            Path(parent).mkdir(parents=True, exist_ok=True)
            shutil.copy2(str(path), remote_path)
            safe_chmod(Path(remote_path), 0o664)
        else:
            # Ensure parent directory exists on the remote host
            ssh(host, f"mkdir -p {remote_q(parent)}")
            remote_tmp = f"{remote_path}.tmp.$$"
            # Upload to a temporary file then move into place for (some) atomicity
            subprocess.run(["scp", str(path), f"{host}:{remote_tmp}"], check=True)
            ssh(host, f"mv {remote_q(remote_tmp)} {remote_q(remote_path)}")
        uploaded.append(rel)
        print(f"uploaded audio: {rel}")
    return uploaded


def render_job(local_input: Path, local_out: Path, extra_args: list[str], temp_dir: Path) -> Path:
    local_out.mkdir(parents=True, exist_ok=True)
    intro_logo = INTRO_LOGO if INTRO_LOGO.exists() else None
    intro_args = []
    if intro_logo is not None:
        intro_args = [
            "--intro-logo", str(intro_logo),
            "--intro-logo-alpha", INTRO_LOGO_ALPHA,
        ]
    cmd = [sys.executable, str(MOVMAKER), str(local_input), "--out-dir", str(local_out), "--tmp-dir", str(temp_dir), *intro_args, *extra_args]
    log_dir = local_input / ".logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    log_path = log_dir / f"mvlog_worker_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as exc:
        log_path.write_text(
            "Command failed:\n" + " ".join(cmd) + "\n\nSTDOUT:\n" + (exc.stdout or "") + "\n\nSTDERR:\n" + (exc.stderr or ""),
            encoding="utf-8",
        )
        raise RuntimeError(f"movmaker failed; see {log_path}") from exc
    else:
        if result.stdout or result.stderr:
            log_path.write_text(
                "Command:\n" + " ".join(cmd) + "\n\nSTDOUT:\n" + (result.stdout or "") + "\n\nSTDERR:\n" + (result.stderr or ""),
                encoding="utf-8",
            )
    outputs = sorted(local_out.glob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True)
    if not outputs:
        raise RuntimeError("movmaker finished but no MP4 was produced")
    return outputs[0]


def process_job(host: str | None, remote_root: str, job: str, work_dir: Path, temp_dir: Path, extra_args: list[str], force: bool) -> None:
    job_dir = f"{remote_root}/in-dir/{job}"
    if host is None and not local_dir_accessible(Path(job_dir)):
        print(f"skip inaccessible input dir: {job}")
        return
    article_id = read_article_id(host, job_dir)
    state = read_state(host, job_dir)
    full_requested = remote_file_exists(host, f"{job_dir}/{ENABLED_NAME}")
    preview_requested = remote_file_exists(host, f"{job_dir}/{PREVIEW_NAME}")
    if not full_requested and not preview_requested:
        print(f"disabled: {job}")
        return
    preview_mode = preview_requested and not full_requested
    if not acquire_lock(host, job_dir):
        print(f"skip locked: {job}")
        return
    try:
        fingerprint_before = remote_fingerprint(host, job_dir)
        state = read_state(host, job_dir)
        old_output = state.get("output") if isinstance(state.get("output"), str) else None
        old_preview_output = state.get("preview_output") if isinstance(state.get("preview_output"), str) else None
        if preview_mode:
            if not force and state.get("preview_fingerprint") == fingerprint_before and old_preview_output:
                if remote_file_exists(host, f"{remote_root}/out-dir/{old_preview_output}"):
                    print(f"preview up to date: {job} -> {old_preview_output}")
                    return
        elif not force and state.get("status") == "done" and state.get("fingerprint") == fingerprint_before and old_output:
            if remote_file_exists(host, f"{remote_root}/out-dir/{old_output}"):
                print(f"up to date: {job} -> {old_output}")
                return

        processing_state = dict(state)
        processing_state.update({
            "status": "processing",
            "mode": "preview" if preview_mode else "full",
            "started_at": now(),
        })
        if article_id:
            processing_state["article_id"] = article_id
        if preview_mode:
            processing_state["preview_fingerprint"] = fingerprint_before
        else:
            processing_state["fingerprint"] = fingerprint_before
            processing_state["output"] = old_output
        processing_state.pop("message", None)
        write_state(host, job_dir, processing_state)

        local_job = work_dir / job
        if local_job.exists():
            shutil.rmtree(local_job)
        local_input = local_job / "input"
        local_out = local_job / "out"
        copy_job_from_remote(host, job_dir, local_input)
        render_args = [*extra_args, "--preview"] if preview_mode else extra_args
        print(f"rendering {'preview: ' if preview_mode else ''}{job}")
        local_mp4 = render_job(local_input, local_out, render_args, temp_dir)
        if preview_mode:
            preview_mp4 = local_mp4.with_name(f"{local_mp4.stem}_preview{local_mp4.suffix}")
            local_mp4.rename(preview_mp4)
            local_mp4 = preview_mp4

        fingerprint_after = remote_fingerprint(host, job_dir)
        if fingerprint_after != fingerprint_before:
            stale_state = dict(state)
            stale_state.update({
                "status": "stale",
                "updated_at": now(),
                "message": "Input changed during render; will regenerate on next run.",
            })
            if article_id:
                stale_state["article_id"] = article_id
            if preview_mode:
                stale_state["preview_fingerprint"] = fingerprint_after
                stale_state["previous_preview_fingerprint"] = fingerprint_before
            else:
                stale_state["fingerprint"] = fingerprint_after
                stale_state["previous_fingerprint"] = fingerprint_before
                stale_state["output"] = old_output
            write_state(host, job_dir, stale_state)
            print(f"stale during render, not publishing: {job}")
            return

        if preview_mode:
            preview_output = publish_output(host, remote_root, local_mp4, old_preview_output)
            # Upload any audio files present in the local input directory to the remote in-dir if missing.
            try:
                upload_audio_files(host, job_dir, local_input)
            except Exception as exc:  # Upload failure should not prevent publishing the preview
                print(f"warning: audio upload failed: {exc}")
            # Recompute fingerprint now that audio files may have been uploaded.
            final_fingerprint = remote_fingerprint(host, job_dir)
            new_state = dict(state)
            new_state.update({
                "status": "done",
                "mode": "preview",
                "preview_fingerprint": final_fingerprint,
                "preview_output": preview_output,
                "preview_generated_at": now(),
            })
            if article_id:
                new_state["article_id"] = article_id
            new_state.pop("message", None)
            write_state(host, job_dir, new_state)
            print(f"published preview: {job} -> {preview_output}")
            
            # Cleanup local cache
            if local_input.exists():
                shutil.rmtree(local_input)
            if local_out.exists():
                shutil.rmtree(local_out)
            return

        output = publish_output(host, remote_root, local_mp4, old_output)
        try:
            upload_audio_files(host, job_dir, local_input)
        except Exception as exc:
            print(f"warning: audio upload failed: {exc}")
        final_fingerprint = remote_fingerprint(host, job_dir)
        new_state = dict(state)
        new_state.update({
            "status": "done",
            "mode": "full",
            "fingerprint": final_fingerprint,
            "output": output,
            "generated_at": now(),
        })
        if article_id:
            new_state["article_id"] = article_id
        new_state.pop("message", None)
        write_state(host, job_dir, new_state)
        print(f"published: {job} -> {output}")
        
        # Cleanup local cache after successful job
        if local_input.exists():
            shutil.rmtree(local_input)
        if local_out.exists():
            shutil.rmtree(local_out)
        return
    except Exception as exc:
        try:
            error_state = {
                "status": "error",
                "updated_at": now(),
                "message": str(exc),
            }
            if article_id:
                error_state["article_id"] = article_id
            write_state(host, job_dir, error_state)
        finally:
            raise
    finally:
        release_lock(host, job_dir)


def main() -> int:
    parser = argparse.ArgumentParser(description="Render changed MVLog input dirs from a remote web server.")
    parser.add_argument("--local", action="store_true", help="Treat remote paths as local; skip SSH")
    parser.add_argument("--host", default="192.168.0.204", help="SSH host for the web server")
    parser.add_argument("--remote-root", default="/var/www/html", help="Remote MVLog web root")
    parser.add_argument("--work-dir", default=str(Path.home() / ".cache" / "mvlog-worker"), help="Local work directory")
    parser.add_argument("--tmp-dir", default=None, help="Large temporary render directory; defaults to WORK_DIR/tmp")
    parser.add_argument("--job", action="append", help="Only process this input dir name; may be repeated")
    parser.add_argument("--force", action="store_true", help="Render even if fingerprint is unchanged")
    parser.add_argument("--movmaker-arg", action="append", default=[], help="Extra argument passed to movmaker; may be repeated")
    args = parser.parse_args()

    work_dir = Path(args.work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)
    temp_dir = Path(args.tmp_dir) if args.tmp_dir else work_dir / "tmp"
    temp_dir.mkdir(parents=True, exist_ok=True)
    os.environ["TMPDIR"] = str(temp_dir)
    os.environ["TEMP"] = str(temp_dir)
    os.environ["TMP"] = str(temp_dir)
    os.environ["MOVMAKER_TMPDIR"] = str(temp_dir)
    jobs = args.job or list_jobs(None if args.local else args.host, args.remote_root)
    for job in jobs:
        process_job(None if args.local else args.host, args.remote_root, job, work_dir, temp_dir, args.movmaker_arg, args.force)
    return 0


if __name__ == "__main__":
    logging.info('Worker finished.')
    raise SystemExit(main())