from image_processor import process_zip_and_generate_images
from load_file import create_streamlit_app 
import streamlit as st
import os
import zipfile
import tempfile
import requests
import shutil
import logging
from PIL import Image

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

examples = [
    ["examples/input/3.jpeg", "This ad creative should be like a facebook post opened in someone's phone.", ["examples/output/image (5).png", "examples/output/image (5).png"]],
    ["examples/input/3.jpeg", "Make this ad creative like a news information on TV", ["examples/output/image (8).png", "examples/output/image (7).png"]],
    ["examples/input/3.jpeg", "The ad creative should comprises of 3-4 images. The images should be like trust building images with no CTA or offer. Moreover, the images should not be much HD it should be of low quality grained image.", ["examples/output/image (3).png", "examples/output/image (4).png"]]
]

custom_css = """
<style>
.stApp {
    max-width: 1200px;
    margin: auto;
    background-color: #fafbfc !important;
}

.prompt-section {
    margin-top: 10px;
    background-color: #f6f6f7;
    padding: 5px;
    border-radius: 8px;
}
.stImage > img, .image-container img {
    height: auto !important;
}
h1, h2, h3, h4, h5 {
    color: #1a1a1a !important;
    margin: 0 !important;
    padding: 5px 0 !important;
    font-weight: 600;
}
body, .stApp {
    background-color: #fafbfc !important;
    color: #24292f !important;
}

[data-testid="stSidebar"] {
    background-color: #fff !important;
    color: #222 !important;
    border-right: 1px solid #e5e7eb;
}
.stTabs [data-baseweb="tab-list"] {
    border-bottom: 2px solid #e5e7eb !important;
}
.stTabs [data-baseweb="tab"].css-1n76uvr {
    color: #24292f !important;
    background: #fff !important;
    border-bottom: 2px solid #ff8900 !important;
    font-weight: bold;
}

.stButton>button {
    background: #fff;
    color: #1a1a1a;
    border: 1px solid #e5e7eb;
    border-radius: 8px;
    font-weight: 500;
    transition: all 0.15s;
}

/* Match card style to form fields */
.example-card {
    border: 1px solid #e5e7eb;
    border-radius: 8px;
    background-color: #f6f6f7;
    color: #24292f;
    margin-bottom: 20px;
    padding: 18px 14px 14px 14px;
    
}

.stSelectbox > div[data-baseweb="select"]:hover {
    border-color: #ff8900 !important; /* orange accent on hover */
    box-shadow: 0 4px 14px rgba(255,137,0,0.08);
}
.stSelectbox label {
    font-weight: 600 !important;
    color: #444 !important;
    margin-bottom: 6px;
    font-size: 1.03rem;
}
.stSelectbox [data-testid="stMarkdownContainer"] {
    color: #888 !important;
}

</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

def zip_gallery_images(gallery_items):
    if not gallery_items:
        logger.warning("No gallery items received for zipping.")
        return None

    image_urls = [item[0] if isinstance(item, tuple) else item for item in gallery_items]

    try:
        temp_dir = tempfile.mkdtemp()
        zip_path = tempfile.NamedTemporaryFile(delete=False, suffix=".zip").name

        for i, url in enumerate(image_urls):
            try:
                ext = url.split('.')[-1]
                file_path = os.path.join(temp_dir, f"image_{i}.{ext}")
                if url.startswith("http://") or url.startswith("https://"):
                    response = requests.get(url, timeout=5)
                    response.raise_for_status()
                    with open(file_path, 'wb') as f:
                        f.write(response.content)
                    logger.info(f"Downloaded and saved remote image: {url}")
                elif os.path.exists(url):
                    shutil.copy(url, file_path)
                    logger.info(f"Copied local image: {url}")
                else:
                    logger.warning(f"Skipping invalid path or URL: {url}")
            except Exception as e:
                logger.error(f"Error processing image {url}: {e}")

        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for file_name in os.listdir(temp_dir):
                zipf.write(os.path.join(temp_dir, file_name), arcname=file_name)
        logger.info(f"Zip archive created at: {zip_path}")
        return zip_path

    except Exception as e:
        logger.critical(f"Failed to create zip: {e}")
        return None

    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)
        logger.debug("Cleaned up temporary directory.")

def check_token(user_token):
    ACCESS_TOKEN = os.getenv("ACCESS_TOKEN")
    if not ACCESS_TOKEN:
        logger.critical("ACCESS_TOKEN not set in environment.")
        return False, "Server error: Access token not configured."
    if user_token == ACCESS_TOKEN:
        logger.info("Access token validated successfully.")
        return True, ""
    logger.warning("Invalid access token attempt.")
    return False, "Invalid token."


def create_main_app():
    st.markdown("# Bulk Creative Generation")
    tabs = st.tabs(["Generate Images", "Creative AI Library"])
    # --- Generate Images Tab ---
    with tabs[0]:
        st.markdown("Bulk Creative Generation is a tool where you can upload a single image or zip of upto 10 images which will create different images for every uploaded image. Make sure the user prompt which you will write will be applicable to each image. Make sure you use this tool wisely. Also refer to the examples below for better understanding.")

        st.markdown("## System Prompt:")
        st.markdown("""You are a top-tier performance digital marketer and creative strategist with 15+ years of expertise in affiliate marketing.
Your objective is to analyze the provided winning ad image, deconstruct its concept, visual composition, and color scheme, and generate a fresh, conversion-focused ad visual tailored for the {category} niche.
The new design should convey a {sentiment} sentiment and incorporate the user instruction: {user_prompt}.
Create a visually compelling ad optimized for {platform} Ads that is scroll-stopping, pattern-interrupting, and designed to drive high CTR and Conversion Rate. Utilize striking color combinations, dynamic contrast levels, and strategic layout compositions to command attention while aligning with the target audience avatar.
Make sure the images should be realistic, not be stocky at all and raw which should look like they are shot from an iPhone with {background}.""")

        # File Uploader and fields
        zip_file = st.file_uploader("Upload Zip or Single File", type=["zip", "png", "jpg", "jpeg"])
        category = st.text_input("Category")

        st.markdown("### Settings")
        col1, col2, col3, col4, col5, col6 = st.columns([2, 2.5, 2.5, 2.5, 2.5, 2])
        with col1:
            blur = st.checkbox("Blur Image")
        with col2:
            size = st.selectbox("Image Size", options=["auto", "1024x1024", "1536x1024", "1024x1536"])
        with col3:
            quality = st.selectbox("Quality", options=["auto", "low", "medium"])
        with col4:
            sentiment = st.selectbox("Sentiment", options=["as original image", "positive", "negative"])
        with col5:
            platform = st.selectbox("Platform", options=["Facebook", "Native", "Newsbreak", "GDN"])
        with col6:
            num_images = st.slider("No. of Images to be generated:", 1, 10, value=2, step=1)


        user_prompt = st.text_area("User Prompt", height=100)

        # Buttons
        colA, colB, colC = st.columns([1,1,1])
        with colA:
            demo_btn = st.button("Generate Demo Image", key="demo_img")
        with colB:
            generate_all_btn = st.button("Generate All Images", key="gen_all_img")
        with colC:
            download_btn = st.button("Download All", key="download_all")

        if "gallery" not in st.session_state:
            st.session_state["gallery"] = []
        gallery = st.session_state["gallery"]

        # Image Generation
        if demo_btn or generate_all_btn:
            if zip_file and category and user_prompt:
                temp_dir = tempfile.mkdtemp()
                if zip_file.name.endswith(".zip"):
                    temp_path = os.path.join(temp_dir, zip_file.name)
                else:
                    temp_path = os.path.join(temp_dir, "input." + zip_file.name.split(".")[-1])
                with open(temp_path, "wb") as f:
                    f.write(zip_file.read())
                try:
                    images = process_zip_and_generate_images(
                        temp_path, category, size, quality, user_prompt, sentiment, platform, num_images, demo_btn, gallery, blur
                    )
                    st.session_state["gallery"] = images
                    cols = st.columns(4)
                    for idx, img_path in enumerate(images):
                        with cols[idx % 4]:
                            st.image(img_path)
                except Exception as e:
                    st.error(f"Error: {e}")
                shutil.rmtree(temp_dir)
            else:
                st.warning("Please upload a file and fill all required fields.")

        # Download All ZIP Event
        if download_btn:
            zip_path = zip_gallery_images(st.session_state["gallery"])
            if zip_path:
                with open(zip_path, "rb") as f:
                    st.download_button(
                        label="Download ZIP",
                        data=f,
                        file_name="generated_images.zip",
                        mime="application/zip"
                    )
                os.remove(zip_path)
            else:
                st.warning("No images to zip and download.")

        # --- Creative AI Library Tab ---
    with tabs[1]:
        MONGO_URI = os.getenv("MONGO_URI")
        if not MONGO_URI:
            logger.warning("MONGO_URI not set; Creative AI Library may fail to load.")
        DB_NAME = "adgenesis"
        COLLECTION_NAME = "records"
        create_streamlit_app(MONGO_URI, DB_NAME, COLLECTION_NAME)  

      

def main():
    st.set_page_config(page_title="Bulk Creative Generation", layout="wide")
    if "authenticated" not in st.session_state:
        st.session_state["authenticated"] = False

    if not st.session_state["authenticated"]:
        st.markdown("## Access Required")
        token_input = st.text_input("Enter Access Token", type="password")
        if st.button("Unlock App"):
            ok, error_msg = check_token(token_input)
            if ok:
                st.session_state["authenticated"] = True
                st.rerun()
            else:
                st.error(error_msg)
    else:
        create_main_app()
        # creative_ai_library_tab()

if __name__ == "__main__":
    try:
        logger.info("Launching Streamlit app...")
        main()
    except Exception as e:
        logger.exception("Unhandled error during app launch.")