import json
import os
import pickle
from pathlib import Path

import gradio as gr
import numpy as np
import pandas as pd
import anthropic

MODEL_PATH = Path("random_forest_regression.pkl")
DATA_PATH = Path("bfs_municipality_and_tax_data.csv")

# Hugging Face Secrets / local environment variables
# Required secret on Hugging Face: ANTHROPIC_API_KEY
LLM_API_KEY = os.getenv("ANTHROPIC_API_KEY", "")
LLM_MODEL = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5")


def load_model_and_data():
    if not MODEL_PATH.exists():
        raise FileNotFoundError(
            f"Missing model file: {MODEL_PATH}. Upload random_forest_regression.pkl to the same folder as app.py."
        )
    if not DATA_PATH.exists():
        raise FileNotFoundError(
            f"Missing data file: {DATA_PATH}. Upload bfs_municipality_and_tax_data.csv to the same folder as app.py."
        )

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

    bfs_data = pd.read_csv(DATA_PATH, sep=",", encoding="utf-8")
    bfs_data["tax_income"] = (
        bfs_data["tax_income"].astype(str).str.replace("'", "", regex=False).astype(float)
    )
    bfs_data["emp"] = bfs_data["emp"].fillna(bfs_data["emp"].median())

    return loaded_model, bfs_data


model, df_bfs_data = load_model_and_data()

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


def match_town(user_town: str):
    """Match user town text to the canonical bfs_name in the CSV."""
    if not user_town:
        return None

    query = str(user_town).strip().lower()
    if not query:
        return None

    # 1) Exact lower-case match
    if query in town_to_row:
        return town_to_row[query]["bfs_name"]

    # 2) Relaxed contains matching
    matches = []
    for town in valid_towns:
        town_lower = str(town).strip().lower()
        if query == town_lower:
            return town
        if query in town_lower or town_lower in query:
            matches.append(town)

    if len(matches) == 1:
        return matches[0]

    return None


def call_llm_json(system_prompt: str, user_prompt: str) -> str:
    """Call Anthropic Claude and require a JSON object response."""
    if not LLM_API_KEY:
        raise RuntimeError(
            "ANTHROPIC_API_KEY fehlt. Bitte unter Hugging Face Settings -> Variables and secrets hinzufügen."
        )

    client = anthropic.Anthropic(api_key=LLM_API_KEY)

    message = client.messages.create(
        model=LLM_MODEL,
        max_tokens=512,
        system=system_prompt,
        messages=[
            {"role": "user", "content": user_prompt},
        ],
    )

    return message.content[0].text


def parse_json_response(raw: str, required_keys: tuple) -> dict:
    cleaned = (raw or "").strip()

    # Remove markdown code fences if present
    if cleaned.startswith("```"):
        cleaned = cleaned.split("```")[1]
        if cleaned.startswith("json"):
            cleaned = cleaned[4:]
        cleaned = cleaned.strip()

    if not cleaned:
        raise ValueError("LLM hat eine leere Antwort zurückgegeben.")

    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise ValueError(
            f"LLM hat kein gültiges JSON zurückgegeben. Erhalten: {cleaned[:300]}"
        ) from exc

    missing_keys = [key for key in required_keys if key not in parsed]
    if missing_keys:
        raise ValueError(
            f"JSON fehlt folgende Schlüssel: {', '.join(missing_keys)}."
        )

    return parsed


def extract_preferences(user_text: str) -> dict:
    """Extract rooms, area_m2 and town from German free text using an LLM."""
    if not user_text or not user_text.strip():
        raise ValueError(
            "Bitte gib einen Wohnungswunsch mit Zimmeranzahl, Fläche und Ort ein."
        )

    system_prompt = """Du bist ein Extraktionssystem für Schweizer Wohnungswünsche.

Extrahiere aus dem Nutzereingabetext exakt diese Informationen:
- rooms: Anzahl Zimmer als Zahl, zum Beispiel 3.5
- area_m2: Wohnfläche in Quadratmetern als Zahl, zum Beispiel 85
- town: Schweizer Gemeinde / Ort als String, zum Beispiel "Winterthur"

Antworte ausschliesslich mit gültigem JSON.
Keine Markdown-Codeblöcke.
Keine Erklärungen.
Keine zusätzlichen Schlüssel.
Wenn ein Wert fehlt, verwende null.

Erwartetes Format:
{"rooms": 3.5, "area_m2": 85, "town": "Winterthur"}"""

    user_prompt = f"""Extrahiere rooms, area_m2 und town aus folgendem deutschen Wohnungswunsch:

{user_text}"""

    raw = call_llm_json(system_prompt, user_prompt)
    parsed = parse_json_response(raw, required_keys=("rooms", "area_m2", "town"))

    rooms = parsed["rooms"]
    area_m2 = parsed["area_m2"]
    town = parsed["town"]

    if rooms is None or area_m2 is None or town is None:
        raise ValueError("Bitte gib Zimmeranzahl, Wohnfläche in m2 und Ort an.")

    try:
        rooms = float(rooms)
        area_m2 = float(area_m2)
    except (TypeError, ValueError) as exc:
        raise ValueError(
            "Zimmeranzahl und Wohnfläche müssen als Zahlen erkannt werden."
        ) from exc

    if rooms <= 0:
        raise ValueError("Die Zimmeranzahl muss grösser als 0 sein.")
    if area_m2 <= 0:
        raise ValueError("Die Wohnfläche muss grösser als 0 sein.")

    canonical_town = match_town(town)
    if canonical_town is None:
        raise ValueError(
            f"Ort '{town}' wurde nicht eindeutig in den BFS-Daten gefunden. "
            "Bitte verwende einen gültigen Schweizer Gemeindenamen."
        )

    return {
        "rooms": rooms,
        "area_m2": area_m2,
        "town": canonical_town,
    }


def predict_apartment_price(rooms: float, area_m2: float, town: str) -> float:
    """Predict monthly rent with exactly these model features:
    [rooms, area_m2, pop, pop_dens, frg_pct, emp, tax_income]
    """
    canonical_town = match_town(town)
    if canonical_town is None:
        raise ValueError(f"Ort '{town}' wurde nicht gefunden.")

    row = town_to_row[canonical_town.lower()]

    features = np.array(
        [[
            float(rooms),
            float(area_m2),
            float(row["pop"]),
            float(row["pop_dens"]),
            float(row["frg_pct"]),
            float(row["emp"]),
            float(row["tax_income"]),
        ]]
    )

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


def generate_explanation(preferences: dict, prediction: float) -> str:
    """Generate a concise German explanation with one uncertainty note using an LLM."""
    system_prompt = """Du bist ein hilfreicher Assistent für Wohnungsmietpreise in der Schweiz.

Deine Aufgabe:
Erkläre die gegebene Modellschätzung kurz und verständlich auf Deutsch.

Wichtige Regeln:
- Verwende die angegebene Prediction.
- Berechne keinen neuen Preis.
- Erfinde keine zusätzlichen Daten.
- Erwähne Zimmeranzahl, Fläche und Ort.
- Füge einen kurzen Unsicherheitshinweis hinzu.
- Sage, dass es sich um eine Schätzung handelt.
- Antworte ausschliesslich mit gültigem JSON ohne Markdown.
- Verwende exakt den Schlüssel "answer".

Erwartetes Format:
{"answer": "Für eine 3.5-Zimmer-Wohnung mit 85 m2 in Winterthur schätzt das Modell die Monatsmiete auf rund 2145 CHF. Die Schätzung ist unsicher, weil Faktoren wie Zustand, Mikrolage und Ausstattung nicht direkt berücksichtigt werden."}"""

    user_prompt = f"""Wohnungsdaten:
{json.dumps(preferences, ensure_ascii=False)}

Modellschätzung der monatlichen Miete in CHF:
{prediction}

Formuliere eine kurze deutsche Antwort."""

    raw = call_llm_json(system_prompt, user_prompt)
    parsed = parse_json_response(raw, required_keys=("answer",))

    answer = parsed["answer"]
    if not answer:
        raise ValueError("LLM Erklärung war leer.")

    return str(answer)


def run_pipeline(user_text: str):
    """End-to-end app pipeline."""
    try:
        preferences = extract_preferences(user_text)
        prediction = predict_apartment_price(
            rooms=preferences["rooms"],
            area_m2=preferences["area_m2"],
            town=preferences["town"],
        )
        final_answer = generate_explanation(
            preferences=preferences,
            prediction=prediction,
        )
        return preferences, prediction, final_answer

    except Exception as exc:
        return (
            {
                "error": str(exc),
                "hint": "Bitte gib Zimmeranzahl, Wohnfläche in m2 und einen Schweizer Ort an.",
            },
            None,
            f"Fehler: {exc}",
        )


with gr.Blocks(title="🏠 Schweizer Wohnungspreisschätzer") as demo:
    gr.Markdown(
        """
        # 🏠 Schweizer Wohnungspreisschätzer

        Beschreibe deinen **Wohnungswunsch auf Deutsch**.

        **Beispiele:**
        - *„Ich suche eine 3.5-Zimmer-Wohnung mit etwa 85 m2 in Winterthur."*
        - *„Wie viel kostet eine 2-Zimmer-Wohnung mit 55 m2 in Zürich?"*
        - *„Ich brauche eine 4.5-Zimmer-Wohnung mit 110 m2 in Bern."*
        """
    )

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

    submit = gr.Button("🔍 Schätzen", variant="primary")

    with gr.Row():
        extracted = gr.JSON(label="📋 Extrahierte Eingaben")
        price = gr.Number(label="💰 Geschätzte Monatsmiete (CHF)")

    response = gr.Textbox(label="💬 Erklärung", lines=5)

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

    gr.Markdown(
        """
        ---
        *Hinweis: Diese Schätzung basiert auf einem Random-Forest-Modell mit BFS-Gemeindedaten.
        Sie dient nur zur Orientierung und ersetzt keine professionelle Mietberatung.*
        """
    )


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