# enrich_subtitles_v02.py
import pandas as pd
import pysrt
from datetime import datetime, timedelta
import tkinter as tk
from tkinter import filedialog, messagebox
import os
import json

CONFIG_FILE = "enrich_subtitles_config.json"

def load_last_dir():
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE, 'r') as f:
            return json.load(f).get("last_dir", ".")
    return "."

def save_last_dir(path):
    with open(CONFIG_FILE, 'w') as f:
        json.dump({"last_dir": path}, f)

def get_video_start_time(srt_path):
    with open(srt_path, 'r', encoding='utf-8') as file:
        for line in file:
            line = line.strip()
            if line and not line.startswith(('0', '1', '2')) and ':' in line:
                try:
                    # Attempt to parse the line as a time
                    return datetime.strptime(line, "%H:%M:%S")
                except ValueError:
                    try:
                        return datetime.strptime(line, "%H:%M:%S.%f")
                    except ValueError:
                        continue
    raise ValueError("No valid time found in subtitle file.")


def enrich_subtitles(original_srt, params_file, telemetry_csv):
    print (original_srt)
    #video_start_time = datetime.strptime("13:50:48.000", "%H:%M:%S.%f")
    #video_start_time = datetime.strptime("13:59:44.000", "%H:%M:%S.%f")
     #srt_path = "original_srt.srt"  # Or whatever path you're using
    video_start_time = get_video_start_time(original_srt)
    print("Video starts at:", video_start_time)
    telemetry = pd.read_csv(telemetry_csv)
    telemetry['Time'] = pd.to_datetime(telemetry['Time'], format="%H:%M:%S.%f")
    telemetry = telemetry.sort_values('Time')

    with open(params_file) as f:
        params = [line.strip() for line in f if line.strip()]

    subs = pysrt.open(original_srt)
    new_subs = pysrt.SubRipFile()

    for sub in subs:
        rel_time = timedelta(
            hours=sub.start.hours,
            minutes=sub.start.minutes,
            seconds=sub.start.seconds,
            milliseconds=sub.start.milliseconds
        )
        sub_abs_time = video_start_time + rel_time
        idx = telemetry['Time'].sub(sub_abs_time).abs().idxmin()
        row = telemetry.loc[idx]

        enriched_lines = [sub_abs_time.strftime('%H:%M:%S')]
        for param in params:
            if param in telemetry.columns:
                value = row[param]
                enriched_lines.append(f"{param}: {value}")

        new_sub = pysrt.SubRipItem(
            index=sub.index,
            start=sub.start,
            end=sub.end,
            text="\n".join(enriched_lines)
        )
        new_subs.append(new_sub)

    base, _ = os.path.splitext(original_srt)
    out_file = f"{base}_enriched.srt"
    new_subs.save(out_file, encoding='utf-8')
    return out_file

def build_gui():
    def browse_file(entry, filetypes):
        path = filedialog.askopenfilename(initialdir=load_last_dir(), filetypes=filetypes)
        if path:
            entry.delete(0, tk.END)
            entry.insert(0, path)
            save_last_dir(os.path.dirname(path))

    def run_process():
        srt = srt_entry.get()
        params = params_entry.get()
        telemetry = telemetry_entry.get()

        if not (srt and params and telemetry):
            messagebox.showerror("Error", "Please select all files.")
            return

        try:
            output = enrich_subtitles(srt, params, telemetry)
            messagebox.showinfo("Success", f"Enriched subtitles saved:\n{output}")
        except Exception as e:
            messagebox.showerror("Processing Error", str(e))

    root = tk.Tk()
    root.title("Enrich Subtitles")

    tk.Label(root, text="timeofday (.srt):").grid(row=0, column=0, sticky="w", padx=5, pady=5)
    srt_entry = tk.Entry(root, width=60)
    srt_entry.grid(row=0, column=1, padx=5)
    tk.Button(root, text="Browse", command=lambda: browse_file(srt_entry, [("SRT files", "*.srt")])).grid(row=0, column=2)

    tk.Label(root, text="Parameters File (.txt):").grid(row=1, column=0, sticky="w", padx=5, pady=5)
    params_entry = tk.Entry(root, width=60)
    params_entry.grid(row=1, column=1, padx=5)
    tk.Button(root, text="Browse", command=lambda: browse_file(params_entry, [("Text files", "*.txt")])).grid(row=1, column=2)

    tk.Label(root, text="Telemetry File (.csv):").grid(row=2, column=0, sticky="w", padx=5, pady=5)
    telemetry_entry = tk.Entry(root, width=60)
    telemetry_entry.grid(row=2, column=1, padx=5)
    tk.Button(root, text="Browse", command=lambda: browse_file(telemetry_entry, [("CSV files", "*.csv")])).grid(row=2, column=2)

    tk.Button(root, text="Process", command=run_process, width=20, bg="#4CAF50", fg="white").grid(row=3, column=1, pady=15)

    root.mainloop()

if __name__ == "__main__":
    build_gui()
