import streamlit as st
import os
import time
try:
    from dotenv import load_dotenv
except:
    load_dotenv = None
from langchain_community.document_loaders import YoutubeLoader

from langchain_groq import ChatGroq

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableBranch, RunnableLambda
from langchain_core.output_parsers import StrOutputParser

from langchain_text_splitters import RecursiveCharacterTextSplitter

import zipfile

# Load env
if load_dotenv:
    load_dotenv()
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")


img_path = "images.jpg"

col1, col2, col3 = st.columns([1,2,1])

with col2:
    if os.path.exists(img_path):
        st.image(img_path, width=1500)
    else:
        st.warning("Image not found")

st.set_page_config(page_title="YouTube to Website Generator", page_icon="🎥")
st.title("🎥 YouTube to Article Website Generator")

url = st.text_input("Enter YouTube URL")

# Groq Model (Free & Fast)
llm = ChatGroq(model="llama-3.3-70b-versatile")

# -------------------------------
# 🔁 Retry Wrapper
# -------------------------------
def invoke_with_retry(prompt, max_retries=5):
    for attempt in range(max_retries):
        try:
            return llm.invoke(prompt)
        except Exception as e:
            if "429" in str(e) or "rate_limit" in str(e).lower():
                wait = 2 ** attempt * 10
                st.warning(f"Rate limited. Retrying in {wait}s... (Attempt {attempt+1}/{max_retries})")
                time.sleep(wait)
            else:
                raise e
    raise Exception("Max retries exceeded. Please try again later.")

# -------------------------------
# 📥 Extract Transcript
# -------------------------------
def extract_transcript(link: str):
    try:
        loader = YoutubeLoader.from_youtube_url(link, language=["en"])
        docs = loader.load()
        return docs[0].page_content if docs and len(docs) > 0 else "Transcript not available."
    except:
        try:
            loader = YoutubeLoader.from_youtube_url(link, language=["te"])
            docs = loader.load()
            return docs[0].page_content if docs and len(docs) > 0 else "Transcript not available."
        except:
            return "Transcript not available."

# -------------------------------
# ✂️ Split Text
# -------------------------------
def get_text_chunks(text, chunk_size=3000, chunk_overlap=200):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )
    return splitter.split_text(text)

# -------------------------------
# 🧠 Recursive Summarization
# -------------------------------
def recursive_summarize(text):
    chunks = get_text_chunks(text)
    summary = ""

    for chunk in chunks:
        prompt = f"""
        Convert the following into a professional blog article.

        Previous Summary:
        {summary}

        New Content:
        {chunk}

        Use headings, subheadings, and clean formatting.
        """

        response = invoke_with_retry(prompt)
        summary = response.content

    return summary

# -------------------------------
# ⚡ Short Summarizer
# -------------------------------
def base_summarizer(link):
    transcript = extract_transcript(link)

    if transcript == "Transcript not available." or len(transcript.strip()) < 50:
        transcript = "This video explains an important topic. Generate a detailed professional article with headings and structured content."

    prompt = f"""
    Convert this YouTube transcript into a detailed professional article with proper headings, subheadings, and structured format:

    {transcript}
    """

    return invoke_with_retry(prompt).content

# -------------------------------
# 🔍 Decide Long or Short
# -------------------------------
def is_long_video(link):
    transcript = extract_transcript(link)
    return len(transcript) > 1500

# -------------------------------
# 🌐 Website Generator Prompt
# -------------------------------
web_prompt = ChatPromptTemplate.from_template("""
You are a Senior Frontend Developer.

STRICT FORMAT:

--html--
<html>
<head>
<link rel="stylesheet" href="style.css">
</head>
<body>
{article}
<script src="script.js"></script>
</body>
</html>
--html--

--css--
/* CSS HERE */
--css--

--js--
// JS HERE
--js--

Do NOT add explanation.
""")

# -------------------------------
# 🔗 Pipeline
# -------------------------------
pipeline = (
    RunnableBranch(
        (RunnableLambda(is_long_video), RunnableLambda(lambda x: recursive_summarize(extract_transcript(x)))),
        RunnableLambda(base_summarizer)
    )
    | RunnableLambda(lambda x: {"article": x})
    | web_prompt
    | llm
    | StrOutputParser()
)

# -------------------------------
# 🧩 Safe Extract
# -------------------------------
def extract_section(text, tag):
    try:
        return text.split(f"--{tag}--")[1].strip()
    except:
        return f"Error extracting {tag}"

# -------------------------------
# ▶️ Run
# -------------------------------
if st.button("Generate Website"):

    if not url:
        st.error("Please enter a YouTube URL")
    else:
        with st.spinner("Processing... ⏳"):
            article = pipeline.invoke(url)

        html = extract_section(article, "html")
        css = extract_section(article, "css")
        js = extract_section(article, "js")

        with open("index.html", "w", encoding="utf-8") as f:
            f.write(html)

        with open("style.css", "w", encoding="utf-8") as f:
            f.write(css)

        with open("script.js", "w", encoding="utf-8") as f:
            f.write(js)

        with zipfile.ZipFile("website.zip", "w") as zipf:
            zipf.write("index.html")
            zipf.write("style.css")
            zipf.write("script.js")

        st.success("✅ Website Generated!")

        with open("website.zip", "rb") as f:
            st.download_button(
                "📥 Download Website ZIP",
                f,
                "website.zip"
            )

        st.subheader("🔍 HTML Preview")
        st.code(html, language="html")

st.markdown("""
<br>
<div style='text-align:center; padding:12px; background-color:#111111; border-radius:10px;'>
<span style='color:#AAAAAA; font-size:16px;'>
Designed & Developed by <b style='color:#CCCCCC;'>Yedeedya Injeti</b><br>
Under <b style='color:#B8860B;'>Innomatics Research Labs</b>
</span>
</div>
<br>
""", unsafe_allow_html=True)