# =========================
# GeoMate V2 - Full App
# =========================
import os
import re
import json
import datetime
from math import floor
from typing import Dict, Any, Tuple, List

import streamlit as st
from streamlit_option_menu import option_menu
from fpdf import FPDF
# ===============================
# Earth Engine Initialization
# ===============================
import ee, os, json
from google.oauth2 import service_account

def init_earth_engine():
    try:
        key_json = os.getenv("EARTH_ENGINE_KEY")
        if not key_json:
            st.error("❌ EARTH_ENGINE_KEY secret is missing. Please configure it under Settings → Secrets.")
            return False
        key_dict = json.loads(key_json)
        creds = service_account.Credentials.from_service_account_info(key_dict)
        ee.Initialize(creds, project=key_dict.get("project_id"))
        st.success("✅ Earth Engine initialized successfully!")
        return True
    except Exception as e:
        st.error(f"Earth Engine init failed: {e}")
        return False

EE_READY = init_earth_engine()
# Optional libs (RAG, embeddings, FAISS, Groq, EE)
# They may not be available on first load; we guard usage at runtime.
try:
    from sentence_transformers import SentenceTransformer
    import faiss  # noqa: F401
    HAVE_EMBED = True
except Exception:
    HAVE_EMBED = False

try:
    from langchain.memory import ConversationBufferMemory
    from langchain.chains import ConversationChain
    from langchain_community.chat_models import ChatGroq
    HAVE_LANGCHAIN = True
except Exception:
    HAVE_LANGCHAIN = False

try:
    import ee
    import geemap
    HAVE_EE = True
except Exception:
    HAVE_EE = False


# ==============================================================
#  GLOBALS & SESSION BOOT
# ==============================================================
st.set_page_config(page_title="GeoMate V2", page_icon="🌍", layout="wide")

# Global/session stores
ss = st.session_state
if "soil_description_site" not in ss:
    # multi-site dict
    ss.soil_description_site: Dict[str, Dict[str, Any]] = {}
if "sites" not in ss:
    ss.sites: List[str] = ["site1"]
if "current_site" not in ss:
    ss.current_site = "site1"
if "MODEL_NAME" not in ss:
    ss.MODEL_NAME = "llama-3.1-70b-versatile"   # <- change freely
if "secrets_status" not in ss:
    ss.secrets_status = {"groq_ok": False, "ee_ok": False}
if "rag_ready" not in ss:
    ss.rag_ready = False
if "rag_memory" not in ss:
    ss.rag_memory = None
if "rag_chain" not in ss:
    ss.rag_chain = None
if "emb_model" not in ss:
    ss.emb_model = None

# Step trackers for chatbot-style flows (per page, per site)
if "steps" not in ss:
    ss.steps = {}
# ensure per-page steppers exist
for key in ["classifier", "reports", "locator"]:
    if key not in ss.steps:
        ss.steps[key] = {}
    if ss.current_site not in ss.steps[key]:
        ss.steps[key][ss.current_site] = {"step": 0}

# Store classifier working inputs per site
if "cls_inputs" not in ss:
    ss.cls_inputs = {}
if ss.current_site not in ss.cls_inputs:
    ss.cls_inputs[ss.current_site] = {}

# Store reports Q&A per site
if "reports_inputs" not in ss:
    ss.reports_inputs = {}
if ss.current_site not in ss.reports_inputs:
    ss.reports_inputs[ss.current_site] = {}


# ==============================================================
#  STARTUP SECRET CHECKS (top-level popups, no crash)
# ==============================================================
def check_secrets_banner():
    groq_key = os.getenv("GROQ_API_KEY", "")
    ee_key = os.getenv("EARTH_ENGINE_KEY", "")

    groq_ok = bool(groq_key)
    ee_ok = bool(ee_key)

    ss.secrets_status["groq_ok"] = groq_ok
    ss.secrets_status["ee_ok"] = ee_ok

    cols = st.columns(2)
    with cols[0]:
        if groq_ok:
            st.success("✅ Groq API key detected.")
        else:
            st.error("❌ Groq API key missing (set `GROQ_API_KEY`). LLM chat will be disabled.")
    with cols[1]:
        if ee_ok:
            st.success("✅ Earth Engine key detected.")
        else:
            st.error("❌ Earth Engine key missing (set `EARTH_ENGINE_KEY`). Locator map will be limited.")


# ==============================================================
#  RAG MEMORY (session-scoped, “like ChatGPT” during this run)
# ==============================================================
def init_rag():
    if not ss.secrets_status["groq_ok"] or not HAVE_LANGCHAIN:
        ss.rag_ready = False
        return

    if ss.rag_memory is None:
        ss.rag_memory = ConversationBufferMemory()
    if ss.rag_chain is None:
        try:
            llm = ChatGroq(model_name=ss.MODEL_NAME, temperature=0.2)
            ss.rag_chain = ConversationChain(llm=llm, memory=ss.rag_memory, verbose=False)
            ss.rag_ready = True
        except Exception:
            ss.rag_ready = False

    # Embeddings for future RAG vectorization if needed
    if HAVE_EMBED and ss.emb_model is None:
        try:
            ss.emb_model = SentenceTransformer("all-MiniLM-L6-v2")
        except Exception:
            ss.emb_model = None


def rag_ask(query: str) -> str:
    """
    Converse naturally; memory retained within this session.
    """
    if not ss.rag_ready or ss.rag_chain is None:
        return "LLM is unavailable (Groq key missing or initialization failed)."
    try:
        return ss.rag_chain.predict(input=query)
    except Exception as e:
        return f"LLM error: {e}"


# ==============================================================
#  SITE STORE
# ==============================================================
def save_site_info(site: str, key: str, value: Any):
    if site not in ss.soil_description_site:
        ss.soil_description_site[site] = {}
    ss.soil_description_site[site][key] = value


# ==============================================================
#  USCS & AASHTO VERBATIM LOGIC
# ==============================================================
ENGINEERING_CHARACTERISTICS = {
    "Gravel": {
        "Settlement": "None",
        "Quicksand": "Impossible",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Possible",
        "Cement_grouting": "Possible",
        "Silicate_bitumen_injections": "Unsuitable",
        "Compressed_air": "Possible (see notes)"
    },
    "Coarse sand": {
        "Settlement": "None",
        "Quicksand": "Impossible",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Possible",
        "Cement_grouting": "Possible only if very coarse",
        "Silicate_bitumen_injections": "Suitable",
        "Compressed_air": "Suitable"
    },
    "Medium sand": {
        "Settlement": "None",
        "Quicksand": "Unlikely",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Suitable",
        "Cement_grouting": "Impossible",
        "Silicate_bitumen_injections": "Suitable",
        "Compressed_air": "Suitable"
    },
    "Fine sand": {
        "Settlement": "None",
        "Quicksand": "Liable",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Suitable",
        "Cement_grouting": "Impossible",
        "Silicate_bitumen_injections": "Not possible in very fine sands",
        "Compressed_air": "Suitable"
    },
    "Silt": {
        "Settlement": "Occurs",
        "Quicksand": "Liable (very coarse silts may behave differently)",
        "Frost-heaving": "Occurs",
        "Groundwater_lowering": "Generally not suitable (electro-osmosis possible)",
        "Cement_grouting": "Impossible",
        "Silicate_bitumen_injections": "Impossible",
        "Compressed_air": "Suitable"
    },
    "Clay": {
        "Settlement": "Occurs",
        "Quicksand": "Impossible",
        "Frost-heaving": "None",
        "Groundwater_lowering": "Impossible (generally)",
        "Cement_grouting": "Only in stiff fissured clay",
        "Silicate_bitumen_injections": "Impossible",
        "Compressed_air": "Used for support only in special cases"
    }
}

def uscs_aashto_verbatim(inputs: Dict[str, Any]) -> Tuple[str, str, str, int, Dict[str, str]]:
    """
    Verbatim USCS & AASHTO classifier.
    Returns: (result_text, uscs, aashto, GI, char_summary)
    """
    opt = str(inputs.get("opt","n")).lower()
    if opt == 'y':
        uscs = "Pt"
        uscs_expl = "Peat / organic soil — compressible, high organic content; poor engineering properties."
        aashto = "Organic (special handling)"
        characteristics = {"summary":"Highly organic peat — large settlement, low strength, not suitable for foundations."}
        return f"USCS: **{uscs}** — {uscs_expl}\n\nAASHTO: **{aashto}**", uscs, aashto, 0, characteristics

    P2 = float(inputs.get("P2", 0.0))
    P4 = float(inputs.get("P4", 0.0))
    D60 = float(inputs.get("D60", 0.0))
    D30 = float(inputs.get("D30", 0.0))
    D10 = float(inputs.get("D10", 0.0))
    LL = float(inputs.get("LL", 0.0))
    PL = float(inputs.get("PL", 0.0))
    PI = LL - PL

    Cu = (D60 / D10) if (D10 > 0 and D60 > 0) else 0
    Cc = ((D30 ** 2) / (D10 * D60)) if (D10 > 0 and D30 > 0 and D60 > 0) else 0

    uscs, uscs_expl = "Unknown", ""

    if P2 <= 50:
        # Coarse soils
        if P4 <= 50:
            # Gravels
            if Cu and Cc:
                if Cu >= 4 and 1 <= Cc <= 3:
                    uscs, uscs_expl = "GW", "Well-graded gravel (good engineering properties)."
                else:
                    uscs, uscs_expl = "GP", "Poorly-graded gravel."
            else:
                if PI < 4 or PI < 0.73 * (LL - 20):
                    uscs, uscs_expl = "GM", "Silty gravel."
                elif PI > 7 and PI > 0.73 * (LL - 20):
                    uscs, uscs_expl = "GC", "Clayey gravel."
                else:
                    uscs, uscs_expl = "GM-GC", "Gravel with mixed fines."
        else:
            # Sands
            if Cu and Cc:
                if Cu >= 6 and 1 <= Cc <= 3:
                    uscs, uscs_expl = "SW", "Well-graded sand."
                else:
                    uscs, uscs_expl = "SP", "Poorly-graded sand."
            else:
                if PI < 4 or PI <= 0.73 * (LL - 20):
                    uscs, uscs_expl = "SM", "Silty sand."
                elif PI > 7 and PI > 0.73 * (LL - 20):
                    uscs, uscs_expl = "SC", "Clayey sand."
                else:
                    uscs, uscs_expl = "SM-SC", "Transition silty/clayey sand."
    else:
        # Fine soils (P2 > 50)
        nDS = int(inputs.get("nDS", 5))
        nDIL = int(inputs.get("nDIL", 6))
        nTG = int(inputs.get("nTG", 6))
        if LL < 50:
            if 20 <= LL < 50 and PI <= 0.73 * (LL - 20):
                if nDS == 1 or nDIL == 3 or nTG == 3:
                    uscs, uscs_expl = "ML", "Silt (low plasticity)."
                elif nDS == 3 or nDIL == 3 or nTG == 3:
                    uscs, uscs_expl = "OL", "Organic silt."
                else:
                    uscs, uscs_expl = "ML-OL", "Mixed silt/organic."
            elif 10 <= LL <= 30 and 4 <= PI <= 7 and PI > 0.72 * (LL - 20):
                if nDS == 1 or nDIL == 1 or nTG == 1:
                    uscs, uscs_expl = "ML", "Silt."
                elif nDS == 2 or nDIL == 2 or nTG == 2:
                    uscs, uscs_expl = "CL", "Clay (low plasticity)."
                else:
                    uscs, uscs_expl = "ML-CL", "Mixed silt/clay."
            else:
                uscs, uscs_expl = "CL", "Clay (low plasticity)."
        else:
            if PI < 0.73 * (LL - 20):
                if nDS == 3 or nDIL == 4 or nTG == 4:
                    uscs, uscs_expl = "MH", "Silt (high plasticity)."
                elif nDS == 2 or nDIL == 2 or nTG == 4:
                    uscs, uscs_expl = "OH", "Organic clay/silt (high plasticity)."
                else:
                    uscs, uscs_expl = "MH-OH", "Mixed high-plasticity."
            else:
                uscs, uscs_expl = "CH", "Clay (high plasticity)."

    # AASHTO
    if P2 <= 35:
        if P2 <= 15 and P4 <= 30 and PI <= 6:
            aashto = "A-1-a"
        elif P2 <= 25 and P4 <= 50 and PI <= 6:
            aashto = "A-1-b"
        elif P2 <= 35 and P4 > 0:
            if LL <= 40 and PI <= 10: aashto = "A-2-4"
            elif LL >= 41 and PI <= 10: aashto = "A-2-5"
            elif LL <= 40 and PI >= 11: aashto = "A-2-6"
            elif LL >= 41 and PI >= 11: aashto = "A-2-7"
            else: aashto = "A-2"
        else: aashto = "A-3"
    else:
        if LL <= 40 and PI <= 10: aashto = "A-4"
        elif LL >= 41 and PI <= 10: aashto = "A-5"
        elif LL <= 40 and PI >= 11: aashto = "A-6"
        else:
            aashto = "A-7-5" if PI <= (LL - 30) else "A-7-6"

    a, b = max(P2-35,0), max(P2-15,0)
    c, d = max(LL-40,0), max(PI-10,0)
    GI = floor(0.2*a + 0.005*a*c + 0.01*b*d)

    aashto_expl = f"{aashto} (GI = {GI})"
    char_summary = ENGINEERING_CHARACTERISTICS.get("Silt", {})
    if uscs.startswith(("G","S")): char_summary = ENGINEERING_CHARACTERISTICS.get("Coarse sand", {})
    if uscs.startswith(("M","C","O","H")): char_summary = ENGINEERING_CHARACTERISTICS.get("Silt", {})

    result_text = f"USCS: **{uscs}** — {uscs_expl}\n\nAASHTO: **{aashto_expl}**"
    return result_text, uscs, aashto, GI, char_summary


# ==============================================================
#  CLASSIFIER CHATBOT (stepwise with retention)
# ==============================================================
# Exact dropdowns (text shown), backend mapping to integers
DILATANCY_OPTS = [
    "Quick slow", "None-Very slow", "Slow", "Slow-none", "None", "Null?"
]
TOUGHNESS_OPTS = [
    "None", "Medium", "Slight?", "Slight-Medium?", "High", "Null"
]
DRY_STRENGTH_OPTS = [
    "None - slight", "Medium - high", "Slight - Medium", "High - Very high", "Null?"
]

DIL_MAP = {txt: i+1 for i, txt in enumerate(DILATANCY_OPTS)}       # 1..6
TOUGH_MAP = {txt: i+1 for i, txt in enumerate(TOUGHNESS_OPTS)}      # 1..6
DRY_MAP = {txt: i+1 for i, txt in enumerate(DRY_STRENGTH_OPTS)}     # 1..5

def classifier_chatbot(site: str):
    st.markdown("🤖 **GeoMate:** Hello there! I am the soil classifier. Ready to start!")

    # Current step
    state = ss.steps["classifier"][site]
    step = state["step"]
    inputs = ss.cls_inputs[site]

    # Helper renderers for “Back/Next”
    def nav_buttons(next_enabled=True, back_enabled=True):
        cols = st.columns(2)
        with cols[0]:
            if back_enabled and step > 0 and st.button("⬅️ Back"):
                state["step"] = max(0, step - 1)
                st.rerun()
        with cols[1]:
            if next_enabled and st.button("➡️ Next"):
                state["step"] = step + 1
                st.rerun()

    # Step 0: Organic
    if step == 0:
        val = st.radio("Is the soil organic?", ["No", "Yes"], index=0 if inputs.get("opt","n")=="n" else 1, key=f"{site}_opt_radio")
        inputs["opt"] = "y" if val == "Yes" else "n"
        save_site_info(site, "Organic", val)
        nav_buttons(next_enabled=True, back_enabled=False)
        return

    # Step 1: % passing #200
    if step == 1:
        p2 = st.number_input("% Passing #200 (0–100)", min_value=0.0, max_value=100.0, value=float(inputs.get("P2", 0.0)), key=f"{site}_P2")
        inputs["P2"] = p2
        save_site_info(site, "% Passing #200", p2)
        nav_buttons()
        return

    # Decision fork by P2:
    P2 = float(inputs.get("P2", 0.0))

    # If organic -> we can classify right away, but still show a summary step
    if inputs.get("opt","n") == "y":
        st.info("Organic soil path selected → USCS: Pt")
        if st.button("🔍 Classify Now"):
            res_text, uscs, aashto, GI, chars = uscs_aashto_verbatim(inputs)
            st.success(res_text)
            save_site_info(site, "USCS", uscs)
            save_site_info(site, "AASHTO", aashto)
            save_site_info(site, "Group Index", GI)
        nav_buttons()
        return

    # Step 2: % passing #4
    if step == 2:
        p4 = st.number_input("% Passing #4 (0–100)", min_value=0.0, max_value=100.0, value=float(inputs.get("P4", 0.0)), key=f"{site}_P4")
        inputs["P4"] = p4
        save_site_info(site, "% Passing #4", p4)
        nav_buttons()
        return

    P4 = float(inputs.get("P4", 0.0))

    # Step 3+: by coarse vs fine
    if P2 <= 50:
        # Coarse soils path
        st.caption("Coarse-grained path (P2 ≤ 50)")

        # Step 3: D-sizes + Atterberg (for fines cases)
        if step == 3:
            c1, c2, c3 = st.columns(3)
            d60 = c1.number_input("D60 (mm)", min_value=0.0, value=float(inputs.get("D60", 0.0)), key=f"{site}_D60")
            d30 = c2.number_input("D30 (mm)", min_value=0.0, value=float(inputs.get("D30", 0.0)), key=f"{site}_D30")
            d10 = c3.number_input("D10 (mm)", min_value=0.0, value=float(inputs.get("D10", 0.0)), key=f"{site}_D10")
            inputs["D60"], inputs["D30"], inputs["D10"] = d60, d30, d10
            save_site_info(site, "D-values (mm)", {"D60": d60, "D30": d30, "D10": d10})
            nav_buttons()
            return

        # Step 4: Atterberg limits
        if step == 4:
            LL = st.number_input("Liquid Limit (LL)", min_value=0.0, max_value=200.0, value=float(inputs.get("LL", 0.0)), key=f"{site}_LL")
            PL = st.number_input("Plastic Limit (PL)", min_value=0.0, max_value=200.0, value=float(inputs.get("PL", 0.0)), key=f"{site}_PL")
            inputs["LL"], inputs["PL"] = LL, PL
            save_site_info(site, "Atterberg Limits", {"LL": LL, "PL": PL})
            nav_buttons()
            return

        # Step 5: Classify
        if step == 5:
            if st.button("🔍 Run Classification"):
                res_text, uscs, aashto, GI, chars = uscs_aashto_verbatim(inputs)
                st.success(res_text)
                save_site_info(site, "USCS", uscs)
                save_site_info(site, "AASHTO", aashto)
                save_site_info(site, "Group Index", GI)
                save_site_info(site, "Engineering Characteristics", chars)
            nav_buttons(next_enabled=False)
            return

    else:
        # Fine soils path (P2 > 50)
        st.caption("Fine-grained path (P2 > 50)")

        # Step 3: Atterberg limits (needed now)
        if step == 3:
            LL = st.number_input("Liquid Limit (LL)", min_value=0.0, max_value=200.0, value=float(inputs.get("LL", 0.0)), key=f"{site}_LL")
            PL = st.number_input("Plastic Limit (PL)", min_value=0.0, max_value=200.0, value=float(inputs.get("PL", 0.0)), key=f"{site}_PL")
            inputs["LL"], inputs["PL"] = LL, PL
            save_site_info(site, "Atterberg Limits", {"LL": LL, "PL": PL})
            nav_buttons()
            return

        # Step 4: Descriptors — text dropdowns mapped to integers
        if step == 4:
            dil_txt = st.selectbox("Dilatancy", DILATANCY_OPTS, index=(list(DIL_MAP).index(inputs.get("dil_txt", DILATANCY_OPTS[1])) if "dil_txt" in inputs else 1), key=f"{site}_DIL")
            tou_txt = st.selectbox("Toughness", TOUGHNESS_OPTS, index=(list(TOUGH_MAP).index(inputs.get("tou_txt", TOUGHNESS_OPTS[0])) if "tou_txt" in inputs else 0), key=f"{site}_TOU")
            dry_txt = st.selectbox("Dry Strength", DRY_STRENGTH_OPTS, index=(list(DRY_MAP).index(inputs.get("dry_txt", DRY_STRENGTH_OPTS[0])) if "dry_txt" in inputs else 0), key=f"{site}_DRY")

            # store both text & numeric codes
            inputs["dil_txt"], inputs["nDIL"] = dil_txt, DIL_MAP[dil_txt]
            inputs["tou_txt"], inputs["nTG"]  = tou_txt, TOUGH_MAP[tou_txt]
            inputs["dry_txt"], inputs["nDS"]  = dry_txt, DRY_MAP[dry_txt]

            save_site_info(site, "Dilatancy", dil_txt)
            save_site_info(site, "Toughness", tou_txt)
            save_site_info(site, "Dry Strength", dry_txt)
            nav_buttons()
            return

        # Step 5: Classify
        if step == 5:
            if st.button("🔍 Run Classification"):
                res_text, uscs, aashto, GI, chars = uscs_aashto_verbatim(inputs)
                st.success(res_text)
                save_site_info(site, "USCS", uscs)
                save_site_info(site, "AASHTO", aashto)
                save_site_info(site, "Group Index", GI)
                save_site_info(site, "Engineering Characteristics", chars)
            nav_buttons(next_enabled=False)
            return


# ==============================================================
#  LOCATOR (chatty minimal)
# ==============================================================
import ee, geemap, json, os, streamlit as st
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.pagesizes import A4

# ----------------- EARTH ENGINE AUTH -----------------
EE_READY = False
try:
    service_json = st.secrets["EARTH_ENGINE_KEY"]
    if isinstance(service_json, str):
        service_json = json.loads(service_json)
    credentials = ee.ServiceAccountCredentials(
        email=service_json["client_email"],
        key_data=json.dumps(service_json)
    )
    ee.Initialize(credentials)
    EE_READY = True
except Exception as e:
    st.error(f"⚠️ Earth Engine init failed: {e}")
    EE_READY = False


# ----------------- SAVE SITE INFO -----------------
def save_site_info(site: str, key: str, value):
    if "soil_description_site" not in st.session_state:
        st.session_state.soil_description_site = {}
    if site not in st.session_state.soil_description_site:
        st.session_state.soil_description_site[site] = {}
    st.session_state.soil_description_site[site][key] = value


# ----------------- ADD EARTH ENGINE LAYERS -----------------
def add_datasets_to_map(m, site, lat, lon):
    try:
        # Soil dataset (OpenLandMap USDA texture class)
        soil = ee.Image("OpenLandMap/SOL/SOL_TEXTURE-CLASS_USDA-TT_M/v02")
        soil_val = soil.sample(ee.Geometry.Point([lon, lat]), scale=250).first().getInfo()
        save_site_info(site, "soil_texture_value", soil_val)

        # Flood dataset (JRC Global Surface Water)
        flood = ee.ImageCollection("JRC/GSW1_4/YearlyHistory").mosaic()
        flood_val = flood.sample(ee.Geometry.Point([lon, lat]), scale=30).first().getInfo()
        save_site_info(site, "flood_risk_value", flood_val)

        # Elevation (SRTM)
        elevation = ee.Image("USGS/SRTMGL1_003")
        elev_val = elevation.sample(ee.Geometry.Point([lon, lat]), scale=30).first().getInfo()
        save_site_info(site, "elevation", elev_val)

        # Seismic/Environmental dataset (SEDAC NDGain placeholder)
        seismic = ee.Image("SEDAC/ND-GAIN/2015")
        seismic_val = seismic.sample(ee.Geometry.Point([lon, lat]), scale=1000).first().getInfo()
        save_site_info(site, "seismic_risk_value", seismic_val)

        # Visualization styles
        soil_vis = {"min": 1, "max": 12, "palette": ["ffffb2","fd8d3c","f03b20","bd0026"]}
        flood_vis = {"palette": ["0000ff"]}
        elev_vis = {"min": 0, "max": 3000, "palette": ["006633","E5FFCC","662A00","DDBB99","FFFFFF"]}
        seismic_vis = {"min": 0, "max": 100, "palette": ["green", "yellow", "red"]}

        # Add layers to map
        m.addLayer(soil, soil_vis, "Soil Texture")
        m.addLayer(flood, flood_vis, "Flood Risk")
        m.addLayer(elevation, elev_vis, "Elevation / Topography")
        m.addLayer(seismic, seismic_vis, "Seismic/Environmental Risk")

        # Center map
        m.setCenter(lon, lat, 8)

    except Exception as e:
        st.error(f"⚠️ Adding datasets failed: {e}")


# ----------------- LOCATOR CHATBOT -----------------
def locator_chat(site: str):
    st.markdown("🤖 **GeoMate:** Share your Area of Interest coordinates.")

    lat = st.number_input("Latitude", value=float(st.session_state.soil_description_site.get(site, {}).get("lat", 0.0)))
    lon = st.number_input("Longitude", value=float(st.session_state.soil_description_site.get(site, {}).get("lon", 0.0)))

    save_site_info(site, "lat", lat)
    save_site_info(site, "lon", lon)

    if EE_READY:
        try:
            m = geemap.Map(center=[lat or 0.0, lon or 0.0], zoom=6)
            add_datasets_to_map(m, site, lat, lon)

            # Save map snapshot for report
            map_path = f"map_{site}.png"
            m.to_image(out_path=map_path, zoom=6, dimensions=(600, 400))
            save_site_info(site, "map_snapshot", map_path)

            # Display map
            m.to_streamlit(height=500)
            st.success("✅ Earth Engine map loaded with soil, flood, seismic, and topography layers.")
        except Exception as e:
            st.error(f"🌍 Earth Engine map failed: {e}")
    else:
        st.warning("⚠️ Earth Engine not available — map disabled.")


# ----------------- PDF REPORT -----------------
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet
import os

def generate_geotech_report(site: str, filename="geotech_report.pdf"):
    if "soil_description_site" not in st.session_state or site not in st.session_state.soil_description_site:
        st.error("❌ No site data available.")
        return

    data = st.session_state.soil_description_site[site]
    doc = SimpleDocTemplate(filename, pagesize=A4)
    styles = getSampleStyleSheet()
    content = []

    # Title
    content.append(Paragraph(f"<b>Geotechnical Report for {site}</b>", styles["Title"]))
    content.append(Spacer(1, 12))

    # General site data
    for key, value in data.items():
        if key not in ["map_snapshot"]:
            content.append(Paragraph(f"<b>{key}:</b> {value}", styles["Normal"]))
            content.append(Spacer(1, 6))

    # Add site map snapshot if available
    if "map_snapshot" in data and os.path.exists(data["map_snapshot"]):
        content.append(Spacer(1, 12))
        content.append(Paragraph("<b>Site Map:</b>", styles["Heading2"]))
        content.append(Image(data["map_snapshot"], width=400, height=300))

    # Add legends section
    content.append(Spacer(1, 12))
    content.append(Paragraph("<b>Legend</b>", styles["Heading2"]))

    legend_data = [
        ["Layer", "Description", "Color Representation"],
        ["Soil Texture", "USDA Texture Classes", "🟨 → 🟧 → 🟥 → 🟥 Dark"],
        ["Flood Risk", "Water extent (JRC GSW)", "🟦"],
        ["Elevation / Topography", "SRTM Elevation", "🟩 → 🟨 → 🟫 → ⬜"],
        ["Seismic Risk", "NDGain Risk Index", "🟩 → 🟨 → 🟥"]
    ]

    table = Table(legend_data, colWidths=[120, 240, 140])
    table.setStyle(TableStyle([
        ("BACKGROUND", (0, 0), (-1, 0), colors.grey),
        ("TEXTCOLOR", (0, 0), (-1, 0), colors.whitesmoke),
        ("ALIGN", (0, 0), (-1, -1), "LEFT"),
        ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
        ("BOTTOMPADDING", (0, 0), (-1, 0), 8),
        ("BACKGROUND", (0, 1), (-1, -1), colors.beige),
        ("GRID", (0, 0), (-1, -1), 0.5, colors.black),
    ]))

    content.append(table)

    # Build PDF
    doc.build(content)
    st.success(f"📄 Full Geotechnical Report generated: {filename}")
    return filename
# ==============================================================
#  GEOmate Ask (LLM + memory + fact capture)
# ==============================================================
FACT_PATTERNS = [
    ("Load Bearing Capacity", r"(?:bearing\s*capacity|q_?ult|qult)\s*[:=]?\s*([\d\.]+)\s*(k?pa|tsf|ksf|psi|mpa)?"),
    ("% Compaction", r"(?:%?\s*compaction|relative\s*compaction)\s*[:=]?\s*([\d\.]+)\s*%"),
    ("Skin Shear Strength", r"(?:skin\s*shear\s*strength|adhesion|α\s*su)\s*[:=]?\s*([\d\.]+)\s*(k?pa|psf|kpa)"),
]

def extract_and_save_facts(site: str, text: str, answer: str):
    lower = f"{text}\n{answer}".lower()
    for key, pattern in FACT_PATTERNS:
        m = re.search(pattern, lower)
        if m:
            val = m.group(1)
            unit = (m.group(2) or "").strip()
            save_site_info(site, key, f"{val} {unit}".strip())

def geomate_ask(site: str):
    st.markdown("🤖 **GeoMate:** Ask anything about your site’s soils or design.")
    if ss.secrets_status["groq_ok"] and HAVE_LANGCHAIN:
        init_rag()

    q = st.text_input("Your question (press Enter to send)", key=f"{site}_ask")
    if q:
        resp = rag_ask(q)
        st.markdown(f"**GeoMate:** {resp}")
        extract_and_save_facts(site, q, resp)


# ==============================================================
#  REPORTS CHATBOT
# ==============================================================
REPORT_QUESTIONS = [
    ("Load Bearing Capacity", "What is the soil bearing capacity? (e.g., 150 kPa)"),
    ("Skin Shear Strength", "What is the skin shear strength? (e.g., 25 kPa)"),
    ("% Compaction", "What is the required % relative compaction? (e.g., 95 %)"),
    ("Rate of Consolidation", "What is the rate of consolidation / settlement time?"),
    ("Nature of Construction", "What is the nature of construction? (e.g., G+1 Residential, Tank, Retaining wall)"),
]

def reports_chatbot(site: str):
    st.markdown("🤖 **GeoMate:** I’ll collect details for a full geotechnical report. You can type **skip** to move on.")

    s = ss.steps["reports"][site]["step"]
    answers = ss.reports_inputs[site]

    # ask current question
    if s < len(REPORT_QUESTIONS):
        key, prompt = REPORT_QUESTIONS[s]
        st.write(f"**Q{s+1}. {prompt}**")
        default_val = str(ss.soil_description_site.get(site, {}).get(key, answers.get(key, "")))
        ans = st.text_input("Your answer", value=default_val, key=f"{site}_rep_{s}")
        cols = st.columns(2)
        with cols[0]:
            if st.button("➡️ Next", key=f"{site}_rep_next_{s}"):
                if ans and ans.strip().lower() != "skip":
                    answers[key] = ans.strip()
                    save_site_info(site, key, ans.strip())
                ss.steps["reports"][site]["step"] = s + 1
                st.rerun()
        with cols[1]:
            if s > 0 and st.button("⬅️ Back", key=f"{site}_rep_back_{s}"):
                ss.steps["reports"][site]["step"] = s - 1
                st.rerun()
    else:
        st.success("All questions answered (or skipped). Generate your report when ready.")
        if st.button("📄 Generate Full Geotechnical Report", key=f"{site}_gen_pdf"):
            fname = generate_report_pdf(site)
            with open(fname, "rb") as f:
                st.download_button("⬇️ Download Report", f, file_name=fname, mime="application/pdf")


# ==============================================================
#  PDF EXPORT
# ==============================================================
def generate_report_pdf(site: str) -> str:
    data = ss.soil_description_site.get(site, {})
    fname = f"{site}_geotechnical_report.pdf"
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", "B", 16)
    pdf.cell(0, 10, "GeoMate Geotechnical Report", ln=True, align="C")
    pdf.ln(6)
    pdf.set_font("Arial", "", 12)
    pdf.cell(0, 8, f"Site: {site}", ln=True)
    pdf.cell(0, 8, f"Date: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}", ln=True)
    pdf.ln(4)

    # Summary sections
    pdf.set_font("Arial", "B", 13)
    pdf.cell(0, 8, "Collected Parameters", ln=True)
    pdf.set_font("Arial", "", 11)
    for k, v in data.items():
        pdf.multi_cell(0, 6, f"{k}: {json.dumps(v) if isinstance(v, (dict, list)) else v}")
    pdf.ln(2)

    # If classification stored, show clearly
    uscs = data.get("USCS", None)
    aashto = data.get("AASHTO", None)
    gi = data.get("Group Index", None)
    if uscs or aashto:
        pdf.set_font("Arial", "B", 13)
        pdf.cell(0, 8, "Classification", ln=True)
        pdf.set_font("Arial", "", 11)
        if uscs: pdf.cell(0, 6, f"USCS: {uscs}", ln=True)
        if aashto: pdf.cell(0, 6, f"AASHTO: {aashto}", ln=True)
        if gi is not None: pdf.cell(0, 6, f"Group Index (GI): {gi}", ln=True)
        pdf.ln(2)

    pdf.output(fname)
    return fname


# ==============================================================
#  PAGE SECTIONS
# ==============================================================
def landing_page():
    st.markdown(
        """
        <style>
        .stApp { background: linear-gradient(180deg, #060606 0%, #0f0f0f 100%); color: #e9eef6; }
        .landing-card { background: linear-gradient(180deg, rgba(255,122,0,0.06), rgba(255,122,0,0.02));
                        border-radius: 12px; padding: 18px; margin-bottom: 14px; border: 1px solid rgba(255,122,0,0.08); }
        </style>
        """,
        unsafe_allow_html=True,
    )
    st.markdown("<h1 style='color:#FF8C00;margin-bottom:0'>GeoMate V2</h1>", unsafe_allow_html=True)
    st.caption("AI geotechnical copilot — soil recognition, classification, locator, RAG, and reports")

    st.markdown("### Startup Status")
    check_secrets_banner()

    col1, col2 = st.columns([2, 1])
    with col1:
        st.markdown("<div class='landing-card'>", unsafe_allow_html=True)
        st.subheader("What GeoMate does")
        st.markdown(
            """
            - **Soil Classifier:** USCS & AASHTO (verbatim logic) with guided, chatbot-style inputs.  
            - **Locator:** (EE) Set AOIs and preview layers (if EE credentials available).  
            - **GeoMate Ask:** Session-memory LLM with fact capture into site variables.  
            - **Reports:** Chatbot gathers all remaining design data and generates a PDF.  
            """
        )
        st.markdown("</div>", unsafe_allow_html=True)

    with col2:
        st.markdown("<div class='landing-card'>", unsafe_allow_html=True)
        st.subheader("Live Project")
        n_sites = len(ss.sites)
        n_cls = sum(1 for s in ss.soil_description_site.values() if "USCS" in s or "AASHTO" in s)
        st.metric("Sites", n_sites)
        st.metric("Classified Sites", n_cls)
        st.markdown("</div>", unsafe_allow_html=True)


def sidebar_controls():
    with st.sidebar:
        st.markdown("<div style='text-align:center; padding:6px 0;'><h3 style='color:#FF8C00; margin:0;'>GeoMate V2</h3></div>", unsafe_allow_html=True)

        # Model selector (variable model name)
        st.subheader("LLM Model")
        ss.MODEL_NAME = st.selectbox(
            "Model",
            [
                "llama-3.1-70b-versatile",
                "llama-3.1-8b-instant",
                "mixtral-8x7b-32768",
                "gemma-7b-it"
            ],
            index=["llama-3.1-70b-versatile","llama-3.1-8b-instant","mixtral-8x7b-32768","gemma-7b-it"].index(ss.MODEL_NAME)
            if ss.MODEL_NAME in ["llama-3.1-70b-versatile","llama-3.1-8b-instant","mixtral-8x7b-32768","gemma-7b-it"] else 0
        )

        # Site selector / creator
        st.subheader("Sites")
        site_choice = st.selectbox("Current site", ss.sites, index=ss.sites.index(ss.current_site))
        if site_choice != ss.current_site:
            ss.current_site = site_choice
            # ensure step entries for this site
            for key in ss.steps:
                if ss.current_site not in ss.steps[key]:
                    ss.steps[key][ss.current_site] = {"step": 0}
            if ss.current_site not in ss.cls_inputs:
                ss.cls_inputs[ss.current_site] = {}
            if ss.current_site not in ss.reports_inputs:
                ss.reports_inputs[ss.current_site] = {}
            st.rerun()

        new_site = st.text_input("New site name")
        if st.button("➕ Add site"):
            ns = new_site.strip() or f"site{len(ss.sites)+1}"
            if ns not in ss.sites:
                ss.sites.append(ns)
                ss.current_site = ns
                ss.soil_description_site.setdefault(ns, {})
                for key in ss.steps:
                    ss.steps[key][ns] = {"step": 0}
                ss.cls_inputs.setdefault(ns, {})
                ss.reports_inputs.setdefault(ns, {})
                st.success(f"Created {ns}")
                st.rerun()

        st.markdown("---")
        pages = ["Landing", "Locator", "Soil Classifier", "GeoMate Ask", "Reports"]
        choice = option_menu(
            menu_title="",
            options=pages,
            icons=["house", "geo-alt", "flask", "robot", "file-earmark-text"],
            default_index=0
        )
        return choice


def locator_page():
    st.header("🌍 Locator")
    locator_chat(ss.current_site)


def classifier_page():
    st.header("🧪 Soil Classifier — USCS & AASHTO")
    classifier_chatbot(ss.current_site)

    # Save classification snapshot (if present)
    data = ss.soil_description_site.get(ss.current_site, {})
    if any(k in data for k in ["USCS", "AASHTO"]):
        st.markdown("### Saved Classification")
        st.json({k: data[k] for k in data if k in ["USCS", "AASHTO", "Group Index", "Engineering Characteristics"]})


def ask_page():
    st.header("🤖 GeoMate Ask — RAG with Session Memory")
    if not ss.secrets_status["groq_ok"] or not HAVE_LANGCHAIN:
        st.warning("LLM not available (Groq key missing or LangChain not installed).")
    geomate_ask(ss.current_site)

    st.markdown("### Current Site Facts")
    st.json(ss.soil_description_site.get(ss.current_site, {}))


def reports_page():
    st.header("📑 Reports")
    reports_chatbot(ss.current_site)


# ==============================================================
#  MAIN
# ==============================================================
def main():
    choice = sidebar_controls()
    if choice == "Landing":
        landing_page()
    elif choice == "Locator":
        locator_page()
    elif choice == "Soil Classifier":
        classifier_page()
    elif choice == "GeoMate Ask":
        ask_page()
    elif choice == "Reports":
        reports_page()


if __name__ == "__main__":
    main()