import streamlit as st
import pandas as pd
import numpy as np
import io
import os
import replicate
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from openai import OpenAI

st.set_page_config(page_title="Construction Estimator", layout="centered")

client = OpenAI(
    api_key=os.getenv("GROQ_API_KEY"),
    base_url="https://api.groq.com/openai/v1"
)
GROQ_MODEL = "llama3-8b-8192"

@st.cache_data
def load_excel(file):
    return pd.read_excel(file)

@st.cache_resource
def embed_chunks(chunks):
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    embeddings = model.encode(chunks)
    return embeddings, model

def query_embedding(user_query, chunks, embeddings, model):
    query_vec = model.encode([user_query])
    similarities = cosine_similarity(query_vec, embeddings)[0]
    top_idx = np.argmax(similarities)
    return chunks[top_idx]

def generate_estimate(context, user_input):
    prompt = f"""You are a construction estimator in Pakistan. Using the following schedule:

{context}

Generate a BOQ with item number, description, quantity, unit, rate, and total amount for:
{user_input}

Output as a markdown table."""

    response = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

def compute_total_quantities(base_df, covered_area, floors):
    base_df = base_df.copy()
    base_df["Adjusted Qty"] = base_df["Qty per 1000 sft"] * (covered_area / 1000)
    floor_factor = 1 + max(0, floors - 1) * 0.8
    base_df["Total Qty"] = base_df["Adjusted Qty"] * floor_factor
    return base_df

def generate_realistic_plan(rooms, baths, living, car_porch):
    prompt = f"floor plan for {rooms} rooms, {baths} bathrooms, {living} living rooms, and {car_porch} car porch in modern style"
    output = replicate.run(
        "cjwbw/floor-plan-generator",
        input={"prompt": prompt}
    )
    return output

st.title("🏗️ Construction Estimator (Material + Cost + BOQ + Sketch)")

quantity_file = st.file_uploader("Upload Material Quantities Excel (per 1000 sft)", type=["xlsx"])
cost_file = st.file_uploader("Upload Material Costs Excel", type=["xlsx"])

if quantity_file and cost_file:
    base_df = load_excel(quantity_file)
    cost_df = load_excel(cost_file)

    if "Material" in base_df.columns and "Qty per 1000 sft" in base_df.columns:
        st.success("Files loaded successfully.")

        rooms = st.number_input("Number of Rooms", min_value=1, value=3)
        baths = st.number_input("Number of Bathrooms", min_value=1, value=2)
        living = st.number_input("Number of Living Rooms", min_value=0, value=1)
        car_porch = st.number_input("Number of Car Porches", min_value=0, value=1)
        covered_area = st.number_input("Total Covered Area (sft)", min_value=100, value=1200)
        floors = st.number_input("Number of Floors", min_value=1, value=1)

        if st.button("Generate Estimate"):
            computed_df = compute_total_quantities(base_df, covered_area, floors)
            boq = computed_df.merge(cost_df, on="Material", how="left")
            boq["Amount"] = boq["Total Qty"] * boq["Rate_per_Unit"]

            st.subheader("📋 Bill of Quantities (BOQ)")
            st.dataframe(boq[["Material", "Total Qty", "Unit", "Rate_per_Unit", "Amount"]])

            total_cost = boq["Amount"].sum()
            st.metric("Total Estimated Cost (Rs)", f"{total_cost:,.0f}")

            st.subheader("🏠 AI-Generated Floor Plan Sketch")
            try:
                image_url = generate_realistic_plan(rooms, baths, living, car_porch)
                st.image(image_url, caption="Generated by AI (Replicate)", use_column_width=True)
            except Exception as e:
                st.warning("Could not generate sketch. Please check Replicate API setup.")
                st.text(str(e))

    else:
        st.error("Quantity sheet must contain 'Material' and 'Qty per 1000 sft' columns.")
else:
    st.info("Please upload both the quantity and cost Excel sheets.")