app/main.py

from __future__ import annotations

import json
import os
import platform
import shutil
import subprocess
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import requests
from dotenv import load_dotenv
from flask import Flask, jsonify, render_template, request, send_file
from werkzeug.utils import secure_filename

from app.estimates import DEFAULT_MAX_PARALLEL, DEFAULT_SEGMENT_SECONDS, estimate_run

ROOT = Path(__file__).resolve().parents[1]
RUNS_DIR = ROOT / "runs"
API_BASE = "https://api.replicate.com/v1"
MODEL_PATH = "heygen/video-translate"

POLL_SECONDS = 8


class RunCancelledError(RuntimeError):
    pass


def ensure_ffmpeg() -> None:
    if shutil.which("ffmpeg") is None:
        raise RuntimeError("`ffmpeg` is not installed or not on PATH.")


@dataclass
class SegmentTaskResult:
    index: int
    status: str
    error: str | None = None


class TranslationRunManager:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._threads: dict[str, threading.Thread] = {}
        self._cancel_flags: dict[str, threading.Event] = {}
        self._load_existing_runs()

    def _load_existing_runs(self) -> None:
        RUNS_DIR.mkdir(parents=True, exist_ok=True)
        for run_dir in RUNS_DIR.iterdir():
            if not run_dir.is_dir():
                continue
            run_file = run_dir / "run.json"
            if not run_file.exists():
                continue
            data = json.loads(run_file.read_text())
            if data.get("status") == "running":
                data["status"] = "interrupted"
                data["error"] = "App restarted while run was in progress."
                self._save_run(data)

    def list_runs(self) -> list[dict[str, Any]]:
        runs = []
        for run_dir in sorted(RUNS_DIR.iterdir(), reverse=True):
            if run_dir.name.startswith("_"):
                continue
            run_file = run_dir / "run.json"
            if run_file.exists():
                runs.append(json.loads(run_file.read_text()))
        return runs

    def get_run(self, run_id: str) -> dict[str, Any] | None:
        run_file = RUNS_DIR / run_id / "run.json"
        if not run_file.exists():
            return None
        return json.loads(run_file.read_text())

    def create_run(
        self,
        video_path: Path,
        output_language: str,
        mode: str,
        segment_seconds: int,
        max_parallel: int,
    ) -> dict[str, Any]:
        run_id = time.strftime("%Y%m%d-%H%M%S") + "-" + uuid.uuid4().hex[:6]
        run_dir = RUNS_DIR / run_id
        run_dir.mkdir(parents=True, exist_ok=False)
        input_dir = run_dir / "input"
        input_dir.mkdir(exist_ok=True)
        input_video = input_dir / secure_filename(video_path.name)
        shutil.copy2(video_path, input_video)
        duration_seconds = self._probe_duration(input_video)
        run_estimate = (
            estimate_run(duration_seconds, mode) if duration_seconds is not None else None
        )

        run_data = {
            "id": run_id,
            "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
            "input_filename": input_video.name,
            "input_duration_seconds": duration_seconds,
            "estimate": run_estimate,
            "status": "queued",
            "error": None,
            "settings": {
                "output_language": output_language,
                "mode": mode,
                "segment_seconds": segment_seconds,
                "max_parallel": max_parallel,
            },
            "paths": {
                "run_dir": str(run_dir),
                "input_video": str(input_video),
                "segments_dir": str(run_dir / "segments"),
                "outputs_dir": str(run_dir / "translated_segments"),
                "merged_output": str(run_dir / "output" / f"{input_video.stem}_translated.mp4"),
                "log_file": str(run_dir / "events.log"),
            },
            "segments": [],
        }
        self._save_run(run_data)
        return run_data

    def start_run(self, run_id: str) -> None:
        run = self.get_run(run_id)
        if not run:
            raise RuntimeError(f"Run `{run_id}` not found.")
        if run.get("status") == "running":
            raise RuntimeError("Run is already active.")

        with self._lock:
            self._cancel_flags[run_id] = threading.Event()
        thread = threading.Thread(target=self._execute_run, args=(run_id, None), daemon=True)
        with self._lock:
            self._threads[run_id] = thread
        thread.start()

    def cancel_run(self, run_id: str) -> dict[str, Any]:
        run = self.get_run(run_id)
        if not run:
            raise RuntimeError(f"Run `{run_id}` not found.")
        if run.get("status") not in {"queued", "running"}:
            raise RuntimeError("Only active runs can be cancelled.")

        with self._lock:
            if run_id not in self._cancel_flags:
                self._cancel_flags[run_id] = threading.Event()
            self._cancel_flags[run_id].set()

        token = os.environ.get("REPLICATE_API_TOKEN")
        if token:
            for segment in run.get("segments", []):
                prediction_id = segment.get("prediction_id")
                if prediction_id and segment.get("status") == "processing":
                    self._cancel_prediction(prediction_id, token)

        self._set_run_status(run_id, "cancelled", "Cancelled by user.")
        self._log(run, "Cancellation requested.")
        updated = self.get_run(run_id)
        if not updated:
            raise RuntimeError(f"Run `{run_id}` not found.")
        return updated

    def delete_run(self, run_id: str) -> None:
        run = self.get_run(run_id)
        if not run:
            raise RuntimeError(f"Run `{run_id}` not found.")
        if run.get("status") == "running":
            raise RuntimeError("Cannot delete a running job. Cancel it first.")
        with self._lock:
            thread = self._threads.get(run_id)
            if thread and thread.is_alive():
                raise RuntimeError("Cannot delete while run is active.")

        run_dir = Path(run["paths"]["run_dir"])
        if run_dir.exists():
            shutil.rmtree(run_dir)
        with self._lock:
            self._threads.pop(run_id, None)
            self._cancel_flags.pop(run_id, None)

    def _is_cancelled(self, run_id: str) -> bool:
        flag = self._cancel_flags.get(run_id)
        return flag is not None and flag.is_set()

    def retry_failed(self, run_id: str) -> None:
        run = self.get_run(run_id)
        if not run:
            raise RuntimeError(f"Run `{run_id}` not found.")
        if run.get("status") == "running":
            raise RuntimeError("Run is currently running.")

        failed_indices = [seg["index"] for seg in run["segments"] if seg["status"] == "failed"]
        if not failed_indices:
            raise RuntimeError("No failed segments to retry.")

        with self._lock:
            self._cancel_flags[run_id] = threading.Event()
        thread = threading.Thread(target=self._execute_run, args=(run_id, failed_indices), daemon=True)
        with self._lock:
            self._threads[run_id] = thread
        thread.start()

    def _execute_run(self, run_id: str, retry_indices: list[int] | None) -> None:
        run = self.get_run(run_id)
        if not run:
            return

        try:
            token = os.environ.get("REPLICATE_API_TOKEN")
            if not token:
                raise RuntimeError("Missing REPLICATE_API_TOKEN in environment.")
            ensure_ffmpeg()

            if retry_indices is None:
                self._split_video(run)
                if self._is_cancelled(run_id):
                    return

            self._set_run_status(run_id, "running", None)
            updated_run = self.get_run(run_id)
            if not updated_run:
                return

            max_parallel = int(updated_run["settings"]["max_parallel"])
            targets = updated_run["segments"]
            if retry_indices is not None:
                targets = [seg for seg in targets if seg["index"] in retry_indices]

            with ThreadPoolExecutor(max_workers=max_parallel) as executor:
                futures = [
                    executor.submit(self._process_segment, run_id, segment["index"], token)
                    for segment in targets
                ]
                for _ in as_completed(futures):
                    if self._is_cancelled(run_id):
                        break

            if self._is_cancelled(run_id):
                return

            final = self.get_run(run_id)
            if not final:
                return
            failed = [seg for seg in final["segments"] if seg["status"] != "succeeded"]
            if failed:
                if not self._is_cancelled(run_id):
                    self._set_run_status(run_id, "failed", "Some segments failed. Retry failed segments.")
                return

            merged_path = self._merge_segments(final)
            current = self.get_run(run_id)
            if not current:
                return
            current["paths"]["merged_output"] = str(merged_path)
            current["status"] = "completed"
            current["error"] = None
            self._save_run(current)
            self._log(current, "Run completed successfully.")
        except Exception as exc:  # pylint: disable=broad-except
            if not self._is_cancelled(run_id):
                self._set_run_status(run_id, "failed", str(exc))
        finally:
            with self._lock:
                self._threads.pop(run_id, None)

    def _split_video(self, run: dict[str, Any]) -> None:
        run_dir = Path(run["paths"]["run_dir"])
        segments_dir = run_dir / "segments"
        outputs_dir = run_dir / "translated_segments"
        output_dir = run_dir / "output"
        for directory in (segments_dir, outputs_dir, output_dir):
            directory.mkdir(parents=True, exist_ok=True)

        input_video = Path(run["paths"]["input_video"])
        segment_seconds = int(run["settings"]["segment_seconds"])
        self._log(run, f"Splitting input into {segment_seconds}s chunks.")
        segment_pattern = str(segments_dir / "seg_%03d.mp4")
        command = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-y",
            "-i",
            str(input_video),
            "-c",
            "copy",
            "-f",
            "segment",
            "-segment_time",
            str(segment_seconds),
            "-reset_timestamps",
            "1",
            segment_pattern,
        ]
        self._run_cmd(command, "Video segmentation failed.")

        segments = sorted(segments_dir.glob("seg_*.mp4"))
        if not segments:
            raise RuntimeError("No segments produced. Check ffmpeg and input file.")

        run["segments"] = [
            {
                "index": idx,
                "name": seg.name,
                "path": str(seg),
                "status": "queued",
                "prediction_id": None,
                "input_video_url": None,
                "output_url": None,
                "downloaded_path": str(outputs_dir / seg.name),
                "error": None,
            }
            for idx, seg in enumerate(segments)
        ]
        self._save_run(run)
        self._log(run, f"Created {len(segments)} segment(s).")

    def _process_segment(self, run_id: str, index: int, token: str) -> SegmentTaskResult:
        run = self.get_run(run_id)
        if not run:
            return SegmentTaskResult(index=index, status="failed", error="Run missing.")
        if self._is_cancelled(run_id):
            return SegmentTaskResult(index=index, status="failed", error="Cancelled")
        segment = run["segments"][index]

        try:
            self._update_segment(run_id, index, status="uploading", error=None)
            if self._is_cancelled(run_id):
                raise RunCancelledError("Run cancelled by user.")
            file_url = self._upload_segment(Path(segment["path"]), token)
            self._update_segment(run_id, index, input_video_url=file_url)

            self._update_segment(run_id, index, status="processing")
            prediction = self._submit_prediction(file_url, run["settings"], token)
            prediction_id = prediction["id"]
            self._update_segment(run_id, index, prediction_id=prediction_id)

            terminal = self._poll_prediction(
                prediction_id,
                token,
                is_cancelled=lambda: self._is_cancelled(run_id),
            )
            if terminal["status"] != "succeeded":
                raise RuntimeError(terminal.get("error") or f"Prediction {terminal['status']}")

            output_url = terminal.get("output")
            if not output_url:
                raise RuntimeError("Prediction succeeded but output URL missing.")

            if self._is_cancelled(run_id):
                raise RunCancelledError("Run cancelled by user.")

            self._update_segment(run_id, index, status="downloading", output_url=output_url)
            output_path = Path(segment["downloaded_path"])
            self._download_file(output_url, output_path)
            self._update_segment(run_id, index, status="succeeded", error=None)
            return SegmentTaskResult(index=index, status="succeeded")
        except RunCancelledError:
            self._update_segment(run_id, index, status="failed", error="Cancelled")
            return SegmentTaskResult(index=index, status="failed", error="Cancelled")
        except Exception as exc:  # pylint: disable=broad-except
            self._update_segment(run_id, index, status="failed", error=str(exc))
            return SegmentTaskResult(index=index, status="failed", error=str(exc))

    def _upload_segment(self, file_path: Path, token: str) -> str:
        with file_path.open("rb") as handle:
            response = requests.post(
                f"{API_BASE}/files",
                headers={"Authorization": f"Token {token}"},
                files={
                    "content": (file_path.name, handle, "video/mp4"),
                    "metadata": ("metadata", "{}", "application/json"),
                },
                timeout=1800,
            )
        if response.status_code >= 400:
            raise RuntimeError(f"Replicate file upload failed: {response.status_code} {response.text}")
        payload = response.json()
        url = (payload.get("urls") or {}).get("get")
        if not url:
            raise RuntimeError("Replicate file upload missing urls.get")
        return url

    def _submit_prediction(self, file_url: str, settings: dict[str, Any], token: str) -> dict[str, Any]:
        body = {
            "input": {
                "video": file_url,
                "output_language": settings["output_language"],
                "mode": settings["mode"],
            }
        }
        response = requests.post(
            f"{API_BASE}/models/{MODEL_PATH}/predictions",
            headers={
                "Authorization": f"Bearer {token}",
                "Content-Type": "application/json",
                "User-Agent": "video-translator-local-app/1.0",
            },
            json=body,
            timeout=120,
        )
        if response.status_code >= 400:
            raise RuntimeError(f"Prediction submit failed: {response.status_code} {response.text}")
        return response.json()

    def _poll_prediction(
        self,
        prediction_id: str,
        token: str,
        is_cancelled: Any = None,
    ) -> dict[str, Any]:
        while True:
            if is_cancelled and is_cancelled():
                self._cancel_prediction(prediction_id, token)
                raise RunCancelledError("Run cancelled by user.")
            response = requests.get(
                f"{API_BASE}/predictions/{prediction_id}",
                headers={"Authorization": f"Bearer {token}", "User-Agent": "video-translator-local-app/1.0"},
                timeout=60,
            )
            if response.status_code >= 400:
                raise RuntimeError(f"Prediction poll failed: {response.status_code} {response.text}")
            payload = response.json()
            if payload["status"] in {"succeeded", "failed", "canceled"}:
                return payload
            time.sleep(POLL_SECONDS)

    def _cancel_prediction(self, prediction_id: str, token: str) -> None:
        try:
            requests.post(
                f"{API_BASE}/predictions/{prediction_id}/cancel",
                headers={"Authorization": f"Bearer {token}", "User-Agent": "video-translator-local-app/1.0"},
                timeout=30,
            )
        except requests.RequestException:
            pass

    def _download_file(self, url: str, destination: Path) -> None:
        destination.parent.mkdir(parents=True, exist_ok=True)
        with requests.get(url, stream=True, timeout=1800) as response:
            if response.status_code >= 400:
                raise RuntimeError(f"Download failed: {response.status_code} {response.text}")
            with destination.open("wb") as handle:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        handle.write(chunk)

    def _merge_segments(self, run: dict[str, Any]) -> Path:
        output_dir = Path(run["paths"]["run_dir"]) / "output"
        output_dir.mkdir(parents=True, exist_ok=True)
        concat_file = Path(run["paths"]["run_dir"]) / "concat.txt"

        lines = []
        for segment in run["segments"]:
            downloaded = Path(segment["downloaded_path"]).resolve()
            if not downloaded.exists():
                raise RuntimeError(f"Missing downloaded segment: {downloaded}")
            escaped = str(downloaded).replace("'", "'\\''")
            lines.append(f"file '{escaped}'")
        concat_file.write_text("\n".join(lines) + "\n")

        merged_output = Path(run["paths"]["merged_output"])
        command = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-y",
            "-f",
            "concat",
            "-safe",
            "0",
            "-i",
            str(concat_file),
            "-c",
            "copy",
            str(merged_output),
        ]
        self._run_cmd(command, "Merging translated segments failed.")
        return merged_output

    def _update_segment(self, run_id: str, index: int, **changes: Any) -> None:
        run = self.get_run(run_id)
        if not run:
            return
        segment = run["segments"][index]
        segment.update(changes)
        self._save_run(run)
        label = segment["name"]
        status = segment.get("status")
        error = segment.get("error")
        if status:
            message = f"[{label}] {status}"
            if error:
                message = f"{message} :: {error}"
            self._log(run, message)

    def _set_run_status(self, run_id: str, status: str, error: str | None) -> None:
        run = self.get_run(run_id)
        if not run:
            return
        run["status"] = status
        run["error"] = error
        self._save_run(run)
        self._log(run, f"Run status -> {status}" + (f" ({error})" if error else ""))

    def _save_run(self, run_data: dict[str, Any]) -> None:
        run_dir = Path(run_data["paths"]["run_dir"])
        run_dir.mkdir(parents=True, exist_ok=True)
        run_file = run_dir / "run.json"
        with self._lock:
            run_file.write_text(json.dumps(run_data, indent=2))

    def _log(self, run_data: dict[str, Any], message: str) -> None:
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        log_line = f"[{timestamp}] {message}\n"
        log_file = Path(run_data["paths"]["log_file"])
        with self._lock:
            log_file.parent.mkdir(parents=True, exist_ok=True)
            with log_file.open("a") as handle:
                handle.write(log_line)

    @staticmethod
    def _probe_duration(video_path: Path) -> float | None:
        if shutil.which("ffprobe") is None:
            return None
        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(video_path),
            ],
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            return None
        try:
            return float(result.stdout.strip())
        except ValueError:
            return None

    @staticmethod
    def _run_cmd(command: list[str], failure_message: str) -> None:
        result = subprocess.run(command, capture_output=True, text=True, check=False)
        if result.returncode != 0:
            stderr = result.stderr.strip() or result.stdout.strip()
            raise RuntimeError(f"{failure_message} {stderr}")


load_dotenv(ROOT / ".env")
app = Flask(__name__, template_folder="templates", static_folder="static")
manager = TranslationRunManager()


@app.get("/")
def index() -> str:
    return render_template("index.html")


@app.get("/api/health")
def health() -> Any:
    token_present = bool(os.environ.get("REPLICATE_API_TOKEN"))
    ffmpeg_present = shutil.which("ffmpeg") is not None
    return jsonify(
        {
            "token_present": token_present,
            "ffmpeg_present": ffmpeg_present,
            "platform": platform.platform(),
            "defaults": {
                "segment_seconds": DEFAULT_SEGMENT_SECONDS,
                "max_parallel": DEFAULT_MAX_PARALLEL,
            },
        }
    )


@app.get("/api/estimate")
def get_estimate() -> Any:
    try:
        duration_seconds = float(request.args.get("duration_seconds", "0"))
    except ValueError:
        return jsonify({"error": "duration_seconds must be a number."}), 400
    if duration_seconds <= 0:
        return jsonify({"error": "duration_seconds must be greater than 0."}), 400

    mode = request.args.get("mode", "precision").strip()
    if mode not in {"precision", "speed"}:
        return jsonify({"error": "Invalid mode."}), 400

    return jsonify(estimate_run(duration_seconds, mode))


@app.get("/api/runs")
def list_runs() -> Any:
    return jsonify({"runs": manager.list_runs()})


@app.get("/api/runs/<run_id>")
def get_run(run_id: str) -> Any:
    run = manager.get_run(run_id)
    if not run:
        return jsonify({"error": "Run not found"}), 404
    return jsonify(run)


@app.post("/api/runs")
def create_run() -> Any:
    if "video" not in request.files:
        return jsonify({"error": "No video uploaded."}), 400
    upload = request.files["video"]
    if not upload.filename:
        return jsonify({"error": "Please choose a video file."}), 400

    temp_dir = ROOT / "runs" / "_incoming"
    temp_dir.mkdir(parents=True, exist_ok=True)
    temp_path = temp_dir / secure_filename(upload.filename)
    upload.save(temp_path)

    output_language = request.form.get("output_language", "English").strip() or "English"
    mode = request.form.get("mode", "precision").strip()
    if mode not in {"precision", "speed"}:
        return jsonify({"error": "Invalid mode."}), 400

    segment_seconds = DEFAULT_SEGMENT_SECONDS
    max_parallel = DEFAULT_MAX_PARALLEL

    try:
        run_data = manager.create_run(
            video_path=temp_path,
            output_language=output_language,
            mode=mode,
            segment_seconds=segment_seconds,
            max_parallel=max_parallel,
        )
        manager.start_run(run_data["id"])
        return jsonify(run_data)
    except Exception as exc:  # pylint: disable=broad-except
        return jsonify({"error": str(exc)}), 500
    finally:
        if temp_path.exists():
            temp_path.unlink()


@app.post("/api/runs/<run_id>/cancel")
def cancel_run(run_id: str) -> Any:
    try:
        run = manager.cancel_run(run_id)
        return jsonify(run)
    except Exception as exc:  # pylint: disable=broad-except
        return jsonify({"error": str(exc)}), 400


@app.delete("/api/runs/<run_id>")
def delete_run(run_id: str) -> Any:
    try:
        manager.delete_run(run_id)
        return jsonify({"ok": True})
    except Exception as exc:  # pylint: disable=broad-except
        return jsonify({"error": str(exc)}), 400


@app.post("/api/runs/<run_id>/retry")
def retry_failed(run_id: str) -> Any:
    try:
        manager.retry_failed(run_id)
        run = manager.get_run(run_id)
        return jsonify(run)
    except Exception as exc:  # pylint: disable=broad-except
        return jsonify({"error": str(exc)}), 400


@app.get("/api/runs/<run_id>/output")
def stream_output(run_id: str) -> Any:
    run = manager.get_run(run_id)
    if not run:
        return jsonify({"error": "Run not found."}), 404
    output = Path(run["paths"]["merged_output"])
    if not output.exists():
        return jsonify({"error": "Output file is not ready."}), 400
    return send_file(output, mimetype="video/mp4", as_attachment=False)


@app.get("/api/runs/<run_id>/download")
def download_output(run_id: str) -> Any:
    run = manager.get_run(run_id)
    if not run:
        return jsonify({"error": "Run not found."}), 404
    output = Path(run["paths"]["merged_output"])
    if not output.exists():
        return jsonify({"error": "Output file is not ready."}), 400
    download_name = run.get("input_filename") or output.name
    if download_name and not download_name.endswith(".mp4"):
        stem = Path(download_name).stem
        download_name = f"{stem}_translated.mp4"
    return send_file(output, mimetype="video/mp4", as_attachment=True, download_name=download_name)


@app.post("/api/runs/<run_id>/open-output")
def open_output(run_id: str) -> Any:
    run = manager.get_run(run_id)
    if not run:
        return jsonify({"error": "Run not found."}), 404
    output = Path(run["paths"]["merged_output"])
    if not output.exists():
        return jsonify({"error": "Output file is not ready."}), 400

    try:
        if platform.system() == "Darwin":
            subprocess.run(["open", str(output.parent)], check=False)
        elif platform.system() == "Windows":
            os.startfile(str(output.parent))  # type: ignore[attr-defined]
        else:
            subprocess.run(["xdg-open", str(output.parent)], check=False)
    except Exception as exc:  # pylint: disable=broad-except
        return jsonify({"error": f"Unable to open folder: {exc}"}), 500

    return jsonify({"ok": True, "path": str(output)})


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5050, debug=False)