# === IMPORTS ===

import streamlit as st
from PIL import Image
from openai import OpenAI
from pathlib import Path


# ✅ 1. Page Configuration — MUST be first Streamlit command
st.set_page_config(page_title="AI Architecture Assistant", layout="centered")


# ✅ 2. Inline custom CSS with Poppins font and olive green theme
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;500;700&display=swap');
.stApp {
    background-color: #f4f7f5 !important;
    font-family: 'Poppins', sans-serif;
    color: #1e1e1e;
    padding: 1rem;
}
/* Headings */
h1 {
    font-family: 'Poppins', sans-serif;
    color: #3a5a40 !important;
    font-weight: 700;
    font-size: 2.3rem;
}
h2, h3 {
    font-family: 'Poppins', sans-serif;
    color: #344e41 !important;
    font-weight: 600;
    font-size: 1.4rem;
}
/* Inputs */
.stTextInput input, .stTextArea textarea {
    background-color: #ffffff !important;
    color: #1e1e1e !important;
    border: 1px solid #c8d6c1;
    border-radius: 10px;
    padding: 14px;
    font-size: 1rem;
    font-family: 'Poppins', sans-serif;
}
/* Buttons */
.stButton > button {
    background-color: #588157;
    color: white;
    font-weight: 600;
    font-family: 'Poppins', sans-serif;
    padding: 0.7rem 1.6rem;
    font-size: 1rem;
    border: none;
    border-radius: 10px;
    transition: all 0.3s ease;
    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
}
.stButton > button:hover {
    background-color: #3a5a40;
    box-shadow: 0 4px 10px rgba(0,0,0,0.15);
}
/* Links */
a {
    color: #3a5a40;
    font-family: 'Poppins', sans-serif;
    text-decoration: none;
    font-weight: 500;
}
a:hover {
    text-decoration: underline;
    color: #2c4030;
}
</style>
""", unsafe_allow_html=True)


# ✅ 3. Logo Display (must come after styling)
try:
    logo = Image.open("ThinkTiny.jpg")
    st.image(logo, width=300)
except FileNotFoundError:
    st.warning("Logo not found. Please ensure 'ThinkTiny.jpg' is in the same folder as app.py.")


# ✅ 4. Connect to OpenAI
client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])


# ✅ 5. Session Setup
if "prompt" not in st.session_state:
    st.session_state.prompt = ""


# ✅ 6. Sample Prompt
sample_prompt = (
    "A highly detailed, realistic architectural floor plan for a modern tiny home. "
    "The design features an open-concept layout with a multi-functional living space "
    "that combines a compact living area, dining space, and efficient kitchen with smart storage. "
    "A cozy loft bedroom is accessible via a sleek staircase or ladder. The minimalist bathroom "
    "includes a shower. Emphasize large windows with natural light, clean lines, and neutral tones "
    "with subtle accents. Annotate key dimensions and furniture placements. Professional architectural rendering style."
)


# === UI START ===
st.title("🏠 AI Architecture Assistant")
st.caption("Design smarter. Visualize faster. Build better.")

# === 1. Prompt Input ===
st.markdown("### ✍️ Enter Your Floor Plan Prompt")
st.text_area("Prompt", key="prompt_input", value=st.session_state.prompt, height=200)

col1, col2 = st.columns(2)
with col1:
    if st.button("Insert Sample Prompt"):
        st.session_state.prompt = sample_prompt
        st.rerun()

st.session_state.prompt = st.session_state.prompt_input

st.text_area("Current Prompt Being Used", st.session_state.prompt, height=200, disabled=True)

# === 2. Generate Image ===
st.markdown("### 🖼️ Generate a Floor Plan Image")

if st.button("Generate Image"):
    if st.session_state.prompt.strip():
        with st.spinner("Generating image..."):
            try:
                response = client.images.generate(
                    model="dall-e-3",
                    prompt=st.session_state.prompt,
                    size="1024x1024",
                    quality="standard",
                    n=1
                )
                image_url = response.data[0].url
                st.image(image_url, caption="Generated Floor Plan", use_column_width=True)
                st.markdown(f"[Download Image]({image_url})", unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Error: {e}")
    else:
        st.warning("Please enter a valid prompt.")

# === 3. Recommend Materials ===
st.markdown("### 🧱 Recommended Building Materials")

if st.button("Get Material Recommendations"):
    if st.session_state.prompt.strip():
        with st.spinner("Analyzing design and recommending materials..."):
            def generate_material_recommendations(prompt):
                system_instruction = (
                    "You are a smart building material recommender. Based on the architectural prompt, "
                    "return a list of 5 to 10 building materials implied by the layout or concept. "
                    "Respond in clear, plain English using bullet points. Do not return JSON."
                )

                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": system_instruction},
                        {"role": "user", "content": prompt}
                    ],
                    temperature=0.7
                )
                return response.choices[0].message.content.strip()

            result = generate_material_recommendations(st.session_state.prompt)
            st.markdown(result)
    else:
        st.warning("Please enter a prompt first.")
st.markdown("### 💰 Cost Estimate by Region")

region = st.selectbox("Select a region in Canada:", [
    "Alberta", "British Columbia", "Ontario", "Quebec", "Atlantic Canada", "Prairies"
])

if st.button("Get Cost Estimate"):
    if st.session_state.prompt.strip():
        with st.spinner("Calculating cost estimate..."):
            def generate_cost_estimate(prompt, region):
                system_instruction = (
                    "You are a residential construction cost estimator for Canadian regions. "
                    "Based on the given architectural prompt and selected region, estimate a typical low-to-high cost range "
                    "for building the described design. Mention 4–5 key cost categories. Use Canadian dollars. "
                    "Adjust prices to be realistic for the selected region, and keep your response brief and clear."
                )

                user_input = f"Region: {region}\n\nPrompt:\n{prompt}"

                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": system_instruction},
                        {"role": "user", "content": user_input}
                    ],
                    temperature=0.7
                )
                return response.choices[0].message.content.strip()

            estimate = generate_cost_estimate(st.session_state.prompt, region)
            st.markdown(estimate)
    else:
        st.warning("Please enter a prompt first.")
        st.markdown("### 🏗️ Building Code Comparison")

st.markdown("---")
st.markdown("### 🧾 Building Code Comparison Between Provinces")

province_1 = st.selectbox("Select first province:", [
    "British Columbia", "Alberta", "Ontario", "Quebec", "Atlantic Canada", "Prairies"
], key="prov1")

province_2 = st.selectbox("Select second province:", [
    "British Columbia", "Alberta", "Ontario", "Quebec", "Atlantic Canada", "Prairies"
], key="prov2")

if st.button("Compare Building Codes"):
    if province_1 != province_2:
        with st.spinner("Comparing codes across provinces..."):
            def compare_building_codes(p1, p2, prompt):
                system_instruction = (
                    "You are a Canadian building code expert. Compare the key differences in building code requirements "
                    f"between {p1} and {p2} for residential construction, especially related to small or modern homes like the one described below. "
                    "Focus on differences in energy efficiency, foundations, insulation, accessibility, and design limitations. "
                    "Keep your response clear and under 10 bullet points."
                )

                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": system_instruction},
                        {"role": "user", "content": prompt}
                    ],
                    temperature=0.7
                )
                return response.choices[0].message.content.strip()

            comparison = compare_building_codes(province_1, province_2, st.session_state.prompt)
            st.markdown(f"### 🏛️ {province_1} vs. {province_2}")
            st.markdown(comparison)
    else:
        st.warning("Please select two different provinces.")
        st.markdown("### ➡️ What would you like to do next?")

# === DESIGN IDEAS SECTION ===

st.markdown("### ➡️ What would you like to do next?")

next_step = st.selectbox(
    "Now that you have some information and a layout in mind, would you like help with...",
    [
        "Select an option",
        "🎨 Explore design ideas (color palettes, materials for your climate)",
        "🔍 Continue researching (codes, build types, regulations)"
    ]
)

if next_step == "🎨 Explore design ideas (color palettes, materials for your climate)":
    st.markdown("#### 🎨 Design Recommendations")
    st.markdown("We'll give you smart and stylish ideas based on your region and layout.")

    # Select region
    design_region = st.selectbox("Where will your home be built?", [
        "British Columbia", "Alberta", "Ontario", "Quebec", "Atlantic Canada", "Prairies"
    ], key="design_region")

    # Select style
    design_style = st.selectbox("What style do you like?", [
        "Scandinavian", "Modern Minimalist", "Rustic Cabin", "Japandi", "West Coast Contemporary"
    ], key="design_style")

    # Select design advice categories
    selected_categories = st.multiselect(
        "What kind of design advice would you like?",
        [
            "Climate-Responsive Material Suggestions",
            "Interior Color Palette",
            "Lighting & Window Placement Tips",
            "Space-Saving Built-ins"
        ],
        default=["Climate-Responsive Material Suggestions", "Interior Color Palette"]
    )

# Generate design ideas with GPT
if st.button("Generate Design Ideas"):
        with st.spinner("Designing your dream space..."):

            def generate_design_ideas(prompt, region, style, categories):
                system_instruction = (
                    "You are a residential design assistant. Based on the home description, region, and preferred design style, "
                    "give tailored recommendations for the selected categories. Keep it concise, stylish, and easy to implement for a small-scale home."
                )

                user_input = (
                    f"Prompt:\n{prompt}\n\n"
                    f"Region: {region}\n"
                    f"Style: {style}\n"
                    f"Categories: {', '.join(categories)}"
                )

                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": system_instruction},
                        {"role": "user", "content": user_input}
                    ],
                    temperature=0.7
                )
                return response.choices[0].message.content.strip()

            design_results = generate_design_ideas(
                st.session_state.prompt,
                design_region,
                design_style,
                selected_categories
            )

            st.markdown("### 💡 Design Suggestions")
            st.markdown(design_results)

            # Store summary in session state
            st.session_state.design_summary = {
                "Prompt": st.session_state.prompt,
                "Region": design_region,
                "Style": design_style,
                "Categories": selected_categories,
                "Design Suggestions": design_results,
                "Materials": st.session_state.get("material_recommendations", "Not generated"),
                "Image URL": st.session_state.get("generated_image_url", "Not generated")
            }

# === Download Summary ===
import io

st.markdown("---")
st.markdown("### 📥 Download Your Concept Summary")

if st.button("Download Summary as Text"):
    summary = st.session_state.design_summary
    content = f"""🏠 AI Architecture Design Summary
Prompt:
{summary['Prompt']}
Region: {summary['Region']}
Style: {summary['Style']}
Selected Categories: {', '.join(summary['Categories'])}
🎨 Design Suggestions:
{summary['Design Suggestions']}
🧱 Recommended Materials:
{summary['Materials']}
🖼️ Floor Plan Image:
{summary['Image URL']}
"""

    st.download_button(
        label="Download Summary",
        data=io.StringIO(content),
        file_name="design_summary.txt",
        mime="text/plain"
    )

# === Concept Image Generation ===
st.markdown("### 🖼️ Visualizing Your Design Concept")

if st.button("Generate Concept Image from Design Summary"):
    with st.spinner("Creating a personalized visual..."):
        refined_prompt = (
    f"Based on the following tiny home project details:\n\n"
    f"Region: {design_region}\n"
    f"Style: {design_style}\n"
    f"Categories: {', '.join(selected_categories)}\n"
    f"Design Suggestions: {st.session_state.design_summary['Design Suggestions']}\n\n"
    "Create a professional architectural rendering that visually reflects this concept.\n"
    "Focus on realistic materials, lighting, and color palette. Use a 3D rendering style with natural surroundings."
)


        try:
            response = client.images.generate(
                model="dall-e-3",
                prompt=refined_prompt,
                size="1024x1024",
                quality="standard",
                n=1
            )
            concept_image_url = response.data[0].url
            st.image(concept_image_url, caption="Generated Concept Image", use_column_width=True)
            st.markdown(f"[Download Image]({concept_image_url})", unsafe_allow_html=True)

            st.session_state.generated_concept_image_url = concept_image_url

        except Exception as e:
            st.error(f"Image generation failed: {e}")
            
if next_step == "🔍 Continue researching (codes, build types, regulations)":
    st.markdown("#### 🔍 Research Hub")
    st.info("Here are some useful topics and tools to help you dig deeper before you build.")

    st.markdown("##### 🧠 Suggested Topics to Explore")
    research_topics = {
        "Prefab vs. Traditional Construction": "Compare pros and cons of prefabricated (modular) vs. traditional stick-built homes.",
        "Land Use & Zoning": "Understand how land use bylaws and zoning affect what and where you can build.",
        "Permit & Code Requirements": "Find out what permits are typically needed for a small home project in Canada.",
        "Sustainable Building Materials": "Explore environmentally-friendly materials that are affordable and effective.",
        "Off-Grid Living Regulations": "Learn the legal considerations if you want to live off-grid.",
        "Energy Efficiency Standards": "Understand how to meet or exceed Canadian energy code requirements."
    }

    selected_topic = st.selectbox("Choose a topic to learn more:", ["Select a topic"] + list(research_topics.keys()))
    if selected_topic != "Select a topic":
        with st.spinner("Researching..."):
            response = client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a Canadian construction research assistant."},
                    {"role": "user", "content": research_topics[selected_topic]}
                ],
                temperature=0.7
            )
            st.markdown(f"### 📖 {selected_topic}")
            st.markdown(response.choices[0].message.content.strip())

    st.markdown("##### 💬 Ask Your Own Research Question")
user_question = st.text_input("What else would you like to learn about?")

if st.button("Ask"):
    if user_question:
        with st.spinner("Getting the answer..."):
            response = client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You're a helpful Canadian homebuilding advisor. Keep answers clear and based on current standards."},
                    {"role": "user", "content": user_question}
                ],
                temperature=0.7
            )
            st.markdown("### 💡 Answer")
            st.markdown(response.choices[0].message.content.strip())
    else:
        st.warning("Please type a question first.")

# === Helpful Resource Links (shown no matter what) ===
st.markdown("##### 🔗 Helpful Resources")
st.markdown("""
- [National Building Code of Canada (2020)](https://nrc.canada.ca/en/certifications-evaluations-standards/codes-canada/codes-canada-publications/national-building-code-canada-2020)
- [BC Building Code](https://www.bccodes.ca/)
- [Alberta Building Codes & Standards](https://www.alberta.ca/building-codes-and-standards.aspx)
- [Ontario Building Code](https://www.ontario.ca/page/building-code)
- [Quebec Construction Code](https://www.rbq.gouv.qc.ca/en/technical-buildings/buildings/construction-code.html)
- [Tiny Home Legal Info (Canada)](https://tinyhousecanada.ca/)
""")
# === 9. Eco-Friendly Upgrades Agent ===

st.markdown("---")
st.markdown("### 🌱 Suggest Eco-Friendly Upgrades")

if st.button("Get Eco-Friendly Suggestions"):
    if st.session_state.prompt.strip():
        with st.spinner("Analyzing your layout for green upgrades..."):

            def suggest_eco_upgrades(prompt):
                system_instruction = (
                    "You are a sustainability expert in residential architecture. "
                    "Based on the following home design prompt, suggest 5 to 7 eco-friendly upgrades. "
                    "Be specific and include options like energy efficiency, water conservation, material choices, and renewable energy use. "
                    "Avoid general advice—tailor suggestions to the design and layout described."
                )

                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": system_instruction},
                        {"role": "user", "content": prompt}
                    ],
                    temperature=0.7
                )

                return response.choices[0].message.content.strip()

            eco_suggestions = suggest_eco_upgrades(st.session_state.prompt)

            st.markdown("### ✅ Recommended Eco-Friendly Features")
            st.markdown(eco_suggestions)

            # Save to session for summary
            st.session_state.eco_suggestions = eco_suggestions

    else:
        st.warning("Please enter a prompt first.")