import gradio as gr
import joblib
import json
import math
import re
from deep_translator import GoogleTranslator
import warnings

warnings.filterwarnings("ignore")

# Load models
from xgboost import Booster
xgb = Booster()
xgb.load_model("xgb_model.json")

rf = joblib.load("rf_model.pkl")  # Needs scikit-learn

# Load tile catalog and sizes
with open("tile_catalog.json", "r", encoding="utf-8") as f:
    tile_catalog = json.load(f)

with open("tile_sizes.json", "r", encoding="utf-8") as f:
    tile_sizes = json.load(f)

tile_size_map = {s["label"].lower().replace(" ", ""): s["area_sqft"] for s in tile_sizes}

def translate(text, lang="en"):
    try:
        return GoogleTranslator(source="auto", target=lang).translate(text)
    except:
        return text

def extract_tile_area(msg, unit):
    msg = msg.lower().replace("×", "x").replace("into", "x").replace("*", "x")
    msg = msg.replace("mm", "").replace("ft", "").replace("feet", "").strip()
    
    if "x" in msg:
        parts = re.split(r"x", msg)
        if len(parts) == 2:
            try:
                val1 = float(re.sub(r"[^\d.]", "", parts[0]))
                val2 = float(re.sub(r"[^\d.]", "", parts[1]))
                if unit == "mm":
                    sqft = (val1 * val2) / 92903.04
                else:
                    sqft = val1 * val2
                return round(sqft, 2)
            except:
                return None
    else:
        # Single value fallback for ft
        try:
            val = float(re.sub(r"[^\d.]", "", msg))
            return val if unit == "ft" else None
        except:
            return None

def chat_fn(message, history, user_state={}):
    message = message.strip().lower()

    if "lang" not in user_state:
        try:
            user_state["lang"] = GoogleTranslator(source="auto", target="en").detect(message)
        except:
            user_state["lang"] = "en"
    lang = user_state["lang"]
    def reply(text): return translate(text, lang)

    # Greeting handler
    if message in ["hi", "hello", "hey", "start", "help"]:
        user_state.clear()
        return reply("👋 Hello! Are you planning for Floor or Wall tiles?"), None, user_state

    if message in ["floor", "wall"]:
        user_state.clear()
        user_state["tile_type"] = message.capitalize()
        user_state["step"] = "get_length"
        return reply(f"Great! You chose {user_state['tile_type']} tiles.\nPlease enter the **length** of the space in feet:"), None, user_state

    # Stepwise logic
    if user_state.get("step") == "get_length":
        try:
            user_state["length"] = float(re.sub(r"[^\d.]", "", message))
            user_state["step"] = "get_width"
            return reply("Now enter the **width** of the space in feet:"), None, user_state
        except:
            return reply("Please enter a valid number for length (e.g. 12.5):"), None, user_state

    if user_state.get("step") == "get_width":
        try:
            user_state["width"] = float(re.sub(r"[^\d.]", "", message))
            user_state["area"] = round(user_state["length"] * user_state["width"], 2)
            user_state["step"] = "get_unit"
            return reply(f"Your total area is {user_state['area']} sq.ft.\nWould you like to enter tile size in **mm** or **ft**?"), None, user_state
        except:
            return reply("Please enter a valid number for width (e.g. 10):"), None, user_state

    if user_state.get("step") == "get_unit":
        if message not in ["mm", "ft"]:
            return reply("Please type either **mm** or **ft** to choose your preferred unit."), None, user_state
        user_state["unit"] = message
        user_state["step"] = "get_tile_size"
        return reply(f"Enter the tile size in {message.upper()} (e.g. 600x600 or 2x2):"), None, user_state

    if user_state.get("step") == "get_tile_size":
        area_per_tile = extract_tile_area(message, user_state["unit"])
        if not area_per_tile:
            return reply("I couldn't understand that size. Try something like `600x600`, `2x2`, or just `4` if square in ft."), None, user_state

        user_state["tile_area"] = area_per_tile
        user_state["step"] = "done"

        area_needed = user_state["area"]
        tile_type = user_state["tile_type"]
        tile_needed = math.ceil((area_needed / area_per_tile) * 1.1)

        # Suggest tiles
        best = []
        for tile in tile_catalog:
            if tile["type"].lower() == tile_type.lower():
                per_box = tile["coverage"] / area_per_tile
                if per_box > 0:
                    boxes = math.ceil(tile_needed / per_box)
                    total = boxes * tile["price"]
                    best.append({
                        "name": tile["name"],
                        "size": tile["size"],
                        "price": tile["price"],
                        "boxes": boxes,
                        "total": total,
                        "url": tile["url"]
                    })

        best.sort(key=lambda x: x["total"])
        top3 = best[:3]

        summary = f"""🧱 Tile Type: {tile_type}
📏 Area: {area_needed} sq.ft
🧮 Tile Size Area: {round(area_per_tile, 2)} sq.ft
🔢 Tiles Needed: {tile_needed} (10% extra included)

🎯 Top Recommendations:
"""
        for i, t in enumerate(top3, 1):
            summary += f"{i}. {t['name']} ({t['size']}) – ₹{t['price']} per box, ~{t['boxes']} boxes\n   🔗 {t['url']}\n"

        summary += "\nType 'Floor' or 'Wall' to start over."

        user_state["summary"] = summary
        return reply(summary), None, user_state

    return reply("👋 Please type 'Floor' or 'Wall' to begin tile estimation."), None, user_state

# UI setup
with gr.Blocks() as demo:
    gr.ChatInterface(
        fn=chat_fn,
        title="🧱 TileBot – Smart Tile Estimator",
        description=(
            "Plan your tile needs easily. Just type `hi` to begin or start with `Floor` or `Wall`.\n"
            "I’ll ask for room size, tile size, and recommend the best fit tiles."
        ),
        theme=gr.themes.Soft(),
        retry_btn=None,
        undo_btn=None
    )

demo.launch()