app/core/preflight.py

import json
import shutil
import subprocess
import sys
import time
from pathlib import Path


class BatchPreflight:
    def __init__(self, base_dir, output_dir, temp_dir, ffmpeg_path):
        self.base_dir = Path(base_dir)
        self.output_dir = Path(output_dir)
        self.temp_dir = Path(temp_dir)
        self.ffmpeg_path = Path(ffmpeg_path)
        self.python_exe = self._resolve_python_executable()

    def _resolve_python_executable(self):
        env_python = self.base_dir / "env" / "Scripts" / "python.exe"
        if env_python.exists():
            return env_python
        return Path(sys.executable)

    def _gpu_env(self):
        import os
        env = os.environ.copy()
        env["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
        env["PYTHONUTF8"] = "1"
        nvidia_root = self.base_dir / "env" / "Lib" / "site-packages" / "nvidia"
        extra_paths = []
        if nvidia_root.exists():
            for child in nvidia_root.iterdir():
                for sub in ("bin", "lib"):
                    p = child / sub
                    if p.exists():
                        extra_paths.append(str(p))
        if extra_paths:
            # FIX: dùng os.pathsep (Linux/Docker là ":", Windows là ";")
            env["PATH"] = os.pathsep.join(extra_paths + [env.get("PATH", "")])
        prev_py = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = str(self.base_dir) + (os.pathsep + prev_py if prev_py else "")
        return env

    def run(self, asr_engine, trans_engine, tts_engine, strict_gpu, log_fn=None, required_checks=None):
        required_checks = set(required_checks or [])
        checks = []
        checks.append(self._check_writable("output_writable", self.output_dir, True))
        checks.append(self._check_writable("temp_writable", self.temp_dir, True))
        checks.append(self._check_disk_space("disk_space", self.output_dir, min_free_gb=5, required=True))
        checks.append(self._check_nvenc(required="ffmpeg_h264_nvenc" in required_checks))
        checks.append(self._check_python_cuda("torch_cuda", required="torch_cuda" in required_checks))
        checks.append(self._check_asr_cuda(required="faster_whisper_ctranslate2_cuda" in required_checks))
        if "PaddleOCR" in asr_engine or "Quét chữ" in asr_engine:
            checks.append(self._check_paddle_cuda(required="paddleocr_cuda" in required_checks))
        checks.append(self._check_onnx_cuda(required="onnxruntime_cuda_provider" in required_checks))
        checks.append(self._check_provider_config(trans_engine, required=True))

        report = {
            "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
            "strict_gpu": strict_gpu,
            "asr_engine": asr_engine,
            "trans_engine": trans_engine,
            "tts_engine": tts_engine,
            "checks": checks,
            "status": "PASS" if all(c["ok"] or not c["required"] for c in checks) else "FAILED",
        }
        self._write_reports(report)
        if log_fn:
            for c in checks:
                state = "PASS" if c["ok"] else ("FAIL" if c["required"] else "WARN")
                log_fn(f"[PREFLIGHT] {state}: {c['name']} - {c['message']}")
            log_fn(f"[PREFLIGHT] report: {self.output_dir / 'preflight_report.json'}")
        return report

    def _run_cmd(self, cmd, timeout=20):
        startupinfo = None
        if sys.platform == "win32":
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        try:
            res = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                encoding="utf-8",
                errors="ignore",
                timeout=timeout,
                startupinfo=startupinfo,
                env=self._gpu_env(),
            )
            return res.returncode, (res.stdout or ""), (res.stderr or "")
        except Exception as e:
            return 99, "", str(e)

    def _check_writable(self, name, path, required):
        try:
            path.mkdir(parents=True, exist_ok=True)
            probe = path / ".preflight_write_test"
            probe.write_text("ok", encoding="utf-8")
            probe.unlink(missing_ok=True)
            return self._result(name, True, required, f"{path} writable")
        except Exception as e:
            return self._result(name, False, required, str(e))

    def _check_disk_space(self, name, path, min_free_gb, required):
        try:
            usage = shutil.disk_usage(path)
            free_gb = usage.free / (1024 ** 3)
            return self._result(name, free_gb >= min_free_gb, required, f"{free_gb:.1f} GB free")
        except Exception as e:
            return self._result(name, False, required, str(e))

    def _check_nvenc(self, required):
        cmd = [
            str(self.ffmpeg_path), "-y", "-f", "lavfi", "-i", "color=c=black:s=256x256",
            "-t", "1", "-c:v", "h264_nvenc", "-f", "null", "-"
        ]
        code, out, err = self._run_cmd(cmd, timeout=15)
        msg = "h264_nvenc OK" if code == 0 else (err or out)[-300:]
        return self._result("ffmpeg_h264_nvenc", code == 0, required, msg)

    def _check_python_cuda(self, name, required):
        code = (
            "import torch; "
            "assert torch.cuda.is_available(), 'torch cuda unavailable'; "
            "x=torch.randn(1, device='cuda'); "
            "print(torch.cuda.get_device_name(0))"
        )
        rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=30)
        msg = out.strip() if rc == 0 else (err or out)[-300:]
        return self._result(name, rc == 0, required, msg)

    def _check_asr_cuda(self, required):
        code = (
            "import os, tempfile, wave; "
            "import ctranslate2; "
            "n=ctranslate2.get_cuda_device_count(); "
            "assert n>0, 'ctranslate2 cuda device count is 0'; "
            "from faster_whisper import WhisperModel; "
            "fd,path=tempfile.mkstemp(suffix='.wav'); os.close(fd); "
            "wf=wave.open(path,'wb'); wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(16000); "
            "wf.writeframes(b'\\x00\\x00'*16000); wf.close(); "
            "model=WhisperModel('base', device='cuda', compute_type='float16'); "
            "segments,info=model.transcribe(path, vad_filter=True, beam_size=1); "
            "list(segments); "
            "os.remove(path); "
            "print('ctranslate2_cuda_transcribe_ok', ctranslate2.__version__, 'devices', n)"
        )
        rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=90)
        msg = out.strip() if rc == 0 else (err or out)[-300:]
        return self._result("faster_whisper_ctranslate2_cuda", rc == 0, required, msg)

    def _check_paddle_cuda(self, required):
        code = (
            "import os, json; "
            "os.environ.setdefault('PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION','python'); "
            "from app.core.ocr_worker_cli import register_ocr_gpu_dll_paths, _run_ocr_quality_probe; "
            "register_ocr_gpu_dll_paths(); "
            "import paddle; "
            "ok=paddle.device.is_compiled_with_cuda(); "
            "assert ok, 'paddle is not compiled with cuda'; "
            "paddle.set_device('gpu:0'); "
            "from paddleocr import PaddleOCR; "
            "import inspect; "
            "kw=dict(use_angle_cls=True, lang='ch', show_log=False); "
            "kw.update({'device': 'gpu:0'} if 'device' in inspect.signature(PaddleOCR).parameters else {'use_gpu': True}); "
            "ocr=PaddleOCR(**kw); "
            "quality_ok, details=_run_ocr_quality_probe(ocr); "
            "assert quality_ok, details; "
            "print(paddle.device.get_device(), json.dumps(details, ensure_ascii=False))"
        )
        rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=60)
        msg = out.strip() if rc == 0 else (err or out)[-300:]
        return self._result("paddleocr_cuda", rc == 0, required, msg)

    def _check_onnx_cuda(self, required):
        code = (
            "import onnxruntime as ort; "
            "providers=ort.get_available_providers(); "
            "assert 'CUDAExecutionProvider' in providers, providers; "
            "print(providers)"
        )
        rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=15)
        msg = out.strip() if rc == 0 else (err or out)[-300:]
        return self._result("onnxruntime_cuda_provider", rc == 0, required, msg)

    @staticmethod
    def _has_pool_key(cfg):
        """ApiPool đọc key từ environ + root .env (OPENROUTER/GROQ/NVIDIA/GEMINI/..._KEY_N),
        config.json (hf_key/groq_key/nvidia_key) và gate .env. Check đúng nguồn pipeline dùng."""
        import os as _os
        prefixes = ("OPENROUTER_KEY", "GROQ_KEY", "NVIDIA_KEY", "GEMINI_KEY",
                    "SILICONFLOW_KEY", "TOGETHERAI_KEY", "XKIRO_KEY",
                    "SUPER_AI_GATE_KEY", "NIM_API_KEY")
        for k, v in _os.environ.items():
            if v and str(v).strip() and k.startswith(prefixes):
                return True
        for name in ("hf_key", "groq_key", "nvidia_key", "nim_api_key",
                     "gemini_key", "openrouter_key"):
            if cfg.get(name):
                return True
        for env_file in (Path(__file__).resolve().parents[2] / ".env",
                         Path("D:/TOOL GOM API/.env")):
            try:
                if env_file.exists():
                    for line in env_file.read_text(encoding="utf-8", errors="ignore").splitlines():
                        line = line.strip()
                        if not line or line.startswith("#") or "=" not in line:
                            continue
                        k, v = line.split("=", 1)
                        if v.strip().strip('"').strip("'") and k.strip().startswith(prefixes):
                            return True
            except Exception:
                pass
        return False

    def _check_provider_config(self, trans_engine, required):
        config_path = self.base_dir / "config.json"
        try:
            cfg = json.loads(config_path.read_text(encoding="utf-8")) if config_path.exists() else {}
            eng = str(trans_engine or "")
            # FIX: "Google" Reserve cho bản miễn phí thật (Google Free/Dịch web).
            # Gemini (dù tên có chữ Google) vẫn cần GEMINI_KEY.
            if "Gemini" in eng and "Web" not in eng:
                if self._has_pool_key(cfg):
                    return self._result("translation_provider_config", True, required, "Gemini key present")
                return self._result("translation_provider_config", False, required, "missing GEMINI_KEY for Gemini engine")
            if "Google" in eng:
                return self._result("translation_provider_config", True, required, "Google (free) selected")
            if "Ollama" in eng:
                model = cfg.get("ollama_model") or cfg.get("translation", {}).get("ollama_model")
                return self._result("translation_provider_config", bool(model), required, "Ollama model configured" if model else "missing Ollama model")
            # FIX: Ultimate/Pool/9Router/SuperAI/Nvidia/Groq/... chạy qua ApiPool —
            # check đúng nguồn key của pool thay vì chỉ hf_key/groq_key.
            if any(k in eng for k in ("Pool", "Ultimate", "Tối thượng", "9Router", "Super AI",
                                      "Nvidia", "NIM", "Groq", "OpenRouter", "SiliconFlow",
                                      "TogetherAI", "Nemotron", "DeepSeek", "xKiro", "Custom Trans")):
                if self._has_pool_key(cfg):
                    return self._result("translation_provider_config", True, required, "API Pool key present")
                return self._result("translation_provider_config", False, required, "missing API Pool key (.env *_KEY_N)")
            has_key = bool(cfg.get("hf_key") or cfg.get("groq_key"))
            return self._result("translation_provider_config", has_key, required, "API key present" if has_key else "missing provider API key")
        except Exception as e:
            return self._result("translation_provider_config", False, required, str(e))

    def _write_reports(self, report):
        self.output_dir.mkdir(parents=True, exist_ok=True)
        (self.output_dir / "preflight_report.json").write_text(
            json.dumps(report, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        lines = [
            "=== PREFLIGHT REPORT ===",
            f"Status: {report['status']}",
            f"Strict GPU: {report['strict_gpu']}",
        ]
        for c in report["checks"]:
            state = "PASS" if c["ok"] else ("FAIL" if c["required"] else "WARN")
            lines.append(f"- {state} {c['name']}: {c['message']}")
        (self.output_dir / "preflight_report.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")

    @staticmethod
    def _result(name, ok, required, message):
        return {
            "name": name,
            "ok": bool(ok),
            "required": bool(required),
            "message": str(message).strip(),
        }