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

import json
import logging
import re
import sys
from pathlib import Path

from _gemini_runner import GeminiRunner, load_repo_module


BASE = load_repo_module("qwen_multi_event_base", "qwen/generate_multi_event_annotations.py")


def strip_detailed_caption_from_global_prompt(prompt: str) -> str:
    cleaned = prompt
    cleaned = cleaned.replace(
        '- The detailed caption should be about 500-800 words.\n',
        "",
    )
    cleaned = cleaned.replace(
        '- The detailed caption may synthesize and reuse the structured fields plus stable recurring evidence from the event summaries, but it must not turn into a segment-by-segment narration.\n',
        "",
    )
    cleaned = re.sub(
        r',\n(\s*)"detailed_caption": "About 500-800 words\."\n',
        "\n",
        cleaned,
    )
    return cleaned


_ORIGINAL_BUILD_GLOBAL_CONTEXT_PROMPT = BASE.build_global_context_prompt
_ORIGINAL_BUILD_GLOBAL_REFINE_PROMPT = BASE.build_global_refine_prompt
_ORIGINAL_COERCE_GLOBAL_PAYLOAD = BASE.coerce_global_payload


def _gemini_build_global_context_prompt(*args, **kwargs):
    prompt = _ORIGINAL_BUILD_GLOBAL_CONTEXT_PROMPT(*args, **kwargs)
    return (
        f"{prompt}\n\n"
        "IMPORTANT FOR GEMINI:\n"
        "- Return one single JSON object for the entire stitched video.\n"
        "- Never return a JSON array.\n"
        "- Never return one object per clip, per segment, or per event.\n"
        "- The stitched video may contain multiple chronological event clips, but this context pass must summarize the whole stitched video into one object only.\n"
        "- Do not return object-detection output, tracking output, grounding output, OCR output, or visual tool output.\n"
        "- Never return keys such as box_2d, bbox, label, labels, mask, polygon, points, score, or tracking_id unless the requested schema explicitly asks for them.\n"
        "- Never return timestamped clip-local records such as event_clip_id, start_time_seconds, or end_time_seconds in this context pass.\n"
        "- If uncertain, still fill the requested schema directly with high-confidence whole-video annotations rather than switching to detection-style results.\n"
    )


def _gemini_build_global_refine_prompt(*args, **kwargs):
    return strip_detailed_caption_from_global_prompt(
        _ORIGINAL_BUILD_GLOBAL_REFINE_PROMPT(*args, **kwargs)
    )


def _gemini_coerce_global_payload(*args, **kwargs):
    payload = _ORIGINAL_COERCE_GLOBAL_PAYLOAD(*args, **kwargs)
    payload.pop("detailed_caption", None)
    return payload


BASE.build_global_context_prompt = _gemini_build_global_context_prompt
BASE.build_global_refine_prompt = _gemini_build_global_refine_prompt
BASE.coerce_global_payload = _gemini_coerce_global_payload


def main() -> int:
    args = BASE.parse_args()
    BASE.setup_logging(args.log_level)

    input_path = Path(args.input).expanduser().resolve()
    output_path = Path(args.output).expanduser().resolve()
    camera_motion_guide_path = Path(args.camera_motion_guide).expanduser().resolve()

    suffixes = BASE.parse_video_suffixes(args.video_suffixes)
    allow_uuids = BASE.load_uuid_allow_list(args.uuids, args.uuids_file)
    camera_motion_guide = BASE.load_camera_motion_guide(camera_motion_guide_path)

    groups = BASE.collect_segment_groups(
        input_path=input_path,
        recursive=args.recursive,
        suffixes=suffixes,
        allow_uuids=allow_uuids,
    )
    if not groups:
        raise RuntimeError("No valid segmented videos were found.")

    uuids = list(groups)
    if args.limit is not None:
        uuids = uuids[: args.limit]

    pending_for_print = uuids
    if args.resume and output_path.exists():
        existing_videos_for_print = BASE.load_existing_output(
            output_path=output_path,
            overwrite=False,
            resume=True,
        )
        pending_for_print = [uuid for uuid in uuids if uuid not in existing_videos_for_print]

    if args.print_uuids_json:
        print(json.dumps(pending_for_print, ensure_ascii=False))
        return 0

    uuids = BASE.apply_uuid_shard(
        uuids=pending_for_print if args.resume else uuids,
        num_shards=args.num_shards,
        shard_index=args.shard_index,
    )

    logging.info("Discovered %d UUID groups.", len(uuids))
    for uuid in uuids:
        logging.info("UUID %s has %d segments.", uuid, len(groups[uuid]))

    if args.dry_run:
        logging.info("Dry run complete. Gemini model was not called.")
        return 0

    existing_videos = BASE.load_existing_output(
        output_path=output_path,
        overwrite=args.overwrite,
        resume=args.resume,
    )
    pending_uuids = [uuid for uuid in uuids if uuid not in existing_videos]
    if not pending_uuids:
        logging.info("Nothing to do. All requested UUIDs already exist in %s.", output_path)
        return 0

    runner = GeminiRunner(
        model_path=Path(args.model_path).expanduser().resolve(),
        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,
        top_p=args.top_p,
        top_k=args.top_k,
        repetition_penalty=args.repetition_penalty,
        greedy=args.greedy,
    )

    try:
        failures: list[str] = []
        for index, uuid in enumerate(pending_uuids, start=1):
            logging.info("Processing UUID %s (%d/%d)", uuid, index, len(pending_uuids))
            try:
                video_payload = BASE.annotate_uuid(
                    uuid=uuid,
                    segments=groups[uuid],
                    runner=runner,
                    camera_motion_guide=camera_motion_guide,
                    args=args,
                )
                existing_videos[uuid] = video_payload
                BASE.write_output(output_path, existing_videos)
                logging.info("Wrote %s", output_path)
            except Exception as exc:
                if not args.continue_on_error:
                    raise
                logging.exception("Failed to process %s: %s", uuid, exc)
                failures.append(uuid)

        if failures:
            logging.error("Completed with failures: %s", ", ".join(failures))
            return 1
        logging.info("Finished %d UUIDs.", len(pending_uuids))
        return 0
    finally:
        runner.close()


if __name__ == "__main__":
    sys.exit(main())