import json
import os
import pickle
import re

import gradio as gr
import numpy as np
import pandas as pd
from openai import OpenAI

# ──────────────────────────────────────────────────────────────
# Load prediction model once at startup
# ──────────────────────────────────────────────────────────────
with open("apartment_rf_model.pkl", "rb") as f:
    _model_data = pickle.load(f)

_MODEL = _model_data["model"]
_FEATURES = _model_data["features"]

# Location defaults (city of Zurich vs. greater canton)
_POP_DENS_CITY = 4729.0
_POP_DENS_OUTSIDE = 1328.0
_FRG_PCT = 25.0
_LAT_CITY, _LON_CITY = 47.380402, 8.530496
_LAT_OUTSIDE, _LON_OUTSIDE = 47.424900, 8.638663


# ──────────────────────────────────────────────────────────────
# Prediction logic (mirrors the existing Gradio form app)
# ──────────────────────────────────────────────────────────────
def predict_price(params: dict) -> int:
    rooms = float(params.get("rooms", 3.0))
    area = float(params.get("area", 80))
    tax_income = float(params.get("tax_income", 80000))
    luxurious = bool(params.get("luxurious", False))
    furnished = bool(params.get("furnished", False))
    temporary = bool(params.get("temporary", False))
    zurich_city = bool(params.get("zurich_city", True))
    attika = bool(params.get("attika", False))
    loft = bool(params.get("loft", False))
    seesicht = bool(params.get("seesicht", False))
    kreis = str(params.get("kreis", "None"))

    pop_dens = _POP_DENS_CITY if zurich_city else _POP_DENS_OUTSIDE
    lat = _LAT_CITY if zurich_city else _LAT_OUTSIDE
    lon = _LON_CITY if zurich_city else _LON_OUTSIDE

    log_area = np.log1p(area)
    log_rooms = np.log1p(rooms)
    rooms_area_ratio = rooms / (area + 1)
    room_per_m2 = area / rooms if rooms > 0 else area
    income_density_score = tax_income * pop_dens / 1e6
    log_income_dens = np.log1p(income_density_score)
    log_pop_dens = np.log1p(pop_dens)
    area_rooms_interact = area * rooms

    kreis_cols = {f"Kreis {i}": 0 for i in range(1, 13)}
    if kreis in kreis_cols:
        kreis_cols[kreis] = 1

    row = {
        "rooms": rooms,
        "area": area,
        "log_area": log_area,
        "log_rooms": log_rooms,
        "pop_dens": pop_dens,
        "log_pop_dens": log_pop_dens,
        "frg_pct": _FRG_PCT,
        "tax_income": tax_income,
        "income_density_score": income_density_score,
        "log_income_dens": log_income_dens,
        "rooms_area_ratio": rooms_area_ratio,
        "area_rooms_interact": area_rooms_interact,
        "lat": lat,
        "lon": lon,
        "luxurious": int(luxurious),
        "furnished": int(furnished),
        "temporary": int(temporary),
        "zurich_city": int(zurich_city),
        "room_per_m2": room_per_m2,
        **kreis_cols,
        "(ATTIKA)": int(attika),
        "(LOFT)": int(loft),
        "(SEESICHT)": int(seesicht),
        "(LUXURIÖS)": 0,
        "(POOL)": 0,
        "(EXKLUSIV)": 0,
    }

    df = pd.DataFrame([row])[_FEATURES]
    pred_log = _MODEL.predict(df)[0]
    return round(np.expm1(pred_log))


# ──────────────────────────────────────────────────────────────
# JSON extraction helper
# ──────────────────────────────────────────────────────────────
def extract_json(text: str) -> dict | None:
    match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group(1))
        except json.JSONDecodeError:
            return None
    return None


# ──────────────────────────────────────────────────────────────
# System prompt for the extraction step
# ──────────────────────────────────────────────────────────────
_EXTRACTION_SYSTEM = """\
You are a helpful real estate assistant specialising in Zurich, Switzerland apartment rentals.
Your task: understand the user's apartment description and extract structured parameters.

## When you have BOTH rooms AND area from the user

Output a JSON block with exactly these keys:
```json
{
  "rooms": <float — number of rooms, e.g. 3.5>,
  "area": <float — living area in m², e.g. 80>,
  "tax_income": <float — municipal median tax income CHF, default 80000 if unknown>,
  "luxurious": <bool — true if described as luxury/high-end>,
  "furnished": <bool — true if furnished>,
  "temporary": <bool — true if temporary/short-term>,
  "zurich_city": <bool — true = within City of Zurich, false = canton but outside city>,
  "attika": <bool — true if penthouse / attika apartment>,
  "loft": <bool — true if loft>,
  "seesicht": <bool — true if lake view>,
  "kreis": <string — "Kreis 1" … "Kreis 12" or "None">
}
```
After the JSON block, add one short sentence listing any assumptions you made.

## When rooms OR area are missing

Ask the user conversationally for the missing detail. Do NOT output a JSON block.

## Parameter guidelines
- Half-rooms count (e.g. 3.5-Zimmer → rooms: 3.5)
- kreis is only valid for zurich_city=true; use "None" otherwise
- All boolean flags default to false unless the user mentions them
- Respond in the same language the user writes in
"""

_EXPLANATION_SYSTEM = "You are a helpful real estate assistant for Zurich, Switzerland. Be concise and friendly."


# ──────────────────────────────────────────────────────────────
# Main chat function
# ──────────────────────────────────────────────────────────────
def respond(message: str, history: list[dict]) -> str:
    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    # Build full message list for the extraction call
    msgs = [{"role": "system", "content": _EXTRACTION_SYSTEM}] + list(history) + [{"role": "user", "content": message}]

    # Step 1 — extract parameters (or ask for missing info)
    r1 = client.chat.completions.create(
        model="gpt-4o-mini",
        max_tokens=1024,
        messages=msgs,
    )
    step1_text = r1.choices[0].message.content
    params = extract_json(step1_text)

    if params is None:
        # Model is asking a clarifying question — return as-is
        return step1_text

    # Step 2 — run the prediction model
    try:
        price = predict_price(params)
    except Exception as exc:
        return step1_text + f"\n\n⚠️ Prediction error: {exc}"

    # Step 3 — ask GPT to write a friendly natural-language explanation
    explanation_prompt = (
        f"The prediction model returned **CHF {price:,} / month** for the apartment "
        "described by the JSON above.\n\n"
        "Write a short, friendly reply (3–5 sentences) that:\n"
        "1. States the predicted rent clearly and prominently\n"
        "2. Briefly contextualises it (location, size, market context)\n"
        "3. Notes it is a model estimate and actual prices may vary\n"
        "4. Invites the user to adjust parameters or ask follow-up questions\n\n"
        "Do NOT repeat the JSON — it is shown separately above your text."
    )

    r2 = client.chat.completions.create(
        model="gpt-4o-mini",
        max_tokens=512,
        messages=[
            {"role": "system", "content": _EXPLANATION_SYSTEM},
            *msgs[1:],  # history without the extraction system prompt
            {"role": "assistant", "content": step1_text},
            {"role": "user", "content": explanation_prompt},
        ],
    )
    explanation = r2.choices[0].message.content

    # Compose the full assistant reply: JSON + prediction + explanation
    return (
        "**Extracted Parameters**\n"
        f"```json\n{json.dumps(params, indent=2)}\n```\n\n"
        f"**Predicted Monthly Rent: CHF {price:,}**\n\n"
        f"{explanation}"
    )


# ──────────────────────────────────────────────────────────────
# Gradio UI
# ──────────────────────────────────────────────────────────────
demo = gr.ChatInterface(
    fn=respond,
    type="messages",
    title="Zurich Apartment Rent Assistant",
    description=(
        "Describe the apartment you're interested in and I'll estimate the monthly rent. "
        "Just tell me about the number of rooms, size, location, and any special features — "
        "no forms needed!"
    ),
    examples=[
        "Ich suche eine 3-Zimmer-Wohnung, ca. 80 m², im Kreis 3.",
        "Was würde eine möblierte 1.5-Zimmer-Wohnung in der Stadt Zürich, ca. 40 m², kosten?",
        "Ich möchte ein 4.5-Zimmer-Penthouse mit Seeblick in Kreis 8, ca. 150 m². Es ist luxuriös.",
    ],
    theme=gr.themes.Soft(),
)

if __name__ == "__main__":
    demo.launch()