# -*- coding: utf-8 -*-

import math
import argparse
import json
import os
import glob
import pandas as pd
from typing import Any, Dict, List, Tuple

def ts_to_seconds(ts: Dict[str, Any]) -> float:
    if ts is None:
        return math.nan
    return float(ts.get("seconds", 0) or 0) + float(ts.get("nanos", 0) or 0) / 1e9

def to_int_sec(x: float, mode: str = "floor") -> int:
    if math.isnan(x):
        return 0
    if mode == "ceil":
        return int(math.ceil(x))
    if mode == "round":
        return int(round(x))
    return int(math.floor(x))

def build_timestamp_list_from_shots(data: Dict[str, Any],
                                   rounding: str = "floor") -> List[List[int]]:
    ann = data.get("annotation_results", [])
    shots_list = (ann[0].get("shot_annotations") if len(ann) > 0 and isinstance(ann[0], dict) else []) or []

    segments: List[List[int]] = []
    prev_end: int | None = None

    for i, shot in enumerate(shots_list):
        raw_end = ts_to_seconds((shot or {}).get("end_time_offset"))
        if math.isnan(raw_end):
            continue

        end_i = to_int_sec(raw_end, mode=rounding)

        if prev_end is None:
            start_i = 1
        else:
            start_i = prev_end + 1

        if end_i <= start_i:
            end_i = start_i + 1

        segments.append([start_i, end_i])
        prev_end = end_i

    return segments

def collect_all_words(data: Dict[str, Any]) -> List[Dict[str, Any]]:
    ann = data.get("annotation_results", [])
    speech_list = (ann[1].get("speech_transcriptions") if len(ann) > 1 and isinstance(ann[1], dict) else []) or []

    all_words: List[Dict[str, Any]] = []
    for item in speech_list:
        alts = (item or {}).get("alternatives", []) or []
        if not alts:
            continue
        all_words.extend((alts[0] or {}).get("words", []) or [])
    return all_words

def assign_words_to_segments(all_words: List[Dict[str, Any]],
                             segments: List[List[int]],
                             mode: str = "overlap") -> List[Dict[str, Any]]:
    words: List[Tuple[float, float, str]] = []
    for w in all_words:
        s = ts_to_seconds(w.get("start_time"))
        e = ts_to_seconds(w.get("end_time"))
        txt = (w.get("word") or "").strip()
        if math.isnan(s) or math.isnan(e) or not txt:
            continue
        words.append((s, e, txt))

    words.sort(key=lambda x: x[0])

    out: List[Dict[str, Any]] = []
    i = 0
    n = len(words)

    for seg_start, seg_end in segments:
        seg_start_f = float(seg_start)
        seg_end_f = float(seg_end)

        while i < n and words[i][1] <= seg_start_f:
            i += 1

        j = i
        seg_words: List[str] = []
        while j < n and words[j][0] < seg_end_f:
            w_start, w_end, w_text = words[j]

            if mode == "contained":
                keep = (w_start >= seg_start_f) and (w_end <= seg_end_f)
            else:  
                keep = (w_start < seg_end_f) and (w_end > seg_start_f)

            if keep:
                seg_words.append(w_text)
            j += 1

        if not seg_words:
            continue

        out.append({
            "sentences": " ".join(seg_words),
            "start_time": int(seg_start),   
            "end_time": int(seg_end)
        })

    return out

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--json_dir",
        type=str,
        required=True,
        help="Directory containing <VIDEO_ID>.json files (or a glob like data/json/*.json)"
    )
    parser.add_argument(
        "--out_dir",
        type=str,
        required=True,
        help="Output directory where <VIDEO_ID>.csv will be written"
    )
    parser.add_argument(
        "--pattern",
        type=str,
        default="*.json",
        help="Filename pattern to match inside json_dir (default: *.json)"
    )

    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    # Allow either a directory OR a glob passed in json_dir
    if any(ch in args.json_dir for ch in ["*", "?", "["]):
        json_paths = sorted(glob.glob(args.json_dir))
    else:
        json_paths = sorted(glob.glob(os.path.join(args.json_dir, args.pattern)))

    if not json_paths:
        print(f"No JSON files found. json_dir={args.json_dir} pattern={args.pattern}")
        return

    print(f"Found {len(json_paths)} JSON files.")
    print(f"Writing CSVs to: {args.out_dir}")

    n_ok, n_fail = 0, 0

    for jp in json_paths:
        video_id = os.path.splitext(os.path.basename(jp))[0]
        out_csv = os.path.join(args.out_dir, f"{video_id}.csv")

        try:
            with open(jp, "r", encoding="utf-8") as f:
                data = json.load(f)

            timestamp_list = build_timestamp_list_from_shots(data)
            all_words = collect_all_words(data)
            seg_texts = assign_words_to_segments(all_words, timestamp_list, mode="overlap")

            if not seg_texts:
                print(f"[SKIP] {video_id}: no segments with words")
                continue

            df = pd.DataFrame(seg_texts)
            df.to_csv(out_csv, index=False)

            n_ok += 1
            print(f"[OK] {video_id} -> {out_csv}  (segments={len(seg_texts)})")

        except Exception as e:
            n_fail += 1
            print(f"[FAIL] {video_id} ({jp}): {e}")

    print(f"\nDone. Wrote {n_ok} CSVs. Failed {n_fail} files.")

if __name__ == "__main__":
    main()