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

import pandas as pd
import os, traceback
import json
from PIL import Image
import argparse
import torch
from transformers import AutoTokenizer, AutoModel, AutoImageProcessor

CACHE_DIR = "ENCODERS/hf_cache"
os.makedirs(CACHE_DIR, exist_ok=True)

# Defaults = "Main/Intermediate" encoders
DEFAULT_TEXT_MODEL = "bert-large-uncased"
DEFAULT_VISUAL_MODEL = "google/vit-base-patch16-224-in21k"


def safe_name(model_id: str) -> str:
    return model_id.replace("/", "__").replace(":", "_")


def choose_output_root(text_model: str, visual_model: str) -> str:
    if text_model == DEFAULT_TEXT_MODEL and visual_model == DEFAULT_VISUAL_MODEL:
        return os.path.join("Embeddings", "Intermediate")
    return os.path.join("Embeddings", "Alternate")


def get_max_len(tokenizer, requested: int = 1024) -> int:
    try:
        return min(int(tokenizer.model_max_length), int(requested))
    except Exception:
        return requested


def load_models(text_model_name: str, visual_model_name: str, cache_dir: str, local_only: bool = True):
    try:
        tokenizer = AutoTokenizer.from_pretrained(text_model_name, cache_dir=cache_dir, local_files_only=local_only)
        text_model = AutoModel.from_pretrained(text_model_name, cache_dir=cache_dir, local_files_only=local_only)
        visual_processor = AutoImageProcessor.from_pretrained(visual_model_name, cache_dir=cache_dir, local_files_only=local_only)
        visual_model = AutoModel.from_pretrained(visual_model_name, cache_dir=cache_dir, local_files_only=local_only)

        text_model.eval()
        visual_model.eval()

        print("[OK] Text hidden_size:", getattr(text_model.config, "hidden_size", "NA"))
        print("[OK] Vision hidden_size:", getattr(visual_model.config, "hidden_size", "NA"))

        return tokenizer, text_model, visual_processor, visual_model

    except Exception as e:
        print("ERROR loading models:", repr(e))
        traceback.print_exc()
        raise


def save_text_features(text, save_dir, index, tokenizer, text_model, max_length: int = 1024):
    special_characters = ",!?...-_*@#_/\\\"'():;~|^=[]{}<>+=`"
    cleaned_sentence = str(text).translate(str.maketrans("", "", special_characters))

    effective_max_len = get_max_len(tokenizer, max_length)
    inputs = tokenizer(cleaned_sentence, max_length=effective_max_len, truncation=True, return_tensors="pt")

    with torch.no_grad():
        outputs = text_model(**inputs)

    hidden_states = outputs.last_hidden_state
    torch.save(hidden_states, os.path.join(save_dir, f"text_embed_{index}.pt"))

    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze().tolist())
    with open(os.path.join(save_dir, f"text_tokens_{index}.json"), "w", encoding="utf-8") as f:
        json.dump(tokens, f)


def save_frame_features(image_path, save_dir, frame_idx, visual_processor, visual_model):
    save_file = os.path.join(save_dir, f"frame_embed_{frame_idx}.pt")
    if os.path.exists(save_file):
        return

    if not os.path.exists(image_path):
        return

    try:
        image = Image.open(image_path).convert("RGB")
        inputs = visual_processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = visual_model(**inputs)
        torch.save(outputs.last_hidden_state, save_file)
    except Exception as e:
        print(f"Error processing {image_path}: {e}")


def process_video_features(csv_path, frames_root, embedding_root, tokenizer, text_model, visual_processor, visual_model):
    file_mod = os.path.basename(csv_path).split(".")[0]
    df = pd.read_csv(csv_path)

    vid_embed_path = os.path.join(embedding_root, file_mod)
    text_embed_path = os.path.join(vid_embed_path, "Text")
    frame_embed_path = os.path.join(vid_embed_path, "Frames")

    os.makedirs(text_embed_path, exist_ok=True)
    os.makedirs(frame_embed_path, exist_ok=True)

    print(f"Processing {file_mod}...")

    for i in range(len(df)):
        if "sentences" not in df.columns:
            raise ValueError(f"'sentences' column missing in {csv_path}")

        if pd.isna(df.loc[i, "sentences"]) or len(str(df.loc[i, "sentences"])) == 0:
            continue

        # 1) Text
        sentence = str(df.loc[i, "sentences"])
        save_text_features(sentence, text_embed_path, i, tokenizer, text_model, max_length=1024)

        # 2) Frames
        start = int(0 if pd.isna(df.loc[i, "start_time"]) else df.loc[i, "start_time"])
        end = int(0 if pd.isna(df.loc[i, "end_time"]) else df.loc[i, "end_time"])
        if end == start:
            end += 1
        end += 1

        frame_dir = os.path.join(frames_root, file_mod)
        for frame_idx in range(start, end):
            img_path = os.path.join(frame_dir, f"frame{frame_idx}.jpg")
            save_frame_features(img_path, frame_embed_path, frame_idx, visual_processor, visual_model)


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--manifest", type=str, required=True,
                        help="Path to video_manifest.txt (one CSV path per line)")
    parser.add_argument("--video_root", type=str, default="Video_Frames",
                        help="Root directory containing extracted video frames")

    parser.add_argument("--text_model", type=str, default=DEFAULT_TEXT_MODEL,
                        help=f"Text encoder HF id (default: {DEFAULT_TEXT_MODEL})")
    parser.add_argument("--visual_model", type=str, default=DEFAULT_VISUAL_MODEL,
                        help=f"Vision encoder HF id (default: {DEFAULT_VISUAL_MODEL})")

    parser.add_argument("--intermediate_root", type=str, default=os.path.join("Embeddings", "Intermediate"),
                        help="Base output root for default encoders (BERT+ViT)")
    parser.add_argument("--alternate_root", type=str, default=os.path.join("Embeddings", "Alternate", "Intermediate"),
                        help="Base output root for non-default encoders (e.g., Longformer/DINOv2)")

    parser.add_argument("--local_only", action="store_true", default=True,
                        help="Load models from local cache only (default: True). Use --no-local_only to allow downloads.")
    parser.add_argument("--no-local_only", dest="local_only", action="store_false",
                        help="Allow downloads from Hugging Face if files are missing.")

    # Optional sharding
    parser.add_argument("--start", type=int, default=0,
                        help="Start line (0-indexed) in manifest (default: 0)")
    parser.add_argument("--end", type=int, default=None,
                        help="End line (exclusive) in manifest (default: process to end)")

    args = parser.parse_args()

    PAIR_NAME = {
    (DEFAULT_TEXT_MODEL, DEFAULT_VISUAL_MODEL): "BERT_ViT",
    ("allenai/longformer-large-4096", "facebook/dinov2-base"): "LONG_DINO",
    (DEFAULT_TEXT_MODEL, "facebook/dinov2-base"): "BERT_DINO",
    ("allenai/longformer-large-4096", DEFAULT_VISUAL_MODEL): "LONG_ViT",
    }

    run_name = PAIR_NAME.get(
    (args.text_model, args.visual_model),
    f"text_{safe_name(args.text_model)}__vision_{safe_name(args.visual_model)}"
    )   
    
    output_root = os.path.join("Embeddings", run_name, "Intermediate")
    os.makedirs(output_root, exist_ok=True)

    print("TEXT_MODEL:", args.text_model)
    print("VISUAL_MODEL:", args.visual_model)
    print("video_root: ", args.video_root)
    print("output_root:", output_root)
    print("local_only:", args.local_only)

    # Load models once
    tokenizer, text_model, visual_processor, visual_model = load_models(
        args.text_model, args.visual_model, CACHE_DIR, local_only=args.local_only
    )

    # Read manifest lines
    with open(args.manifest, "r", encoding="utf-8") as f:
        filenames = [line.strip() for line in f if line.strip() and not line.strip().startswith("#")]

    if not filenames:
        print(f"No CSV entries found in manifest: {args.manifest}")
        return

    start = max(args.start, 0)
    end = len(filenames) if args.end is None else min(args.end, len(filenames))
    if start >= end:
        print(f"Invalid range: start={start}, end={end}, total={len(filenames)}")
        return

    manifest_dir = os.path.dirname(os.path.abspath(args.manifest))

    print(f"Found {len(filenames)} CSVs in manifest.")
    print(f"Processing lines [{start}, {end})")

    for i in range(start, end):
        target_csv = filenames[i]

        if not os.path.isabs(target_csv):
            candidate = os.path.join(manifest_dir, target_csv)
            if os.path.exists(candidate):
                target_csv = candidate

        if not os.path.exists(target_csv):
            print(f"[SKIP] Missing CSV: {target_csv}")
            continue

        video_id = os.path.splitext(os.path.basename(target_csv))[0]
        print(f"[{i}/{end-1}] Processing: {video_id}")

        try:
            process_video_features(
                target_csv,
                args.video_root,
                output_root,
                tokenizer,
                text_model,
                visual_processor,
                visual_model,
            )
        except Exception as e:
            print(f"[ERROR] {video_id} failed: {e}")
            traceback.print_exc()


if __name__ == "__main__":
    main()