from __future__ import annotations

import asyncio
import hashlib
import io
import json
import sys
import tarfile
import types
from pathlib import Path

import httpx
import pytest
from imagefind.config import Settings
from imagefind.main import create_app
from imagefind.models import ModelManager
from imagefind.speech import SpeechService


def _bundle(*, unsafe: bool = False) -> bytes:
    output = io.BytesIO()
    with tarfile.open(fileobj=output, mode="w:gz") as archive:
        files = {
            "manifest.json": json.dumps({"version": "manual-test"}).encode(),
            "visual/image/modules.json": json.dumps([{"idx": 0, "path": "", "type": "Transformer"}]).encode(),
            "visual/text/modules.json": json.dumps([{"idx": 0, "path": "", "type": "Transformer"}]).encode(),
        }
        if unsafe:
            files["../outside"] = b"unsafe"
        for name, content in files.items():
            member = tarfile.TarInfo(name)
            member.size = len(content)
            archive.addfile(member, io.BytesIO(content))
    return output.getvalue()


def _app(tmp_path: Path):
    settings = Settings(
        data_dir=tmp_path / "data",
        embedding_backend="hash",
        upload_reserve_gb=0,
        model_upload_gb=1,
    )
    settings.prepare()
    app = create_app(settings)
    _, token = app.state.services.auth.create_api_token("models")
    return app, {"Authorization": f"Bearer {token}"}


def _write_audio_variant(root: Path) -> None:
    root.mkdir(parents=True, exist_ok=True)
    for name in (
        "config.json",
        "preprocessor_config.json",
        "tokenizer_config.json",
        "tokenizer.json",
        "openvino_encoder_model.xml",
        "openvino_encoder_model.bin",
        "openvino_decoder_model.xml",
        "openvino_decoder_model.bin",
    ):
        (root / name).write_text("{}", encoding="utf-8")


def test_manual_model_bundle_upload_and_atomic_install(tmp_path: Path):
    app, headers = _app(tmp_path)
    bundle = _bundle()

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.post(
                "/api/v1/models/upload",
                headers={
                    **headers,
                    "Content-Type": "application/octet-stream",
                    "X-Model-Filename": "imagefind-models.tar.gz",
                },
                content=bundle,
            )
            assert response.status_code == 202
            payload = response.json()
            assert payload["size_bytes"] == len(bundle)
            assert payload["sha256"] == hashlib.sha256(bundle).hexdigest()

        with app.state.services.db.read() as conn:
            row = conn.execute("SELECT payload_json FROM jobs WHERE id=?", (payload["job_id"],)).fetchone()
        job_payload = json.loads(row["payload_json"])
        archive_path = Path(job_payload["archive_path"])
        result = app.state.services.models.install(
            archive_path=archive_path,
            expected_sha256=job_payload["sha256"],
        )
        assert result["manifest"]["version"] == "manual-test"
        assert (app.state.services.settings.models_dir / "visual" / "image" / "modules.json").is_file()
        assert (app.state.services.settings.models_dir / "visual" / "text" / "modules.json").is_file()

    asyncio.run(scenario())


def test_medium_audio_variant_uses_isolated_verification(tmp_path: Path, monkeypatch):
    app, _headers = _app(tmp_path)
    service = app.state.services
    _write_audio_variant(service.settings.models_dir / "audio" / "small")
    _write_audio_variant(service.settings.models_dir / "audio" / "medium")
    service.settings.audio_model_variant = "small"
    service.speech.accelerator.gpu_device = "GPU.0"
    called: list[str] = []

    def forbidden_main_process_verify():
        raise AssertionError("Medium verification must not run in the API process")

    def isolated_verify(variant):
        called.append(f"{variant}:{service.settings.audio_model_variant}")
        service.speech.accelerator.mark_ready("audio", "GPU.0", ["GPU.0"])
        return service.speech.accelerator.status()["components"]["audio"]

    monkeypatch.setattr(service.speech, "verify_acceleration", forbidden_main_process_verify)
    monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)

    selected = service.models.set_audio_variant("medium")

    assert selected["active"] is True
    assert called == ["medium:small"]
    assert service.settings.audio_model_variant == "medium"


def test_medium_audio_variant_verification_failure_rolls_back(tmp_path: Path, monkeypatch):
    app, _headers = _app(tmp_path)
    service = app.state.services
    _write_audio_variant(service.settings.models_dir / "audio" / "small")
    _write_audio_variant(service.settings.models_dir / "audio" / "medium")
    service.settings.audio_model_variant = "small"
    service.speech.accelerator.gpu_device = "GPU.0"

    def isolated_verify(_variant):
        raise RuntimeError("worker crashed")

    monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)

    with pytest.raises(RuntimeError, match="worker crashed"):
        service.models.set_audio_variant("medium")

    assert service.settings.audio_model_variant == "small"


def test_speech_config_medium_failure_returns_conflict_and_keeps_small(tmp_path: Path, monkeypatch):
    app, headers = _app(tmp_path)
    service = app.state.services
    _write_audio_variant(service.settings.models_dir / "audio" / "small")
    _write_audio_variant(service.settings.models_dir / "audio" / "medium")
    service.settings.audio_model_variant = "small"
    service.speech.accelerator.gpu_device = "GPU.0"

    def isolated_verify(_variant):
        raise RuntimeError("worker crashed")

    monkeypatch.setattr(service.speech, "verify_variant_acceleration_isolated", isolated_verify)

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.patch(
                "/api/v1/speech/config",
                headers=headers,
                json={"model_variant": "medium"},
            )
            assert response.status_code == 409
            config = await client.get("/api/v1/speech/config", headers=headers)
            assert config.json()["model_variant"] == "small"

    asyncio.run(scenario())


def test_old_medium_fp16_variant_requires_int8_update(tmp_path: Path):
    settings = Settings(data_dir=tmp_path / "data")
    settings.prepare()
    medium = settings.models_dir / "audio" / "medium"
    _write_audio_variant(medium)
    manifest = {
        "components": {
            "audio": {
                "variants": {
                    "medium": {
                        "version": "OpenVINO/whisper-medium-fp16-ov",
                        "source": "https://hf-mirror.com/OpenVINO/whisper-medium-fp16-ov",
                    }
                },
                "active_variant": "small",
            }
        }
    }
    (settings.models_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
    manager = ModelManager.__new__(ModelManager)
    manager.settings = settings
    manager.speech = SpeechService(settings)
    manager._manifest_cached_mtime_ns = -1
    manager._manifest_cached = {}

    variant = manager.audio_variants()["medium"]

    assert variant["installed"] is False
    assert variant["needs_update"] is True
    assert variant["expected_version"] == "OpenVINO/whisper-medium-int8-ov"
    assert variant["size_bytes"] > 0


def test_model_status_defaults_to_light_manifest_and_allows_full_manifest(tmp_path: Path):
    app, headers = _app(tmp_path)
    app.state.services.settings.models_dir.mkdir(parents=True, exist_ok=True)
    (app.state.services.settings.models_dir / "manifest.json").write_text(
        json.dumps({"version": "test", "components": {}, "files": {"audio/model.bin": "hash"}}),
        encoding="utf-8",
    )

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            light = await client.get("/api/v1/models", headers=headers)
            assert light.status_code == 200
            assert "files" not in light.json()["manifest"]

            full = await client.get("/api/v1/models?summary=false", headers=headers)
            assert full.status_code == 200
            assert full.json()["manifest"]["files"] == {"audio/model.bin": "hash"}

    asyncio.run(scenario())


def test_manual_model_bundle_rejects_unsafe_archive(tmp_path: Path):
    app, _ = _app(tmp_path)
    archive_path = tmp_path / "unsafe.tar.gz"
    archive_path.write_bytes(_bundle(unsafe=True))

    try:
        app.state.services.models.install(archive_path=archive_path)
    except ValueError as exc:
        assert "不安全路径" in str(exc)
    else:
        raise AssertionError("unsafe archive must be rejected")
    assert not (tmp_path / "outside").exists()


def test_default_model_install_uses_official_repositories(tmp_path: Path, monkeypatch):
    app, _ = _app(tmp_path)
    manager = app.state.services.models
    rapid_root = tmp_path / "rapidocr"
    rapid_models = rapid_root / "models"
    rapid_models.mkdir(parents=True)
    for name in (
        "ch_PP-OCRv4_det_infer.onnx",
        "ch_PP-OCRv4_rec_infer.onnx",
        "ch_ppocr_mobile_v2.0_cls_infer.onnx",
    ):
        (rapid_models / name).write_bytes(name.encode())
    rapid_module = types.ModuleType("rapidocr_onnxruntime")
    rapid_module.__file__ = str(rapid_root / "__init__.py")
    hub_module = types.ModuleType("huggingface_hub")
    requests_module = types.ModuleType("requests")

    class Session:
        def __init__(self):
            self.trust_env = True
            self.proxies = {}

    def snapshot_download(repo_id, *, local_dir, **_):
        target = Path(local_dir)
        target.mkdir(parents=True, exist_ok=True)
        (target / "modules.json").write_text(
            json.dumps([{"idx": 0, "path": "", "type": "Transformer", "repo": repo_id}]),
            encoding="utf-8",
        )
        return str(target)

    hub_module.snapshot_download = snapshot_download
    hub_module.configure_http_backend = lambda **_: None
    requests_module.Session = Session
    monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", rapid_module)
    monkeypatch.setitem(sys.modules, "huggingface_hub", hub_module)
    monkeypatch.setitem(sys.modules, "requests", requests_module)

    def download_file(url, destination):
        destination.parent.mkdir(parents=True, exist_ok=True)
        if destination.suffix == ".xml":
            destination.write_text("<net><layers/></net>", encoding="utf-8")
        else:
            destination.write_bytes(b"\0" * (64 * 1024))

    monkeypatch.setattr(manager, "_download_file", download_file)

    def export_audio_model(_source, destination):
        destination.mkdir(parents=True, exist_ok=True)
        for name in (
            "config.json",
            "preprocessor_config.json",
            "tokenizer_config.json",
            "vocab.json",
            "openvino_encoder_model.xml",
            "openvino_encoder_model.bin",
            "openvino_decoder_model.xml",
            "openvino_decoder_model.bin",
            "openvino_decoder_with_past_model.xml",
            "openvino_decoder_with_past_model.bin",
        ):
            (destination / name).write_text("{}" if name.endswith(".json") else "model")

    monkeypatch.setattr(manager, "_export_audio_model", export_audio_model)
    monkeypatch.setattr(manager, "repair_component", lambda *_args, **_kwargs: {})
    progress = []

    result = manager.install(progress=lambda value, message: progress.append((value, message)))

    assert result["manifest"]["source"] == "official-repositories"
    assert (manager.settings.models_dir / "ocr" / "det.onnx").is_file()
    assert (manager.settings.models_dir / "faces" / "detector.xml").is_file()
    assert (manager.settings.models_dir / "audio" / "openvino_encoder_model.xml").is_file()
    assert any("多语言文本模型" in message for _, message in progress)


def test_component_install_preserves_other_components_and_manifest(tmp_path: Path, monkeypatch):
    app, _ = _app(tmp_path)
    manager = app.state.services.models
    root = manager.settings.models_dir
    (root / "faces").mkdir(parents=True)
    (root / "faces" / "sentinel.bin").write_bytes(b"existing-face-index-runtime")
    (root / "manifest.json").write_text(
        json.dumps(
            {
                "format_version": 2,
                "version": "existing",
                "components": {"faces": {"version": "faces-existing"}},
                "files": {"faces/sentinel.bin": hashlib.sha256(b"existing-face-index-runtime").hexdigest()},
            }
        ),
        encoding="utf-8",
    )
    rapid_root = tmp_path / "component-rapidocr"
    rapid_models = rapid_root / "models"
    rapid_models.mkdir(parents=True)
    for name in (
        "ch_PP-OCRv4_det_infer.onnx",
        "ch_PP-OCRv4_rec_infer.onnx",
        "ch_ppocr_mobile_v2.0_cls_infer.onnx",
    ):
        (rapid_models / name).write_bytes(name.encode())
    rapid_module = types.ModuleType("rapidocr_onnxruntime")
    rapid_module.__file__ = str(rapid_root / "__init__.py")
    monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", rapid_module)

    progress: list[tuple[float, str]] = []
    result = manager.install(component="ocr", progress=lambda value, message: progress.append((value, message)))

    assert result["component"] == "ocr"
    assert (root / "faces" / "sentinel.bin").read_bytes() == b"existing-face-index-runtime"
    assert (root / "ocr" / "det.onnx").is_file()
    manifest = json.loads((root / "manifest.json").read_text(encoding="utf-8"))
    assert manifest["components"]["faces"]["version"] == "faces-existing"
    assert manifest["components"]["ocr"]["version"] == "RapidOCR-PP-OCRv4"
    assert "faces/sentinel.bin" in manifest["files"]
    assert "ocr/det.onnx" in manifest["files"]
    assert any("启用" in message for _, message in progress)


def test_component_activation_rolls_back_directory_and_manifest(tmp_path: Path, monkeypatch):
    app, _ = _app(tmp_path)
    manager = app.state.services.models
    root = manager.settings.models_dir
    old = root / "ocr"
    old.mkdir(parents=True)
    (old / "sentinel.bin").write_bytes(b"old-component")
    original_manifest = {"version": "old", "components": {"ocr": {"version": "old-ocr"}}}
    (root / "manifest.json").write_text(json.dumps(original_manifest), encoding="utf-8")
    staged = tmp_path / "staged-ocr"
    staged.mkdir()
    for name in ("det.onnx", "rec.onnx", "cls.onnx"):
        (staged / name).write_bytes(b"new")

    monkeypatch.setattr(manager, "_write_manifest", lambda _manifest: (_ for _ in ()).throw(OSError("disk")))
    with pytest.raises(OSError, match="disk"):
        manager._activate_component("ocr", staged, {"version": "new-ocr"})

    assert (old / "sentinel.bin").read_bytes() == b"old-component"
    assert json.loads((root / "manifest.json").read_text(encoding="utf-8")) == original_manifest


def test_component_install_api_queues_missing_models_individually_and_reports_progress(tmp_path: Path):
    app, headers = _app(tmp_path)
    stale_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "stale"})

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.post("/api/v1/models/install", headers=headers, json={"component": "all"})
            assert response.status_code == 202
            job_ids = response.json()["job_ids"]
            assert response.json()["cancelled_audio_jobs"] == 1
            # The deterministic hash visual backend is already ready in tests.
            assert list(job_ids) == ["ocr", "faces", "audio"]
            assert response.json()["job_id"] == job_ids["ocr"]
            duplicate = await client.post("/api/v1/models/install", headers=headers, json={"component": "ocr"})
            assert duplicate.json()["job_ids"]["ocr"] == job_ids["ocr"]
            with app.state.services.db.transaction() as conn:
                conn.execute(
                    "UPDATE jobs SET status='running',progress=.42,message='正在复制 OCR 权重' WHERE id=?",
                    (job_ids["ocr"],),
                )
            status = (await client.get("/api/v1/models", headers=headers)).json()
            assert status["installations"]["ocr"] == {
                "job_id": job_ids["ocr"],
                "status": "running",
                "progress": 0.42,
                "message": "正在复制 OCR 权重",
                "error": None,
            }
            custom = await client.post(
                "/api/v1/models/install",
                headers=headers,
                json={"component": "faces", "url": "https://example.invalid/models.tar.gz"},
            )
            assert custom.status_code == 400
        with app.state.services.db.read() as conn:
            audio = conn.execute("SELECT status FROM jobs WHERE id=?", (stale_audio,)).fetchone()
        assert audio["status"] == "cancelled"

    asyncio.run(scenario())


def test_audio_medium_install_is_independent_and_does_not_cancel_active_audio_work(tmp_path: Path):
    app, headers = _app(tmp_path)
    active_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "active"})

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.post(
                "/api/v1/models/install",
                headers=headers,
                json={"component": "audio", "variant": "medium"},
            )
        assert response.status_code == 202
        assert response.json()["cancelled_audio_jobs"] == 0
        job_id = response.json()["job_ids"]["audio:medium"]
        with app.state.services.db.read() as conn:
            job = conn.execute("SELECT payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)).fetchone()
            active = conn.execute("SELECT status FROM jobs WHERE id=?", (active_audio,)).fetchone()
        assert json.loads(job["payload_json"]) == {"component": "audio", "variant": "medium"}
        assert job["dedupe_key"] == "install-model:audio:medium"
        assert active["status"] == "queued"

    asyncio.run(scenario())


def test_pending_audio_install_cancels_stale_audio_jobs_on_startup(tmp_path: Path):
    app, _ = _app(tmp_path)
    stale_audio = app.state.services.jobs.enqueue("transcribe_audio", {"video_id": "stale"})
    app.state.services.jobs.enqueue("install_models", {"component": "audio"}, dedupe_key="install-model:audio")

    assert app.state.services._cancel_audio_for_pending_model_install() == 1

    with app.state.services.db.read() as conn:
        row = conn.execute("SELECT status,message FROM jobs WHERE id=?", (stale_audio,)).fetchone()
    assert row["status"] == "cancelled"
    assert "模型安装优先" in row["message"]


def test_component_install_repairs_ready_runtime_without_redownloading_model(tmp_path: Path, monkeypatch):
    app, headers = _app(tmp_path)
    manager = app.state.services.models
    monkeypatch.setattr(
        manager,
        "status",
        lambda: {
            "visual_ready": True,
            "ocr_ready": False,
            "faces_ready": False,
            "audio_ready": False,
            "operational_components": {"visual": False},
        },
    )

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.post(
                "/api/v1/models/install",
                headers=headers,
                json={"component": "visual"},
            )
        assert response.status_code == 202
        job_id = response.json()["job_ids"]["visual"]
        with app.state.services.db.read() as connection:
            job = connection.execute(
                "SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?",
                (job_id,),
            ).fetchone()
        assert job["kind"] == "prepare_ai_runtime"
        assert json.loads(job["payload_json"]) == {"components": ["visual"]}
        assert job["dedupe_key"] == "prepare-ai-runtime:visual"

    asyncio.run(scenario())


def test_outdated_audio_repair_queues_real_model_reinstall(tmp_path: Path, monkeypatch):
    app, headers = _app(tmp_path)
    manager = app.state.services.models
    monkeypatch.setattr(
        manager,
        "status",
        lambda: {
            "visual_ready": False,
            "ocr_ready": False,
            "faces_ready": False,
            "audio_ready": True,
            "operational_components": {"audio": False},
            "component_health": {"audio": {"state": "outdated", "error": "需要转换为 FP16 模型"}},
        },
    )

    async def scenario():
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.post(
                "/api/v1/models/install",
                headers=headers,
                json={"component": "audio"},
            )
        assert response.status_code == 202
        job_id = response.json()["job_ids"]["audio"]
        with app.state.services.db.read() as connection:
            job = connection.execute("SELECT kind,payload_json,dedupe_key FROM jobs WHERE id=?", (job_id,)).fetchone()
        assert job["kind"] == "install_models"
        assert json.loads(job["payload_json"]) == {"component": "audio"}
        assert job["dedupe_key"] == "install-model:audio"

    asyncio.run(scenario())