#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ast
import json
import re
import subprocess
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Use Qwen3-VL to segment one video into consecutive events."
    )
    parser.add_argument("video", help="Path to the input video.")
    parser.add_argument(
        "--model-path",
        default="Qwen/Qwen3-VL-32B-Instruct",
        help="Local checkpoint path or Hugging Face model id.",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Optional JSON output path. If omitted, prints JSON to stdout.",
    )
    parser.add_argument("--fps", type=float, default=1.0, help="Video sampling FPS.")
    parser.add_argument(
        "--max-pixels",
        type=int,
        default=360000,
        help="Optional Qwen video processor max_pixels hint. Use 0 to disable.",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=1024,
        help="Max new tokens for the model response.",
    )
    parser.add_argument(
        "--device-map",
        default="auto",
        help="Transformers device_map argument.",
    )
    parser.add_argument(
        "--dtype",
        default="bfloat16",
        choices=["auto", "float16", "bfloat16", "float32"],
        help="Torch dtype for model loading.",
    )
    parser.add_argument(
        "--attn-implementation",
        default="",
        help='Optional attn_implementation, e.g. "flash_attention_2". Empty disables it.',
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Sampling temperature. 0 uses greedy decoding.",
    )
    parser.add_argument(
        "--ffprobe-bin",
        default="ffprobe",
        help="ffprobe executable used to read video duration.",
    )
    parser.set_defaults(trust_remote_code=True)
    parser.add_argument(
        "--no-trust-remote-code",
        dest="trust_remote_code",
        action="store_false",
        help="Disable trust_remote_code for model and processor loading.",
    )
    return parser.parse_args()


def probe_duration(video_path: Path, ffprobe_bin: str) -> float | None:
    try:
        result = subprocess.run(
            [
                ffprobe_bin,
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(video_path),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    except (FileNotFoundError, subprocess.CalledProcessError):
        return None
    try:
        duration = float(result.stdout.strip())
    except ValueError:
        return None
    return round(duration, 3) if duration > 0 else None


def extract_json_text(raw_text: str) -> str:
    text = raw_text.strip()
    fenced = re.findall(r"```(?:json)?\s*(.*?)```", text, flags=re.DOTALL | re.I)
    candidates = [item.strip() for item in fenced if item.strip()]
    candidates.append(text)
    for candidate in candidates:
        if candidate.startswith("{") and candidate.endswith("}"):
            return candidate

    start = text.find("{")
    if start < 0:
        raise ValueError(f"No JSON object found in model output: {text[:300]}")

    depth = 0
    in_string = False
    escape = False
    for index in range(start, len(text)):
        char = text[index]
        if in_string:
            if escape:
                escape = False
            elif char == "\\":
                escape = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]
    raise ValueError(f"Unclosed JSON object in model output: {text[:300]}")


def parse_json_payload(raw_text: str) -> dict[str, Any]:
    json_text = extract_json_text(raw_text)
    try:
        payload = json.loads(json_text)
    except json.JSONDecodeError:
        payload = ast.literal_eval(json_text)
    if not isinstance(payload, dict):
        raise ValueError(f"Expected a JSON object, got {type(payload).__name__}.")
    return payload


def coerce_float(value: Any, field_name: str) -> float:
    if isinstance(value, str):
        value = value.strip().rstrip("s")
    try:
        return round(float(value), 3)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"Invalid {field_name}: {value!r}") from exc


def build_prompt(duration: float | None) -> str:
    duration_line = (
        f"The video duration is approximately {duration:.3f} seconds."
        if duration is not None
        else "The video duration is not provided; infer timestamps in seconds from the video timeline."
    )
    return f"""
You are segmenting a single video into consecutive semantic events.

{duration_line}

Definition of an event:
- An event is a continuous temporal interval with one dominant visible action, interaction, camera-framing state, or meaningful scene/object state.
- Split when the dominant action changes, a new subject/object appears or disappears, the focus changes materially, or the scene state changes enough that a separate event caption would be needed.
- Do not over-segment tiny pose changes, ordinary motion within the same action, or brief camera shake.
- Do not leave gaps between events.
- Events must be chronological and non-overlapping.
- The first event should begin at 0.0 seconds.
- If the duration is known, the last event should end at the video duration.
- Use seconds as numbers, not strings.
- Include a short description for each event so the boundary decision is inspectable.

Return strict JSON only. Do not use markdown fences.

Return JSON in exactly this schema:
{{
  "event_count": 3,
  "events": [
    {{
      "event_id": 1,
      "begin_time": 0.0,
      "end_time": 2.5,
      "description": "The first visible event in one concise sentence."
    }},
    {{
      "event_id": 2,
      "begin_time": 2.5,
      "end_time": 5.0,
      "description": "The second visible event in one concise sentence."
    }}
  ]
}}
""".strip()


def normalize_events(payload: dict[str, Any], duration: float | None) -> dict[str, Any]:
    raw_events = payload.get("events")
    if not isinstance(raw_events, list) or not raw_events:
        raise ValueError("Model output must contain a non-empty events list.")

    events: list[dict[str, Any]] = []
    for index, item in enumerate(raw_events, start=1):
        if not isinstance(item, dict):
            raise ValueError(f"Event {index} is not an object.")
        begin_time = coerce_float(item.get("begin_time"), "begin_time")
        end_time = coerce_float(item.get("end_time"), "end_time")
        if duration is not None:
            begin_time = min(max(begin_time, 0.0), duration)
            end_time = min(max(end_time, 0.0), duration)
        if end_time <= begin_time:
            raise ValueError(
                f"Event {index} has non-positive duration: {begin_time} to {end_time}"
            )
        description = re.sub(
            r"\s+", " ", str(item.get("description", "")).strip()
        )
        events.append(
            {
                "event_id": index,
                "begin_time": begin_time,
                "end_time": end_time,
                "description": description,
            }
        )

    events.sort(key=lambda item: (item["begin_time"], item["end_time"]))
    for index, event in enumerate(events, start=1):
        event["event_id"] = index

    if events[0]["begin_time"] != 0.0:
        events[0]["begin_time"] = 0.0
    if duration is not None:
        events[-1]["end_time"] = duration

    return {
        "event_count": len(events),
        "events": events,
    }


class Qwen3VLRunner:
    def __init__(
        self,
        model_path: str,
        device_map: str,
        dtype: str,
        attn_implementation: str,
        trust_remote_code: bool,
        fps: float,
        max_pixels: int,
        temperature: float,
    ) -> None:
        try:
            import torch
            from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "Missing dependencies. Install torch, transformers, and preferably qwen-vl-utils."
            ) from exc

        self.torch = torch
        self.temperature = temperature
        self.fps = fps if fps > 0 else None
        self.max_pixels = max_pixels if max_pixels > 0 else None

        try:
            from qwen_vl_utils import process_vision_info  # type: ignore
        except ModuleNotFoundError:
            process_vision_info = None
        self.process_vision_info = process_vision_info

        dtype_map = {
            "auto": "auto",
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
        }
        load_kwargs: dict[str, Any] = {
            "device_map": device_map,
            "torch_dtype": dtype_map[dtype],
            "trust_remote_code": trust_remote_code,
        }
        if attn_implementation:
            load_kwargs["attn_implementation"] = attn_implementation

        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
            model_path,
            **load_kwargs,
        )
        self.processor = AutoProcessor.from_pretrained(
            model_path,
            trust_remote_code=trust_remote_code,
        )

    def _build_messages(self, video_path: Path, prompt: str) -> list[dict[str, Any]]:
        video_item: dict[str, Any] = {
            "type": "video",
            "video": str(video_path.resolve()),
        }
        if self.fps is not None:
            video_item["fps"] = self.fps
        if self.max_pixels is not None:
            video_item["max_pixels"] = self.max_pixels
        return [
            {
                "role": "system",
                "content": "You are a careful video event segmentation assistant. Return strict JSON only.",
            },
            {
                "role": "user",
                "content": [video_item, {"type": "text", "text": prompt}],
            },
        ]

    def _prepare_inputs(self, messages: list[dict[str, Any]]):
        if self.process_vision_info is not None:
            text = self.processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
            image_inputs, video_inputs = self.process_vision_info(messages)
            inputs = self.processor(
                text=[text],
                images=image_inputs,
                videos=video_inputs,
                padding=True,
                return_tensors="pt",
            )
        else:
            try:
                inputs = self.processor.apply_chat_template(
                    messages,
                    tokenize=True,
                    add_generation_prompt=True,
                    return_dict=True,
                    return_tensors="pt",
                )
            except Exception as exc:
                raise RuntimeError(
                    "Failed to prepare video inputs. Install qwen-vl-utils or use a transformers build with Qwen3-VL video chat support."
                ) from exc
        return inputs.to(self.model.device)

    def segment(self, video_path: Path, prompt: str, max_new_tokens: int) -> dict[str, Any]:
        messages = self._build_messages(video_path, prompt)
        inputs = self._prepare_inputs(messages)
        generate_kwargs: dict[str, Any] = {"max_new_tokens": max_new_tokens}
        if self.temperature and self.temperature > 0:
            generate_kwargs.update({"do_sample": True, "temperature": self.temperature})
        else:
            generate_kwargs["do_sample"] = False

        with self.torch.inference_mode():
            generated_ids = self.model.generate(**inputs, **generate_kwargs)
        generated_ids_trimmed = [
            out_ids[len(in_ids) :]
            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        raw_text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]
        return parse_json_payload(raw_text)


def main() -> None:
    args = parse_args()
    video_path = Path(args.video).expanduser().resolve()
    if not video_path.is_file():
        raise FileNotFoundError(f"Video not found: {video_path}")

    duration = probe_duration(video_path, args.ffprobe_bin)
    runner = Qwen3VLRunner(
        model_path=args.model_path,
        device_map=args.device_map,
        dtype=args.dtype,
        attn_implementation=args.attn_implementation,
        trust_remote_code=args.trust_remote_code,
        fps=args.fps,
        max_pixels=args.max_pixels,
        temperature=args.temperature,
    )
    payload = runner.segment(
        video_path=video_path,
        prompt=build_prompt(duration),
        max_new_tokens=args.max_new_tokens,
    )
    result = normalize_events(payload, duration)

    output_text = json.dumps(result, ensure_ascii=False, indent=2) + "\n"
    if args.output:
        output_path = Path(args.output).expanduser().resolve()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(output_text, encoding="utf-8")
    else:
        print(output_text, end="")


if __name__ == "__main__":
    main()