import streamlit as st
import os
import re
import base64
import tempfile
import time

# ✅ NEW Gemini SDK
from google import genai

# Local imports
from agents import get_content_enhancer_agent
from tools.image_tool import ImageGenerationTool
from local_knowledge_base import load_knowledge_base

# --- Page Config ---
st.set_page_config(page_title="Video Analyzer", layout="wide")
st.title("🎥 Video Analysis & Content Generation")
st.markdown("Provide a video file or a YouTube link to generate an enriched summary with AI-generated images.")

# --- Gemini Client Initialization ---
try:
    client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
except Exception:
    st.error("GOOGLE_API_KEY not found. Add it to your environment/secrets.")
    st.stop()

# --- Load Knowledge Base ---
@st.cache_resource(show_spinner="Loading Local Knowledge Base...")
def get_local_kb():
    local_kn_b = load_knowledge_base()
    if local_kn_b:
        local_kn_b.load(recreate=False)
    return local_kn_b

local_kb = get_local_kb()

# --- Core Logic ---
def enrich_and_generate_images(video_summary: str) -> str:
    st.write("### Step 2: Enhancing Content with AI Agent...")

    with st.spinner("Enhancing content..."):
        content_agent = get_content_enhancer_agent(knowledge_base=local_kb)
        enriched_text_with_placeholders = "".join(list(content_agent.run(video_summary)))

    st.success("✅ Content enhancement complete.")

    with st.expander("See Enriched Text"):
        st.text(enriched_text_with_placeholders)

    # --- Image Generation ---
    st.write("### Step 3: Generating Images...")
    final_output = enriched_text_with_placeholders

    try:
        image_tool = ImageGenerationTool(api_key=os.getenv("CLIPDROP_API_KEY"))
    except Exception:
        st.error("CLIPDROP_API_KEY missing.")
        st.stop()

    placeholders = re.findall(r'\[IMAGE: caption="(.*?)"\]', final_output)

    if not placeholders:
        st.warning("No image placeholders found.")
        return final_output

    progress_bar = st.progress(0)

    for i, caption in enumerate(placeholders):
        progress_bar.progress(i / len(placeholders))

        image_path = image_tool.generate_image(prompt=caption)

        placeholder = f'[IMAGE: caption="{caption}"]'

        if "Error:" in image_path or not os.path.exists(image_path):
            replacement = f"\n> Image failed: {caption}\n"
        else:
            with open(image_path, "rb") as f:
                encoded = base64.b64encode(f.read()).decode()

            replacement = f"""
<div align="center">
<img src="data:image/png;base64,{encoded}" style="max-width:100%; max-height:500px;">
<br><em>{caption}</em>
</div>
"""

        final_output = final_output.replace(placeholder, replacement, 1)

        time.sleep(3)

    progress_bar.progress(1.0)
    st.success("✅ Images generated.")

    return final_output


# --- YouTube Workflow ---
def run_youtube_analysis_workflow(yt_url: str):
    st.write("### Step 1: Analyzing YouTube Video...")

    try:
        prompt = (
            "Analyze this video. Extract 2-3 core topics and provide a detailed summary in order."
        )

        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=[
                {
                    "role": "user",
                    "parts": [
                        {"text": prompt},
                        {"file_data": {"file_uri": yt_url}}
                    ]
                }
            ]
        )

        video_summary = response.text

        st.success("✅ YouTube analysis complete.")

        with st.expander("Raw Summary"):
            st.markdown(video_summary)

        return enrich_and_generate_images(video_summary)

    except Exception as e:
        st.error(f"YouTube analysis failed: {e}")
        return None


# --- Local File Workflow ---
def run_local_file_analysis_workflow(video_path: str):
    uploaded_file = None

    try:
        st.write("### Step 1: Uploading Video...")

        uploaded_file = client.files.upload(file=video_path)

        st.info("Processing video...")

        prompt = (
            "Analyze this video. Extract key topics and generate a detailed summary."
        )

        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=[
                {
                    "role": "user",
                    "parts": [
                        {"text": prompt},
                        {"file_data": {"file_uri": uploaded_file.uri}}
                    ]
                }
            ]
        )

        video_summary = response.text

        st.success("✅ Local video analysis complete.")

        with st.expander("Raw Summary"):
            st.markdown(video_summary)

        return enrich_and_generate_images(video_summary)

    except Exception as e:
        st.error(f"Local analysis failed: {e}")
        return None

    finally:
        if uploaded_file:
            try:
                client.files.delete(name=uploaded_file.name)
            except Exception:
                pass


# --- UI ---
st.divider()
st.subheader("Select Video Source")

tab1, tab2 = st.tabs(["YouTube Link", "Upload Video"])

video_path = None
final_response = None

with tab1:
    yt_url = st.text_input("YouTube URL")

with tab2:
    video_file = st.file_uploader("Upload video", type=["mp4", "mov", "avi", "webm"])

if st.button("🚀 Analyze", use_container_width=True):

    if yt_url:
        final_response = run_youtube_analysis_workflow(yt_url)

    elif video_file:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
            tmp.write(video_file.getbuffer())
            video_path = tmp.name

        try:
            final_response = run_local_file_analysis_workflow(video_path)
        finally:
            if video_path and os.path.exists(video_path):
                os.remove(video_path)

    else:
        st.warning("Provide a URL or upload a file.")
        st.stop()

    if final_response:
        st.divider()
        st.markdown("## ✨ Final Output")
        st.markdown(final_response, unsafe_allow_html=True)
    else:
        st.error("No output generated.")