"""
Evoneural MVP - Local 3D Mesh + Skybox Generation
Run: streamlit run app.py
Open: http://localhost:8501
"""

import os
import sys
from pathlib import Path

# Ensure project root is on path
ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import streamlit as st

OUTPUTS = ROOT / "outputs"
OUTPUTS.mkdir(exist_ok=True)


def main() -> None:
    st.set_page_config(
        page_title="Evoneural MVP - Mesh & Skybox",
        page_icon="🎮",
        layout="wide",
    )
    st.title("Evoneural MVP – Local Mesh & Skybox")
    st.caption("Text → 3D mesh (TripoSR) and Text → 360° skybox (Stable Diffusion). Runs on localhost.")

    # Sidebar: model setup (token + download)
    with st.sidebar:
        st.subheader("Stable Diffusion model")
        from scripts.skybox_generator import _default_local_weights_dir
        local_model = _default_local_weights_dir()
        if local_model:
            st.success(f"Local model: found")
            st.caption(os.path.basename(local_model))
        else:
            st.warning("No local model. Download below or need internet on first generate.")
        hf_token = st.text_input(
            "Hugging Face token (optional, if behind firewall)",
            type="password",
            key="hf_token",
            placeholder="hf_...",
            help="Get a token at huggingface.co/settings/tokens",
        )
        if hf_token:
            os.environ["HF_TOKEN"] = hf_token
        if st.button("Download model (~4GB to ./weights/sd-v1-5)", key="btn_download"):
            with st.spinner("Downloading model... (may take several minutes)"):
                try:
                    from scripts.download_sd_model import download_sd_model
                    path = download_sd_model(token=hf_token or os.environ.get("HF_TOKEN"))
                    st.success(f"Model saved. Try generating a skybox.")
                    st.rerun()
                except Exception as e:
                    st.error(str(e))
                    st.caption("Set a Hugging Face token above if your network blocks Hugging Face.")

    tab_mesh, tab_skybox = st.tabs(["🟦 Text → 3D Mesh", "🌅 Text → Skybox"])

    with tab_mesh:
        st.subheader("Generate 3D mesh from text")
        st.markdown(
            "Uses **Stable Diffusion** for text→image, then **TripoSR** for image→mesh. "
            "TripoSR repo must be cloned into `./TripoSR` (see README)."
        )
        prompt_mesh = st.text_input(
            "Prompt (e.g. for mesh)",
            value="A highly detailed, sci-fi mechanical drone with glowing blue accents.",
            key="mesh_prompt",
        )
        col1, col2 = st.columns(2)
        with col1:
            mesh_format = st.selectbox("Mesh format", ["glb", "obj"], key="mesh_fmt")
            seed_mesh = st.number_input("Seed (optional)", value=42, min_value=0, key="mesh_seed")
        with col2:
            use_image = st.checkbox("Use uploaded image instead of text", value=False, key="use_img")
            uploaded = st.file_uploader("Upload image for mesh", type=["png", "jpg"], key="mesh_upload") if use_image else None

        if st.button("Generate mesh", key="btn_mesh"):
            if not prompt_mesh.strip() and not use_image:
                st.warning("Enter a prompt or upload an image.")
            else:
                with st.spinner("Running pipeline..."):
                    try:
                        from scripts.mesh_generator import (
                            generate_mesh_from_image,
                            generate_mesh_from_text,
                            find_triposr_root,
                        )
                        triposr_root = find_triposr_root(str(ROOT))
                        if not triposr_root:
                            st.error(
                                "TripoSR not found. Clone it: "
                                "`git clone https://github.com/VAST-AI-Research/TripoSR.git TripoSR`"
                            )
                        elif use_image and uploaded:
                            path = os.path.join(OUTPUTS, "uploaded_mesh_input.png")
                            with open(path, "wb") as f:
                                f.write(uploaded.getvalue())
                            mesh_path, elapsed, msg = generate_mesh_from_image(
                                path,
                                output_dir=str(OUTPUTS / "mesh_run"),
                                mesh_format=mesh_format,
                            )
                            if mesh_path:
                                st.success(f"Done in {elapsed:.1f}s. {msg}")
                                with open(mesh_path, "rb") as f:
                                    st.download_button("Download mesh", f, file_name=os.path.basename(mesh_path), key="dl_mesh_upload")
                            else:
                                st.error(msg)
                        else:
                            mesh_path, elapsed, msg = generate_mesh_from_text(
                                prompt_mesh,
                                output_dir=str(OUTPUTS),
                                mesh_format=mesh_format,
                                seed=seed_mesh,
                            )
                            if mesh_path:
                                st.success(f"Done in {elapsed:.1f}s. {msg}")
                                with open(mesh_path, "rb") as f:
                                    st.download_button("Download mesh", f, file_name=os.path.basename(mesh_path), key="dl_mesh")
                            else:
                                st.error(msg)
                    except Exception as e:
                        st.exception(e)

    with tab_skybox:
        st.subheader("Generate 2:1 equirectangular skybox")
        st.markdown(
            "Uses **Stable Diffusion 2.1** at 2:1 aspect (e.g. 1024×512). "
            "Optional seamless check compares left/right edges."
        )
        prompt_sky = st.text_input(
            "Prompt (e.g. for skybox)",
            value="Cyberpunk city skyline at dusk, neon reflections, cinematic lighting.",
            key="sky_prompt",
        )
        col1, col2 = st.columns(2)
        with col1:
            width = st.selectbox("Width", [1024, 2048], key="sky_w")
            height = width // 2
            seed_sky = st.number_input("Seed (optional)", value=42, min_value=0, key="sky_seed")
        with col2:
            check_seamless = st.checkbox("Run seamless edge check", value=True, key="seamless")

        if st.button("Generate skybox", key="btn_sky"):
            if not prompt_sky.strip():
                st.warning("Enter a prompt.")
            else:
                with st.spinner("Generating skybox..."):
                    try:
                        from scripts.skybox_generator import generate_skybox
                        from scripts.check_seamless import check_seamless as run_seamless

                        out_path, elapsed, vram_mb = generate_skybox(
                            prompt_sky,
                            output_dir=str(OUTPUTS),
                            width=width,
                            height=height,
                            seed=seed_sky,
                        )
                        st.success(f"Done in {elapsed:.1f}s. Peak VRAM: {vram_mb:.0f} MB")
                        st.image(out_path, use_container_width=True)
                        with open(out_path, "rb") as f:
                            st.download_button("Download skybox", f, file_name=os.path.basename(out_path), key="dl_sky")

                        if check_seamless:
                            result = run_seamless(out_path)
                            st.info(result["message"])
                    except Exception as e:
                        st.exception(e)

    st.divider()
    st.caption("Evoneural AI – Local ML Deployment MVP. Models run locally (no API).")


if __name__ == "__main__":
    main()