import json
import math
import os
from pathlib import Path

import pickle

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

MODEL_PATH = Path("model.pkl")

# Support both variable naming schemes used in notebook/app and Hugging Face settings.
LLM_API_KEY = os.getenv("LLM_API_KEY") or os.getenv("OPENAI_API_KEY") or ""
LLM_MODEL = os.getenv("LLM_MODEL") or os.getenv("OPENAI_MODEL") or ""

if not MODEL_PATH.exists():
    fallback_model = Path("model.pkl")
    if fallback_model.exists():
        MODEL_PATH = fallback_model

with open(MODEL_PATH, "rb") as model_file:
    model_package = pickle.load(model_file)

if isinstance(model_package, dict):
    model = model_package["model"]
    scaler = model_package.get("scaler", None)
    model_features = model_package.get("features", None)
else:
    model = model_package
    scaler = None
    model_features = None

df_bfs_data = pd.read_csv("bfs_municipality_and_tax_data.csv", sep=",", encoding="utf-8")
df_bfs_data["tax_income"] = (
    df_bfs_data["tax_income"].astype(str).str.replace("'", "", regex=False).astype(float)
)

town_to_row = {
    str(row["bfs_name"]).lower(): row
    for _, row in df_bfs_data.iterrows()
}
valid_towns = list(df_bfs_data["bfs_name"].sort_values().unique())


# Core Pipeline Functions

def match_town(user_town: str):
    """Return the canonical town name from the dataset, or None."""
    if not user_town or not user_town.strip():
        return None
    
    user_town_lower = str(user_town).strip().lower()
    
    # Exact lower-case match
    if user_town_lower in town_to_row:
        return town_to_row[user_town_lower]["bfs_name"]
    
    # Relaxed contains-match over valid_towns
    for town in valid_towns:
        if user_town_lower in str(town).lower():
            return town
    
    return None


def create_features(rooms, area, town):
    """Create the full feature row expected by the saved model."""
    town_lower = str(town).lower()
    if town_lower not in town_to_row:
        raise ValueError(f"Town '{town}' not found in dataset.")

    town_data = town_to_row[town_lower]
    area = float(area)
    rooms = float(rooms)
    pop = float(town_data["pop"])
    emp = float(town_data["emp"])
    tax_income = float(town_data["tax_income"])

    features = {
        "rooms": rooms,
        "area": area,
        "pop": pop,
        "pop_dens": float(town_data["pop_dens"]),
        "frg_pct": float(town_data["frg_pct"]),
        "emp": emp,
        "tax_income": tax_income,
        "rooms_per_sqm": rooms / area if area else 0.0,
        "wealth_index": (tax_income / 100000.0) * (emp / 100000.0),
        "is_zurich_city": 1 if town_lower == "zürich" or town_lower == "zurich" else 0,
        "pop_emp_ratio": pop / (emp + 1.0),
        "log_area": math.log1p(area),
        "log_pop": math.log1p(pop),
        "log_tax_income": math.log1p(tax_income),
    }

    return pd.DataFrame([features])


def call_llm_json(system_prompt: str, user_prompt: str) -> str:
    """Call LLM with system and user prompts, return JSON response text."""
    if not LLM_API_KEY or not LLM_MODEL:
        raise ValueError("LLM_API_KEY and LLM_MODEL environment variables are required.")
    
    client = OpenAI(api_key=LLM_API_KEY)

    response = client.chat.completions.create(
        model=LLM_MODEL,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        max_tokens=500,
    )

    return (response.choices[0].message.content or "").strip()


# Validate the LLM response before the rest of the app depends on it.
# Why this helps:
# - LLMs sometimes return empty text, Markdown, or incomplete JSON.
# - Early validation makes the app more stable and easier to debug.
# - This is a strong general design habit: check external input before using it.
def parse_json_response(raw: str, required_keys: tuple[str, ...]) -> dict:
    cleaned = (raw or "").strip()

    if not cleaned:
        raise ValueError("LLM returned an empty response instead of JSON.")

    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise ValueError(
            f"LLM did not return valid JSON. Received: {cleaned[:300]}"
        ) from exc

    missing_keys = [key for key in required_keys if key not in parsed]
    if missing_keys:
        raise ValueError(
            f"LLM JSON is missing required keys: {', '.join(missing_keys)}."
        )

    return parsed


def extract_preferences(user_text: str) -> dict:
    """Extract rooms, area_m2, and town from free text using LLM."""
    system_prompt = """Du bist ein Assistent, der Wohnungswünsche in strukturierte Daten umwandelt.
Extrahiere aus der Benutzereingabe die drei Parameter:
- rooms: Anzahl der Zimmer (als Dezimalzahl, z.B. 3.5)
- area_m2: Wohnfläche in Quadratmetern (als Zahl)
- town: Name der Stadt oder Gemeinde

Antworte ausschließlich mit gültigem JSON in diesem Format:
{"rooms": <number>, "area_m2": <number>, "town": "<string>"}"""
    
    raw_response = call_llm_json(system_prompt, user_text)
    parsed = parse_json_response(raw_response, ("rooms", "area_m2", "town"))
    
    matched_town = match_town(parsed["town"])
    if not matched_town:
        raise ValueError(f"Town '{parsed['town']}' not found in dataset.")
    
    parsed["town"] = matched_town
    return parsed


def predict_apartment_price(rooms: float, area_m2: float, town: str) -> float:
    """Predict monthly rent using the loaded model."""
    feature_frame = create_features(rooms, area_m2, town)

    if model_features is not None:
        missing = [name for name in model_features if name not in feature_frame.columns]
        if missing:
            raise ValueError(f"Missing model features: {', '.join(missing)}")
        feature_frame = feature_frame[model_features]

    if scaler is not None:
        feature_input = scaler.transform(feature_frame)
    else:
        feature_input = feature_frame

    prediction = model.predict(feature_input)[0]
    return round(prediction, 2)


def generate_explanation(preferences: dict, prediction: float) -> str:
    """Generate a user-friendly explanation using LLM."""
    system_prompt = "Du bist ein hilfsbereiter Assistent für Immobilienvorhersagen."
    
    prompt = f"""Erkläre das folgende Miet-Schätzungsresultat in einfachen deutschen Worten.
    
Wohnungswunsch:
- Zimmer: {preferences.get('rooms', 'N/A')}
- Fläche: {preferences.get('area_m2', 'N/A')} m²
- Ort: {preferences.get('town', 'N/A')}

Geschätzte Monatsmiete: {prediction} CHF

Gib eine kurze Erklärung (1-2 Sätze) auf Deutsch und eine Unsicherheitsnote.
Antworte ausschließlich mit gültigem JSON in diesem Format:
{{"answer": "<deine Erklärung>"}}"""
    
    raw_response = call_llm_json(system_prompt, prompt)
    parsed = parse_json_response(raw_response, ("answer",))
    
    return parsed["answer"]


def run_pipeline(user_text: str):
    """End-to-end pipeline: extract -> predict -> explain."""
    try:
        preferences = extract_preferences(user_text)
        prediction = predict_apartment_price(
            preferences["rooms"],
            preferences["area_m2"],
            preferences["town"]
        )
        explanation = generate_explanation(preferences, prediction)
        return (preferences, prediction, explanation)
    except Exception as e:
        raise Exception(f"Pipeline error: {str(e)}") from e


with gr.Blocks(title="Apartment Wishes -> Prediction") as demo:
    gr.Markdown(
        """
        # Apartment Predictor
        Beschreibe den Wohnungswunsch bitte auf Deutsch.
        Beispiel: "Ich suche eine 3.5-Zimmer-Wohnung mit etwa 85 m2 in Winterthur."
        """
    )

    user_text = gr.Textbox(
        label="Wohnungswunsch",
        lines=4,
        placeholder="Beschreibe Zimmer, Fläche in m2 und Ort auf Deutsch...",
    )
    submit = gr.Button("Schätzen")

    extracted = gr.JSON(label="Extrahierte Eingaben")
    price = gr.Number(label="Geschätzte Monatsmiete (CHF)")
    response = gr.Textbox(label="Antwort", lines=6)

    submit.click(
        fn=run_pipeline,
        inputs=[user_text],
        outputs=[extracted, price, response],
    )

demo.launch()