from __future__ import annotations

import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR / "vllm"))

from vllm_caption_runtime import build_vllm_request, load_llm, load_processor, sampling_params


# =============================================================================
# Paths
# =============================================================================

INPUT_FOLDER = "/workspace/videos_to_caption"
OUTPUT_FOLDER = "/workspace/captions_out"
MODEL_DIR = str(SCRIPT_DIR / "model")

RECURSIVE = True
OVERWRITE_EXISTING = False
OUTPUT_EXTENSION = ".txt"
ERROR_EXTENSION = ".error.txt"
VIDEO_EXTENSIONS = (".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v")


# =============================================================================
# Prompt Settings
# =============================================================================

PROMPT_OVERRIDE = ""
CAPTION_LENGTH = "very large"
INCLUDE_WATERMARK_INFO = False
HAS_THINKING = True

VULGARITY = "low"
UNCERTAINTY = "low"
CHARACTER_NAMES = "none"
FLUFF = "none"
HAS_REPETITION = False
SPECULATION = "low"
TEMPORAL_DETAIL = "medium"
VISUAL_SPECIFICITY = "moderate"
CAMERA_DETAIL = "medium"
CAPTION_STYLE = "plain"


# =============================================================================
# vLLM / Generation Hyperparameters
# =============================================================================

NUM_FRAMES = 12
SAMPLING_RATE = 16_000
MAX_MODEL_LEN = 4096
MAX_NUM_SEQS = 20
BATCH_SIZE = 20
GPU_MEMORY_UTILIZATION = 0.88
DTYPE = "bfloat16"
# Online vLLM FP8 for Gemma linear weights. Parakeet is left unchanged.
USE_FP8 = True
FP8_QUANTIZATION = "fp8_per_tensor"
ENFORCE_EAGER = False
ENABLE_PREFIX_CACHING = False
TRUST_REMOTE_CODE = False

MAX_TOKENS = 1200
TEMPERATURE = 0.0
TOP_P = 0.9
REPETITION_PENALTY = 1.1


def prompt_settings() -> dict[str, object]:
    return {
        "caption_length": CAPTION_LENGTH,
        "include_watermark_info": INCLUDE_WATERMARK_INFO,
        "has_thinking": HAS_THINKING,
        "vulgarity": VULGARITY,
        "uncertainty": UNCERTAINTY,
        "character_names": CHARACTER_NAMES,
        "fluff": FLUFF,
        "has_repetition": HAS_REPETITION,
        "speculation": SPECULATION,
        "temporal_detail": TEMPORAL_DETAIL,
        "visual_specificity": VISUAL_SPECIFICITY,
        "camera_detail": CAMERA_DETAIL,
        "caption_style": CAPTION_STYLE,
    }


def find_videos(input_folder: Path) -> list[Path]:
    iterator = input_folder.rglob("*") if RECURSIVE else input_folder.glob("*")
    videos = [
        path
        for path in iterator
        if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS
    ]
    return sorted(videos)


def output_path_for(video_path: Path, input_folder: Path, output_folder: Path) -> Path:
    relative = video_path.relative_to(input_folder)
    return (output_folder / relative).with_suffix(OUTPUT_EXTENSION)


def write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text.rstrip() + "\n", encoding="utf-8")


def chunks(items: list[Path], size: int) -> list[list[Path]]:
    return [items[index : index + size] for index in range(0, len(items), size)]


def find_json_end(text: str) -> int | None:
    if not text.startswith("{"):
        return None

    depth = 0
    in_string = False
    escape = False
    for index, char in enumerate(text):
        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 index + 1
    return None


def strip_caption_label(text: str) -> str:
    caption = text.strip()
    lowered = caption.lower()
    for label in ("final caption", "final_caption", "caption", "final"):
        if lowered == label:
            return ""
        if lowered.startswith(label):
            rest = caption[len(label) :]
            stripped = rest.lstrip()
            if stripped.startswith(":"):
                return stripped[1:].lstrip()
            if stripped.startswith("- "):
                return stripped[1:].lstrip()
            if rest[:1].isspace():
                return rest.lstrip()
    return caption


def final_caption_from_raw_output(text: str) -> str:
    raw = text.strip()
    lowered = raw.lower()
    for marker in ("thought", "thinking", "reasoning"):
        if lowered == marker:
            return ""
        if lowered.startswith(marker):
            rest = raw[len(marker) :].lstrip()
            if rest.startswith(":"):
                rest = rest[1:].lstrip()
            raw = rest
            break

    json_end = find_json_end(raw)
    if json_end is not None:
        return strip_caption_label(raw[json_end:])
    return strip_caption_label(raw)


def main() -> None:
    quantization = FP8_QUANTIZATION if USE_FP8 else None
    input_folder = Path(INPUT_FOLDER)
    output_folder = Path(OUTPUT_FOLDER)
    if not input_folder.is_dir():
        raise RuntimeError(f"INPUT_FOLDER does not exist or is not a directory: {input_folder}")

    videos = find_videos(input_folder)
    if not videos:
        print(f"No videos found in {input_folder}")
        return

    processor = load_processor(MODEL_DIR)
    llm = load_llm(
        model_dir=MODEL_DIR,
        max_model_len=MAX_MODEL_LEN,
        max_num_seqs=MAX_NUM_SEQS,
        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
        dtype=DTYPE,
        quantization=quantization,
        enforce_eager=ENFORCE_EAGER,
        enable_prefix_caching=ENABLE_PREFIX_CACHING,
        trust_remote_code=TRUST_REMOTE_CODE,
    )
    params = sampling_params(
        temperature=TEMPERATURE,
        max_tokens=MAX_TOKENS,
        top_p=TOP_P,
        repetition_penalty=REPETITION_PENALTY,
    )

    done = 0
    skipped = 0
    failed = 0
    settings = prompt_settings()

    for batch in chunks(videos, BATCH_SIZE):
        requests = []
        request_videos = []
        for video_path in batch:
            output_path = output_path_for(video_path, input_folder, output_folder)
            if output_path.exists() and not OVERWRITE_EXISTING:
                skipped += 1
                continue
            try:
                requests.append(
                    build_vllm_request(
                        processor=processor,
                        model_dir=MODEL_DIR,
                        video_path=video_path,
                        num_frames=NUM_FRAMES,
                        sampling_rate=SAMPLING_RATE,
                        prompt_override=PROMPT_OVERRIDE,
                        prompt_settings=settings,
                    )
                )
                request_videos.append(video_path)
            except Exception as exc:
                failed += 1
                error_path = output_path.with_suffix(ERROR_EXTENSION)
                write_text(error_path, f"{type(exc).__name__}: {exc}")
                print(f"FAILED preprocess {video_path}: {exc}")

        if not requests:
            continue

        try:
            outputs = llm.generate(requests, sampling_params=params)
        except Exception as exc:
            failed += len(request_videos)
            for video_path in request_videos:
                output_path = output_path_for(video_path, input_folder, output_folder)
                error_path = output_path.with_suffix(ERROR_EXTENSION)
                write_text(error_path, f"{type(exc).__name__}: {exc}")
            print(f"FAILED batch starting {request_videos[0]}: {exc}")
            continue

        for video_path, output in zip(request_videos, outputs, strict=True):
            output_path = output_path_for(video_path, input_folder, output_folder)
            raw_text = output.outputs[0].text if output.outputs else ""
            text = final_caption_from_raw_output(raw_text)
            write_text(output_path, text)
            done += 1
            print(f"WROTE {output_path}")

    print(f"Done. wrote={done} skipped={skipped} failed={failed}")


if __name__ == "__main__":
    main()