# ======================================================
# HRNet Fashion Landmark Measurement - Streamlit UI
# ======================================================

import os
import sys
import json
import argparse
import numpy as np
import streamlit as st
from PIL import Image
import gdown
import io
import time
from dotenv import load_dotenv

# Load environment variables at the very beginning
load_dotenv()

# Defer heavy imports to optimize startup time
torch = None
cv2 = None
cloudinary = None

# ------------------------------------------------------
# SERVER-SIDE AUTH GUARD
# ------------------------------------------------------
import firebase_config
import auth
import firestore_service

# Initialize Firebase (Server-Side)
firebase_config.initialize_firebase()

is_embedded = st.query_params.get("embed") == "true"

# Auth is optional — try to verify token if provided, but don't block access
if "uid" not in st.session_state:
    token = st.query_params.get("token")
    if token:
        try:
            with st.spinner("Verifying session..."):
                user_uid = auth.verify_user_from_query()
                st.session_state["uid"] = user_uid
            if not is_embedded:
                st.query_params.clear()
                st.rerun()
        except Exception as e:
            # Token invalid — still allow access, just without Firestore save
            st.session_state["uid"] = None
    else:
        st.session_state["uid"] = None


# ------------------------------------------------------
# STREAMLIT CONFIG
# ------------------------------------------------------
st.set_page_config(
    page_title="DASH AI Fit",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# ------------------------------------------------------
# PREMIUM CSS INJECTION
# ------------------------------------------------------
CSS_CODE = r"""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');
    
    html, body, [data-testid="stAppViewContainer"] {
        background-color: #f7f9fc !important;
        color: #09090b !important;
        font-family: 'Inter', sans-serif;
    }
    
    .block-container {
        max-width: 480px !important;
        padding-top: 3rem !important;
        padding-bottom: 5rem !important;
        margin: auto;
    }

    .header-container { text-align: center; margin-bottom: 2rem; }
    .main-subtitle { font-size: 1.1rem; font-weight: 500; color: #52525b; }
    .main-subtitle b { color: #000000; font-weight: 700; }

    .step-header {
        font-size: 0.75rem;
        font-weight: 700;
        color: #18181b !important;
        text-transform: uppercase;
        margin-bottom: 0.8rem;
        margin-top: 2rem;
        letter-spacing: 0.5px;
        display: flex;
        align-items: center;
        gap: 0.5rem;
        border-bottom: 1px solid #f4f4f5;
        padding-bottom: 0.5rem;
    }

    /* ----------------------------------------------------------------
       FILE UPLOADER - PREMIUM BLACK BORDER
       ---------------------------------------------------------------- */
    [data-testid="stFileUploader"] {
        background-color: #ffffff !important;
        border: 2px solid #000000 !important;
        border-radius: 24px !important;
        min-height: 250px !important; 
        position: relative !important;
        box-shadow: 0 10px 30px -10px rgba(0, 0, 0, 0.1) !important;
        padding: 0 !important;
        overflow: hidden !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1);
    }
    
    [data-testid="stFileUploader"]:hover {
        transform: translateY(-2px);
        box-shadow: 0 20px 40px -15px rgba(0, 0, 0, 0.15) !important;
    }
    
    [data-testid="stFileUploader"] section {
        position: absolute !important;
        top: 0;
        left: 0;
        width: 100% !important;
        height: 100% !important;
        z-index: 20;
        opacity: 0 !important;
        cursor: pointer !important;
    }

    /* Custom Arrow */
    [data-testid="stFileUploader"]::before {
        content: "\2B06"; 
        font-size: 4rem;
        position: absolute;
        top: 35%;
        left: 50%;
        transform: translate(-50%, -50%);
        z-index: 10;
        pointer-events: none;
        color: #000000 !important;
        display: block !important;
        visibility: visible !important;
        transition: transform 0.3s ease;
    }
    
    [data-testid="stFileUploader"]:hover::before {
        transform: translate(-50%, -60%);
    }
    
    /* Custom Text - Polished & Simplified */
    [data-testid="stFileUploader"]::after {
        content: "Upload flat laid pant image";
        text-align: center;
        font-family: 'Inter', sans-serif;
        font-weight: 700;
        font-size: 1.1rem;
        color: #000000 !important;
        position: absolute;
        top: 70%;
        left: 50%;
        transform: translate(-50%, -50%);
        width: 90%;
        z-index: 10;
        pointer-events: none;
        display: block !important;
        visibility: visible !important;
    }

    /* ----------------------------------------------------------------
       WAIST INPUT - MINIMALIST
       ---------------------------------------------------------------- */
    [data-testid="stNumberInput"] button {
        display: none !important;
    }
    input[type=number]::-webkit-inner-spin-button, 
    input[type=number]::-webkit-outer-spin-button { 
        -webkit-appearance: none; 
        margin: 0; 
    }
    input[type=number] {
        -moz-appearance: textfield;
    }

    .stNumberInput input {
        text-align: center !important;
        font-size: 1.75rem !important;
        font-weight: 800 !important;
        height: 64px !important;
        border-radius: 16px !important;
        border: 2px solid #e2e8f0 !important;
        background-color: #ffffff !important;
        color: #000000 !important;
        transition: all 0.2s;
    }
    .stNumberInput input:focus {
        border-color: #000000 !important;
        box-shadow: 0 0 0 4px rgba(0,0,0,0.05) !important;
    }

    /* Analyze Button */
    .stButton > button {
        width: 100%;
        background: #000000 !important;
        color: #ffffff !important;
        border: none !important;
        border-radius: 16px !important;
        padding: 1.25rem !important;
        font-size: 1.1rem !important;
        font-weight: 800 !important;
        text-transform: uppercase;
        letter-spacing: 1px;
        margin-top: 2rem;
        box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.2) !important;
        transition: all 0.3s ease;
    }
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 15px 30px -10px rgba(0, 0, 0, 0.3) !important;
    }

    .social-proof {
        text-align: center;
        font-size: 0.9rem;
        color: #64748b !important;
        margin-top: 1.5rem;
        font-weight: 600; 
    }

    #MainMenu, footer, header {visibility: hidden;}
    .stDeployButton {display:none;}
    [data-testid="stHeader"] {display: none;}
</style>
"""
st.markdown(CSS_CODE, unsafe_allow_html=True)

# ------------------------------------------------------
# LAZY IMPORTS & RESOURCE CACHING
# ------------------------------------------------------
def get_cloudinary():
    global cloudinary
    if cloudinary is None:
        import cloudinary
        import cloudinary.uploader
    return cloudinary

def init_cloudinary():
    """Configure Cloudinary using environment variables."""
    cld = get_cloudinary()
    try:
        cld.config(
            cloud_name=os.getenv("CLOUDINARY_CLOUD_NAME"),
            api_key=os.getenv("CLOUDINARY_API_KEY"),
            api_secret=os.getenv("CLOUDINARY_API_SECRET"),
            secure=True
        )
        return True
    except Exception as e:
        st.error(f"Failed to initialize Cloudinary: {e}")
        return False

# ------------------------------------------------------
# HRNET PATHS & SETUP
# ------------------------------------------------------
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
HRNET_ROOT = BASE_DIR
CFG_FILE = os.path.join(HRNET_ROOT, "experiments/deepfashion2/hrnet/w48_384x288_adam_lr1e-3.yaml")
MODEL_FILE = os.path.join(HRNET_ROOT, "lib/models/pose_hrnet-w48_384x288-deepfashion2_mAP_0.7017.pth")
MODEL_URL = "https://drive.google.com/file/d/1Srj8YJLu00aZm368sCtEIWB8Vkpr7kXA/view?usp=sharing"
sys.path.insert(0, os.path.join(HRNET_ROOT, "lib"))

def ensure_model_exists():
    if not os.path.exists(MODEL_FILE):
        os.makedirs(os.path.dirname(MODEL_FILE), exist_ok=True)
        with st.spinner("Initializing AI engine..."):
            gdown.download(MODEL_URL, MODEL_FILE, quiet=True)

@st.cache_resource
def load_hrnet():
    # Helper to load heavy libraries only when needed
    global torch, cv2
    if torch is None:
        import torch
    if cv2 is None:
        import cv2

    ensure_model_exists()
    from config import cfg, update_config
    import models
    
    args = argparse.Namespace(cfg=CFG_FILE, opts=[], modelDir="", logDir="", dataDir="", prevModelDir="")
    update_config(cfg, args)
    cfg.freeze()
    model = models.pose_hrnet.get_pose_net(cfg, is_train=False)
    model.load_state_dict(torch.load(MODEL_FILE, map_location="cpu"), strict=False)
    model.eval()
    return model

# ------------------------------------------------------
# AI HELPERS
# ------------------------------------------------------
def load_image_portrait(pil_img):
    global cv2
    if cv2 is None: import cv2
    
    img = np.array(pil_img.convert("RGB"))
    h, w = img.shape[:2]
    if w > h:
        img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    return img

def round_half(x):
    return round(x * 2) / 2

# Constants
WAIST_L, WAIST_C, WAIST_R = 6, 0, 30
OUTER_ANKLE_L, OUTER_ANKLE_R = 14, 16
INNER_ANKLE_L, INNER_ANKLE_R = 9, 21
INSEAM = 27
FALLBACK_INNER_ANKLE_L = 108

def run_hrnet(img):
    global torch, cv2
    if torch is None: import torch
    if cv2 is None: import cv2
    
    model = load_hrnet() # Get cached model
    
    H, W = img.shape[:2]
    inp = cv2.resize(img, (288, 384)).astype(np.float32) / 255.0
    inp = inp.transpose(2, 0, 1)
    inp = torch.from_numpy(inp).unsqueeze(0)
    with torch.no_grad():
        heatmaps = model(inp).squeeze(0).cpu().numpy()
    pts = {}
    num_joints, hm_h, hm_w = heatmaps.shape
    for jid in range(num_joints):
        y, x = np.unravel_index(np.argmax(heatmaps[jid]), heatmaps[jid].shape)
        pts[jid] = (int(x * W / hm_w), int(y * H / hm_h))
    if FALLBACK_INNER_ANKLE_L in pts and pts[FALLBACK_INNER_ANKLE_L][1] > pts[INNER_ANKLE_L][1]:
        pts[INNER_ANKLE_L] = pts[FALLBACK_INNER_ANKLE_L]
    return pts

def dist(p, a, b):
    x1, y1 = p[a]
    x2, y2 = p[b]
    return np.hypot(x1 - x2, y1 - y2)

def draw_measurements(img, pts):
    global cv2
    if cv2 is None: import cv2
    
    draw = img.copy()
    connections = [
        (WAIST_L, WAIST_C), (WAIST_C, WAIST_R),
        (WAIST_L, OUTER_ANKLE_L), (WAIST_R, OUTER_ANKLE_R),
        (INNER_ANKLE_L, OUTER_ANKLE_L),
        (INNER_ANKLE_R, OUTER_ANKLE_R),
        (INNER_ANKLE_L, INSEAM),
        (INNER_ANKLE_R, INSEAM),
    ]
    for a, b in connections:
        cv2.line(draw, pts[a], pts[b], (60,120,180), 2)
    return draw

def resize_and_compress(image_arr, max_dim=1000, quality=75):
    """Optimizes the image: resizes to max_dim and compresses to JPEG."""
    try:
        pil_img = Image.fromarray(image_arr)
        width, height = pil_img.size
        if max(width, height) > max_dim:
            if width > height:
                new_width = max_dim
                new_height = int(height * (max_dim / width))
            else:
                new_height = max_dim
                new_width = int(width * (max_dim / height))
            pil_img = pil_img.resize((new_width, new_height), Image.BICUBIC)
        buf = io.BytesIO()
        pil_img.save(buf, format="JPEG", quality=quality, optimize=False)
        buf.seek(0)
        return buf
    except Exception as e:
        print(f"Optimization failed: {e}")
        return None

def upload_image_to_cloud(image_obj, folder):
    cld = get_cloudinary()
    try:
        if isinstance(image_obj, (np.ndarray, np.generic)):
            pil_img = Image.fromarray(image_obj)
            buf = io.BytesIO()
            pil_img.save(buf, format="JPEG", quality=80)
            buf.seek(0)
            obj_to_upload = buf
        else:
            obj_to_upload = image_obj

        response = cld.uploader.upload(
            obj_to_upload,
            folder=folder,
            resource_type="image",
            timeout=120
        )
        return response.get("secure_url")
    except Exception as e:
        print(f"Image upload failed: {e}")
        return None

# ------------------------------------------------------
# UI HEADER
# ------------------------------------------------------
st.markdown("""
<div class="header-container">
    <div class="main-subtitle">Find your perfect size in <b>10 seconds.</b></div>
</div>
""", unsafe_allow_html=True)

# ------------------------------------------------------
# STEP 1: NAME
# ------------------------------------------------------
st.markdown('<div class="step-header">STEP 1 — ENTER YOUR NAME</div>', unsafe_allow_html=True)
user_name = st.text_input("Enter your name", placeholder="Enter your name", label_visibility="collapsed")

# ------------------------------------------------------
# STEP 2: IMAGE UPLOAD
# ------------------------------------------------------
st.markdown('<div class="step-header">STEP 2 — upload a flat laid pant image</div>', unsafe_allow_html=True)
uploaded = st.file_uploader("Upload Image", type=["jpg","png","jpeg"], label_visibility="collapsed")

# ------------------------------------------------------
# STEP 3: WAIST SIZE (Editable Only)
# ------------------------------------------------------
st.markdown('<div class="step-header">STEP 3 — ENTER WAIST SIZE</div>', unsafe_allow_html=True)

# Initialize session state for waist if not exists
if 'waist_val' not in st.session_state:
    st.session_state.waist_val = 32

# Callback to sync input with session state
def update_waist_from_input():
    st.session_state.waist_val = st.session_state.waist_input

# Single editable number input
st.number_input(
    "Waist Size", 
    min_value=24, 
    max_value=60, 
    value=st.session_state.waist_val, 
    step=1, 
    key="waist_input",
    on_change=update_waist_from_input,
    label_visibility="collapsed"
)

# ------------------------------------------------------
# SOCIAL PROOF & ACTION
# ------------------------------------------------------
st.markdown('<div class="social-proof">Used by 500+ Users</div>', unsafe_allow_html=True)

run_ai = st.button("ANALYZE MY FIT")

# ------------------------------------------------------
# VALIDATION & LOGIC
# ------------------------------------------------------
if run_ai:
    if not user_name:
        st.warning("Please enter your name.")
        st.stop()
    if not uploaded:
        st.warning("Please upload a pant image.")
        st.stop()
    
    # Init Cloudinary now if needed
    if not init_cloudinary():
        st.warning("Cloudinary config missing.")

    try:
        img_pil = Image.open(uploaded)
        img = load_image_portrait(img_pil)

        with st.spinner("Analyzing your measurements..."):
            # Real AI Processing
            pts = run_hrnet(img)
            result_img = draw_measurements(img, pts)
            
            # Calculate Measurements
            waist_in = st.session_state.waist_val
            waist_px = dist(pts, WAIST_L, WAIST_C) + dist(pts, WAIST_C, WAIST_R)
            outseam_px = (dist(pts, WAIST_L, OUTER_ANKLE_L) + dist(pts, WAIST_R, OUTER_ANKLE_R)) / 2
            ankle_px = (dist(pts, INNER_ANKLE_L, OUTER_ANKLE_L) + dist(pts, INNER_ANKLE_R, OUTER_ANKLE_R)) / 2
            inseam_px = (dist(pts, INSEAM, INNER_ANKLE_L) + dist(pts, INSEAM, INNER_ANKLE_R)) / 2

            half_waist = (waist_in / 2) - 1
            outseam_val = round_half(outseam_px / waist_px * half_waist)
            ankle_val = round_half(ankle_px / waist_px * half_waist)
            inseam_val = round_half(inseam_px / waist_px * half_waist)

        # ------------------------------------------------------
        # RESULTS DISPLAY (Styled Card)
        # ------------------------------------------------------
        st.image(result_img, caption="AI Detection Result", use_container_width=True)
        
        results_html = f"""
        <style>
            .results-card {{
                background: #ffffff;
                border: 1px solid #f4f4f5;
                border-radius: 28px;
                padding: 2.5rem;
                color: #09090b;
                box-shadow: 0 20px 25px -5px rgba(0, 0, 0, 0.05);
                margin-top: 2rem;
            }}
            .results-header {{ font-size: 1.2rem; font-weight: 700; margin-bottom: 1.5rem; }}
            .size-badge {{ font-size: 4rem; font-weight: 800; line-height: 1; margin-bottom: 1.5rem; }}
            .detail-grid {{ display: grid; grid-template-columns: 1fr 1fr; gap: 1.5rem; border-top: 1px solid #f4f4f5; padding-top: 2rem; }}
            .detail-label {{ font-size: 0.7rem; color: #a1a1aa; text-transform: uppercase; font-weight: 800; letter-spacing: 1px; }}
            .detail-value {{ font-size: 1.25rem; font-weight: 700; }}
        </style>
        <div class="results-card">
            <div class="results-header">Hello {user_name},</div>
            <div class="detail-label">Recommended Size</div>
            <div class="size-badge">{int(waist_in)}</div>
            <div class="detail-grid">
                <div class="detail-item">
                    <div class="detail-label">Waist</div>
                    <div class="detail-value">{waist_in}in</div>
                </div>
                <div class="detail-item">
                    <div class="detail-label">Outseam</div>
                    <div class="detail-value">{outseam_val}in</div>
                </div>
                <div class="detail-item">
                    <div class="detail-label">Inseam</div>
                    <div class="detail-value">{inseam_val}in</div>
                </div>
                <div class="detail-item">
                    <div class="detail-label">Ankle</div>
                    <div class="detail-value">{ankle_val}in</div>
                </div>
            </div>
        </div>
        """
        st.markdown(results_html, unsafe_allow_html=True)

        # ------------------------------------------------------
        # POST MESSAGE TO PARENT (auto-save to localStorage)
        # ------------------------------------------------------
        import json as _json
        fit_payload = _json.dumps({
            "type": "dashFitResult",
            "name": user_name,
            "waist": waist_in,
            "outseam": float(outseam_val),
            "inseam": float(inseam_val),
            "ankle": float(ankle_val),
            "recommendedSize": str(int(waist_in)),
            "saved_at": str(__import__('datetime').datetime.utcnow().isoformat())
        })
        post_msg_js = f"""
        <script>
            (function() {{
                var payload = {fit_payload};
                // Aggressively target all possible parent elements in the iframe stack
                try {{ window.top.postMessage(payload, '*'); }} catch(e) {{}}
                try {{ window.parent.postMessage(payload, '*'); }} catch(e) {{}}
                try {{ window.parent.parent.postMessage(payload, '*'); }} catch(e) {{}}
                
                console.log('✅ Fit data emitted from Streamlit:', payload);
            }})();
        </script>
        """
        import streamlit.components.v1 as components
        components.html(post_msg_js, height=0, width=0)
        
        st.success("✅ Measurements saved successfully! You can now access them via 'Custom Fit' on the DASH website.")

        # ------------------------------------------------------
        # AUTOMATIC SAVE (Server-Side, only if logged in)
        # ------------------------------------------------------
        if st.session_state.get("uid"):
            with st.spinner("Syncing to cloud profile..."):
                output_buf = resize_and_compress(result_img)
                if output_buf:
                    image_url = upload_image_to_cloud(output_buf, "dashclothing/scans")
                    if image_url:
                        scan_data = {
                            "name": user_name,
                            "imageUrl": image_url,
                            "waist": waist_in,
                            "outseam": outseam_val,
                            "inseam": inseam_val,
                            "ankle": ankle_val,
                            "status": "completed"
                        }
                        try:
                            doc_id = firestore_service.save_prediction(st.session_state["uid"], scan_data)
                            st.success("Result Saved!")
                        except Exception as e:
                            st.error(f"Failed to save to database: {str(e)}")
                    else:
                        st.warning("Image upload skipped.")
    except Exception as e:
        st.error(f"An error occurred: {e}")