#!/usr/bin/env python3
"""Unit and CPU smoke checks for fixed-frame SegmentQA preprocessing."""

from __future__ import annotations

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

import av

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "scripts"))

import preprocess_segmentqa_fixed_frames as prep  # noqa: E402


PYTHON = "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/conda-envs/medrl/bin/python"
DATA_ROOT = "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge"
OUT = PROJECT_ROOT / "tmp/segmentqa_fixed80_preprocess_smoke_20260826"


def test_env() -> dict[str, str]:
    env = dict(os.environ)
    env["PYTHONPATH"] = str(PROJECT_ROOT / "ms-swift")
    return env


def video_stats(path: Path) -> dict:
    with av.open(str(path)) as container:
        stream = container.streams.video[0]
        pts = []
        dims = set()
        for frame in container.decode(stream):
            pts.append(float(frame.pts * frame.time_base))
            dims.add((frame.width, frame.height))
        return {"count": len(pts), "pts": pts, "dims": dims}


def test_time_and_sampling_helpers() -> None:
    assert prep.parse_time("01:02:03") == 3723
    assert prep.parse_time(12.5) == 12.5
    targets = prep.target_times(10, 18, 4)
    assert targets == [11, 13, 15, 17]
    assert prep.metadata_from_pts([10.1234, 10.5678], 420, 224) == {
        "fps": 1000,
        "frames_indices": [10123, 10568],
        "total_num_frames": 10569,
        "width": 420,
        "height": 224,
        "duration": 10.569,
        "video_backend": "precomputed_absolute_pts",
    }


def test_template_metadata_override_path() -> None:
    from swift.template.templates.qwen import _extract_precomputed_absolute_video_metadata as extract_metadata

    metadata = prep.metadata_from_pts([1.0, 2.0], 420, 224)
    video_inputs = {"video": "clip.mp4", "nframes": 80, "enable_thinking": False, "video_metadata": metadata}
    extracted = extract_metadata(video_inputs, 0)
    assert extracted == metadata
    assert video_inputs == {"video": "clip.mp4", "nframes": 80, "enable_thinking": False}
    legacy_inputs = {"video": "clip.mp4", "fps": 1.0, "min_frames": 16, "max_frames": 64}
    assert extract_metadata(legacy_inputs, 0) is None
    assert legacy_inputs == {"video": "clip.mp4", "fps": 1.0, "min_frames": 16, "max_frames": 64}


def validate_qwen_fetch_video_path(record: dict) -> None:
    from qwen_vl_utils import fetch_video
    from swift.template.templates.qwen import _extract_precomputed_absolute_video_metadata
    from transformers.video_utils import VideoMetadata

    video_inputs = {"video": record["videos"][0], **record["chat_template_kwargs"]}
    metadata = _extract_precomputed_absolute_video_metadata(video_inputs, 0)
    assert metadata is not None
    assert "do_sample_frames" not in metadata
    VideoMetadata(**metadata)
    (video, fetched_metadata), sample_fps = fetch_video(
        video_inputs, return_video_sample_fps=True, return_video_metadata=True, image_patch_size=16)
    assert video.shape[0] == 80
    assert video.shape[1] == 3
    assert video.shape[-2] % 32 == 0
    assert video.shape[-1] % 32 == 0
    assert fetched_metadata["total_num_frames"] >= 80
    assert sample_fps > 0
    assert video_inputs == {"video": record["videos"][0], "nframes": 80, "enable_thinking": False}


def assert_author_direct_segment_prompt(prompt: str) -> None:
    assert "timestamp in the source video" in prompt
    assert "<1234.5 seconds> is 00:20:34" in prompt
    assert "START OF THE VIDEO, not from the start of this" in prompt
    assert (
        "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, "
        "Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent."
    ) in prompt
    assert "Reply with the answer and nothing else -- no reasoning, no preamble, no" in prompt
    assert "A single short line." in prompt


def test_cpu_smoke_preprocess_train() -> None:
    if OUT.exists():
        shutil.rmtree(OUT)
    cmd = [
        PYTHON,
        "scripts/preprocess_segmentqa_fixed_frames.py",
        "--data-root",
        DATA_ROOT,
        "--output-dir",
        str(OUT),
        "--splits",
        "train",
        "--frames",
        "80",
        "--width",
        "420",
        "--height",
        "224",
        "--workers",
        "1",
        "--limit",
        "2",
    ]
    result = subprocess.run(cmd, cwd=PROJECT_ROOT, check=True, text=True, capture_output=True, env=test_env())
    summary = json.loads(result.stdout.strip().splitlines()[-1])
    assert summary["segments"] == 2
    assert summary["success"] == 2
    assert summary["failed"] == 0

    manifest = [json.loads(line) for line in (OUT / "segment_manifest.jsonl").read_text().splitlines()]
    assert {item["dataset"] for item in manifest} == set(prep.DATASET_NAMES)
    for item in manifest:
        stats = video_stats(OUT / item["relpath"])
        assert stats["count"] == 80
        assert stats["dims"] == {(420, 224)}
        assert item["verified_frame_count"] == 80
        assert item["width"] == 420
        assert item["height"] == 224
        assert len(item["selected_source_pts"]) == 80
        assert all(item["start_seconds"] <= pts < item["end_seconds"] for pts in item["selected_source_pts"])
        assert item["video_metadata"]["fps"] == 1000
        assert item["video_metadata"]["frames_indices"] == [round(pts * 1000) for pts in item["selected_source_pts"]]
        assert item["video_metadata"]["total_num_frames"] == max(item["video_metadata"]["frames_indices"]) + 1
        assert item["video_metadata"]["video_backend"] == "precomputed_absolute_pts"
        assert "do_sample_frames" not in item["video_metadata"]

    train_rows = [json.loads(line) for line in (OUT / "train.jsonl").read_text().splitlines()]
    assert len(train_rows) >= 2
    assert all(row["messages"][0]["role"] == "system" for row in train_rows)
    assert_author_direct_segment_prompt(train_rows[0]["messages"][0]["content"])
    assert all(row["messages"][1]["content"].startswith("<video>\n") for row in train_rows)
    assert all(row["messages"][2]["role"] == "assistant" and row["messages"][2]["content"] for row in train_rows)
    assert all(row["chat_template_kwargs"]["nframes"] == 80 for row in train_rows)
    assert all(row["chat_template_kwargs"]["enable_thinking"] is False for row in train_rows)
    validate_qwen_fetch_video_path(train_rows[0])
    assert not (OUT / "test_requests.jsonl").exists()
    report = json.loads((OUT / "validation_report.json").read_text())
    assert report["status"] == "pass"
    assert report["success"] == report["segments"] == 2

    before = {p: p.stat().st_mtime for p in OUT.glob("videos/train/*/*.mp4")}
    resume = subprocess.run(cmd + ["--resume"], cwd=PROJECT_ROOT, check=True, text=True, capture_output=True, env=test_env())
    resume_summary = json.loads(resume.stdout.strip().splitlines()[-1])
    assert resume_summary["success"] == 2
    assert before == {p: p.stat().st_mtime for p in OUT.glob("videos/train/*/*.mp4")}


if __name__ == "__main__":
    test_time_and_sampling_helpers()
    test_template_metadata_override_path()
    test_cpu_smoke_preprocess_train()