import gradio as gr
import whisper
import json
import re
import os


def transcribe_video(file_path):
    result = model.transcribe(file_path, verbose=False)
    return result['text'], result['segments']

def extract_highlights(transcript_segments, user_query, max_segments=100):
    transcript_text = "\n".join([
        f"[{int(seg['start'])} --> {int(seg['end'])}] {seg['text']}"
        for seg in transcript_segments[:max_segments]
    ])

    prompt = f"""
You are a JSON-only assistant.

Return only valid JSON in the following format:

{{
  "segments": [
    {{ "start": 52, "end": 58 }},
    {{ "start": 108, "end": 110 }}
  ]
}}

No markdown, no explanation, no wrapping. Just plain JSON.
User query: {user_query}
Transcript:
{transcript_text}
"""

    completion = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[{"role": "user", "content": prompt}],
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=True,
        stop=None,
    )

    result = ""
    for chunk in completion:
        result += chunk.choices[0].delta.content or ""

    try:
        return json.loads(result.strip())
    except json.JSONDecodeError:
        raise ValueError("Failed to parse LLM output:\n" + result)

def process_video(video_file, user_query):
    steps = []
    steps.append("🎬 Uploaded video.")

    video_path = video_file
    _, transcript_segments = transcribe_video(video_path)

    steps.append("🧠 Extracting highlights...")
    results = extract_highlights(transcript_segments, user_query)

    segments = [(max(0, s["start"] - 5), s["end"] + 5) for s in results["segments"]]
    video = mp.VideoFileClip(video_path)
    clips = [video.subclip(start, end) for start, end in segments]
    final = mp.concatenate_videoclips(clips)

    output_path = "highlight_output.mp4"
    steps.append("🎞️ Rendering...")
    final.write_videofile(output_path, codec="libx264", audio_codec="aac", verbose=False, logger=None)

    steps.append("✅ Done!")
    return "\n".join(steps), output_path

with gr.Blocks() as demo:
    gr.Markdown("# ⚽ Momentify – Sports Highlight Generator")
    with gr.Row():
        video_input = gr.Video(label="Upload Match Video (.mp4)")
        query_input = gr.Textbox(label="Highlight Query", placeholder="e.g. Messi goals")
    output_text = gr.Markdown()
    output_video = gr.Video()
    btn = gr.Button("Generate Highlights")
    btn.click(process_video, [video_input, query_input], [output_text, output_video])

demo.launch(share=True)