import pickle
from pathlib import Path

import pandas as pd
import streamlit as st


st.set_page_config(page_title="UrbanNest Rent Predictor", layout="wide")

ARTIFACT_DIR = Path("models")
INFERENCE_BUNDLE_PATH = ARTIFACT_DIR / "inference_bundle.pkl"
UNKNOWN_CATEGORY_TOKEN = "__UNK__"


@st.cache_resource
def load_inference_artifacts():
    if not INFERENCE_BUNDLE_PATH.exists():
        raise FileNotFoundError(
            "Missing models/inference_bundle.pkl. Run train.ipynb first."
        )

    with open(INFERENCE_BUNDLE_PATH, "rb") as f:
        bundle = pickle.load(f)

    required_keys = {"model", "label_encoders", "feature_columns", "feature_metadata", "model_metadata"}
    missing_keys = required_keys - set(bundle.keys())
    if missing_keys:
        raise KeyError(f"inference_bundle.pkl is missing keys: {sorted(missing_keys)}")

    return (
        bundle["model"],
        bundle["label_encoders"],
        bundle["feature_columns"],
        bundle["feature_metadata"],
        bundle["model_metadata"],
    )


def _build_numeric_input(feature_name, metadata):
    dtype = metadata["dtype"]
    min_value = metadata["min"]
    max_value = metadata["max"]
    default_value = metadata["mean"]

    if "int" in dtype:
        return st.number_input(
            label=feature_name,
            min_value=int(min_value),
            max_value=int(max_value),
            value=int(round(default_value)),
            step=1,
        )

    return st.number_input(
        label=feature_name,
        min_value=float(min_value),
        max_value=float(max_value),
        value=float(default_value),
        step=0.1,
        format="%.4f",
    )


def main():
    st.title("UrbanNest Analytics: Dynamic House Rent Prediction")
    st.write("Provide property details below to estimate monthly rent (INR).")

    try:
        model, label_encoders, feature_columns, feature_metadata, model_metadata = load_inference_artifacts()
    except Exception as exc:
        st.error(f"Failed to load model artifacts: {exc}")
        st.stop()

    st.caption(
        f"Best training method: {model_metadata['best_method']} | "
        f"CV MAE: {model_metadata['best_cv_mae']:.2f} | "
        f"Test MAE: {model_metadata['test_mae']:.2f}"
    )

    user_inputs = {}
    left_col, right_col = st.columns(2)
    column_index = 0

    for feature_name in feature_columns:
        target_column = left_col if column_index % 2 == 0 else right_col
        with target_column:
            if feature_name in label_encoders:
                classes = [str(item) for item in label_encoders[feature_name].classes_]
                if UNKNOWN_CATEGORY_TOKEN in classes:
                    display_options = [c for c in classes if c != UNKNOWN_CATEGORY_TOKEN]
                    display_options.append("Other / Unknown")
                    user_inputs[feature_name] = st.selectbox(feature_name, options=display_options)
                else:
                    user_inputs[feature_name] = st.selectbox(feature_name, options=classes)
            else:
                user_inputs[feature_name] = _build_numeric_input(feature_name, feature_metadata[feature_name])
        column_index += 1

    if st.button("Predict"):
        encoded_row = {}
        for feature_name in feature_columns:
            feature_value = user_inputs[feature_name]
            if feature_name in label_encoders:
                raw_value = str(feature_value)
                if raw_value == "Other / Unknown":
                    raw_value = UNKNOWN_CATEGORY_TOKEN
                encoded_row[feature_name] = int(label_encoders[feature_name].transform([raw_value])[0])
            else:
                dtype = feature_metadata[feature_name]["dtype"]
                if "int" in dtype:
                    encoded_row[feature_name] = int(feature_value)
                else:
                    encoded_row[feature_name] = float(feature_value)

        model_input = pd.DataFrame([encoded_row], columns=feature_columns)
        prediction = float(model.predict(model_input)[0])
        st.success(f"Predicted Monthly Rent: INR {prediction:,.2f}")


if __name__ == "__main__":
    main()