"""
Adeifix — AI Construction Takeoff
Gradio app for HuggingFace Spaces (T4 GPU)
"""

import os
import time
import torch
import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
REPO_ID = "rimashussain/gemma4-cubicasa-floorplan"

SYSTEM_PROMPT = (
    "You are Adeifix, an expert AI construction takeoff assistant. "
    "When shown an architectural floor plan, you analyze it and provide:\n"
    "1. A list of all rooms with their names\n"
    "2. Dimensions and area estimates where visible\n"
    "3. Count of doors and windows per room\n"
    "4. Wall types (interior/exterior) if distinguishable\n"
    "5. Any special features (stairs, elevators, wet areas)\n\n"
    "Be precise, use metric units, and format output clearly."
)

EXAMPLE_PROMPTS = [
    "Analyze this floor plan. List all rooms with dimensions.",
    "Count all doors and windows in this plan.",
    "Identify wet areas (kitchens, bathrooms) and list plumbing fixtures.",
    "Estimate total built-up area and carpet area.",
    "List all rooms and their approximate square meter areas.",
]

# ---------------------------------------------------------------------------
# Load model
# ---------------------------------------------------------------------------
print("🔧 Loading model (this may take a few minutes)...")

processor = AutoProcessor.from_pretrained(REPO_ID)
model = AutoModelForImageTextToText.from_pretrained(
    REPO_ID,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_4bit=True,
)

print("✅ Model loaded and ready!")


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def analyze_floorplan(image: Image.Image, prompt: str) -> str:
    if image is None:
        return "⚠️ Please upload a floor plan image."

    if not prompt.strip():
        prompt = EXAMPLE_PROMPTS[0]

    # Resize large images
    image.thumbnail((1024, 1024), Image.LANCZOS)

    # Build messages
    messages = [
        {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        },
    ]

    start = time.time()

    try:
        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            add_generation_prompt=True,
        ).to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=2048,
                temperature=0.1,
                do_sample=True,
            )

        # Decode only new tokens
        input_len = inputs["input_ids"].shape[-1]
        result = processor.decode(outputs[0][input_len:], skip_special_tokens=True)
        elapsed = time.time() - start

        return f"{result}\n\n---\n⏱ Inference: {elapsed:.1f}s"

    except Exception as e:
        return f"❌ Error during analysis: {str(e)}"


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
css = """
    .gradio-container { max-width: 1100px !important; }
    .header-text { text-align: center; margin-bottom: 8px; }
    .header-text h1 { font-size: 2.2em; font-weight: 700; margin-bottom: 4px; }
    .header-text p { opacity: 0.7; font-size: 1.05em; }
    footer { display: none !important; }
"""

with gr.Blocks(css=css, title="Adeifix — AI Construction Takeoff", theme=gr.themes.Soft()) as demo:

    gr.HTML("""
        <div class="header-text">
            <h1>📐 Adeifix</h1>
            <p>AI-Powered Construction Takeoff — Upload a floor plan for instant analysis</p>
        </div>
    """)

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            image_input = gr.Image(label="Floor Plan", type="pil", height=400)
            prompt_input = gr.Textbox(
                label="Analysis Prompt",
                placeholder="e.g. List all rooms with dimensions...",
                value=EXAMPLE_PROMPTS[0],
                lines=2,
            )
            gr.Examples(
                examples=[[p] for p in EXAMPLE_PROMPTS],
                inputs=[prompt_input],
                label="Quick Prompts",
            )
            analyze_btn = gr.Button("🔍 Analyze Floor Plan", variant="primary", size="lg")

        with gr.Column(scale=1):
            output = gr.Textbox(label="Analysis Output", lines=22, show_copy_button=True)

    analyze_btn.click(fn=analyze_floorplan, inputs=[image_input, prompt_input], outputs=output)

    gr.HTML("""
        <div style="text-align:center; opacity:0.5; margin-top:16px; font-size:0.85em;">
            Powered by fine-tuned Gemma 4 · Trained on CubiCasa5K + ArchCAD
        </div>
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)