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

from geopy import distance
from geopy.geocoders import Nominatim
import srtm
import requests
import requests_cache
import openmeteo_requests
from retry_requests import retry
import plotly.graph_objects as go


# --- GLOBAL SETUP ---

elevation_data = srtm.get_data()

with open("weather_icons_custom.json", "r") as f:
	icons = json.load(f)

cache_session = requests_cache.CachedSession(".cache", expire_after=3600)
retry_session = retry(cache_session, retries=5, backoff_factor=0.2)
openmeteo = openmeteo_requests.Client(session=retry_session)

geolocator = Nominatim(user_agent="snow_finder")

OVERPASS_URL = "https://maps.mail.ru/osm/tools/overpass/api/interpreter"
ICON_URL = "https://raw.githubusercontent.com/basmilius/weather-icons/refs/heads/dev/production/fill/svg/"

DEFAULT_LAT, DEFAULT_LON = 49.6116, 6.1319
MAX_PEAKS = 100


# --- UTILS ---

def compute_bbox(lat, lon, dist_km):
	lat_delta = dist_km / 111.0
	lon_delta = dist_km / (111.0 * np.cos(np.radians(lat)))

	south = max(lat - lat_delta, -90)
	north = min(lat + lat_delta, 90)
	west = lon - lon_delta
	east = lon + lon_delta

	if west < -180:
		west += 360
	if east > 180:
		east -= 360

	return f"{south},{west},{north},{east}"


def get_elevation_from_srtm(lat, lon):
	if lat is None or lon is None:
		return None

	if -56 <= lat <= 60:
		try:
			alt = elevation_data.get_elevation(lat, lon)
			if alt is not None and alt > 0:
				return alt
		except Exception:
			pass

	return None


# --- PEAK FETCHING (REFACTORED LOGIC) ---

def get_peaks_from_overpass(lat, lon, dist_km, min_altitude_m):
	bbox = compute_bbox(lat, lon, dist_km)

	query = f"""
	[out:json];
	(
	  nwr[natural=peak]({bbox});
	  nwr[natural=hill]({bbox});
	);
	out body;
	"""

	try:
		r = requests.get(OVERPASS_URL, params={"data": query}, timeout=30)
		r.raise_for_status()
		data = r.json()
	except Exception as e:
		print(f"Error fetching peaks: {e}")
		return pd.DataFrame()

	rows = []

	for e in data.get("elements", []):
		lat_e, lon_e = e.get("lat"), e.get("lon")
		if lat_e is None or lon_e is None:
			continue

		tags = e.get("tags", {})
		alt = None

		ele = tags.get("ele")
		if ele and str(ele).replace(".", "").replace("-", "").isnumeric():
			alt = float(ele)

		if alt is None or alt <= 0:
			alt = get_elevation_from_srtm(lat_e, lon_e)

		if alt is None or alt < min_altitude_m:
			continue

		rows.append({
			"name": tags.get("name", "Unnamed Peak/Hill"),
			"latitude": lat_e,
			"longitude": lon_e,
			"altitude": int(round(alt, 0)),
		})

	if not rows:
		return pd.DataFrame()

	df = pd.DataFrame(rows)

	df["distance_m"] = df.apply(
		lambda r: distance.distance(
			(r["latitude"], r["longitude"]), (lat, lon)
		).m,
		axis=1
	)

	df = (
		df.sort_values("distance_m")
		  .head(MAX_PEAKS)
		  .reset_index(drop=True)
	)

	return df


# --- WEATHER FETCH ---

def get_weather_for_peaks_iteratively(df_peaks, min_snow_cm, max_results=20, max_requests=100):
	if df_peaks.empty:
		return pd.DataFrame()

	url = "https://api.open-meteo.com/v1/forecast"
	results, requests_made = [], 0

	for _, row in df_peaks.iterrows():
		if len(results) >= max_results or requests_made >= max_requests:
			break

		params = {
			"latitude": str(row["latitude"]),
			"longitude": str(row["longitude"]),
			"elevation": str(row["altitude"]),
			"hourly": "temperature_2m,is_day,weather_code,snow_depth",
			"forecast_days": "1",
			"timezone": "auto",
		}

		try:
			responses = openmeteo.weather_api(url, params=params)
			if not responses:
				continue

			hourly = responses[0].Hourly()
			if hourly is None:
				continue

			idx = 0

			snow_depth_cm = float(hourly.Variables(3).ValuesAsNumpy()[idx]) * 100

			if snow_depth_cm >= min_snow_cm:
				results.append({
					**row.to_dict(),
					"temp_c": float(hourly.Variables(0).ValuesAsNumpy()[idx]),
					"is_day": int(hourly.Variables(1).ValuesAsNumpy()[idx]),
					"weather_code": int(hourly.Variables(2).ValuesAsNumpy()[idx]),
					"snow_depth_m": snow_depth_cm / 100,
					"snow_depth_cm": int(round(snow_depth_cm, 0)),
				})

		except Exception as e:
			print(f"Error fetching weather for {row['name']}: {e}")

		requests_made += 1

	return pd.DataFrame(results)


# --- POST-PROCESSING ---

def format_weather_data(df):
	if df.empty:
		return df

	def icon_mapper(row):
		code = str(int(row["weather_code"]))
		tod = "day" if row["is_day"] == 1 else "night"
		info = icons.get(code, {}).get(tod, {})
		return (
			ICON_URL + info.get("icon", ""),
			info.get("description", "Unknown"),
			info.get("icon", "")
		)

	df[["weather_icon_url", "weather_desc", "weather_icon_name"]] = df.apply(
		icon_mapper, axis=1, result_type="expand"
	)

	df["distance_km"] = (df["distance_m"] / 1000).round(1)
	df["temp_c_str"] = df["temp_c"].round(0).astype(int).astype(str) + "°C"
	return df


def geocode_location(location_text):
	try:
		loc = geolocator.geocode(location_text, timeout=10)
		if loc:
			return loc.latitude, loc.longitude, f"Found: {loc.address}"
		return None, None, f"Location '{location_text}' not found."
	except Exception as e:
		return None, None, f"Geocoding error: {e}"


# --- CORE LOGIC ---

def find_snowy_peaks(min_snow_cm, radius_km, min_altitude_m, lat, lon):
	if lat is None or lon is None:
		fig = create_empty_map(DEFAULT_LAT, DEFAULT_LON)
		fig.update_layout(title_text="Enter valid coordinates.")
		return fig, "Please enter coordinates."

	if not (-90 <= lat <= 90 and -180 <= lon <= 180):
		fig = create_empty_map(DEFAULT_LAT, DEFAULT_LON)
		fig.update_layout(title_text="Invalid coordinates.")
		return fig, "Coordinates out of range."

	df_peaks = get_peaks_from_overpass(lat, lon, radius_km, min_altitude_m)
	if df_peaks.empty:
		fig = create_map_with_center(lat, lon)
		fig.update_layout(title_text=f"No peaks found within {radius_km} km.")
		return fig, f"No peaks found within {radius_km} km."

	df_weather = get_weather_for_peaks_iteratively(df_peaks, min_snow_cm)
	if df_weather.empty:
		fig = create_map_with_center(lat, lon)
		fig.update_layout(title_text=f"No snowy peaks ≥ {min_snow_cm} cm.")
		return fig, f"No peaks met the ≥ {min_snow_cm} cm snow requirement."

	df_final = format_weather_data(df_weather)
	fig = create_map_with_results(lat, lon, df_final)
	fig.update_layout(title_text=f"Found {len(df_final)} snowy peaks!")
	return fig, f"🎉 Showing {len(df_final)} snowy peaks with ≥ {min_snow_cm} cm of snow."


# --- MAP HELPERS ---

def create_empty_map(lat, lon):
	fig = go.Figure()
	fig.update_layout(
		map=dict(style="open-street-map", center={"lat": lat, "lon": lon}, zoom=8),
		margin={"r": 0, "t": 40, "l": 0, "b": 0},
		height=1024,
		width=1024,
	)
	return fig


def create_map_with_center(lat, lon):
	fig = go.Figure(
		go.Scattermap(
			lat=[lat],
			lon=[lon],
			mode="markers",
			marker=dict(size=24, color="white", opacity=0.8),
			hoverinfo="skip",
		)
	)
	fig.add_trace(
		go.Scattermap(
			lat=[lat],
			lon=[lon],
			mode="markers",
			marker=dict(size=12, color="red"),
			text=["Search Center"],
			hoverinfo="text",
		)
	)
	fig.update_layout(
		map=dict(style="open-street-map", center={"lat": lat, "lon": lon}, zoom=8),
		margin={"r": 0, "t": 40, "l": 0, "b": 0},
		height=1024,
		width=1024,
	)
	return fig


def create_map_with_results(lat, lon, df_final):
	fig = go.Figure()

	fig.add_trace(
		go.Scattermap(
			lat=df_final["latitude"],
			lon=df_final["longitude"],
			mode="markers",
			marker=dict(size=24, color="white", opacity=0.8),
			hoverinfo="skip",
		)
	)

	fig.add_trace(
		go.Scattermap(
			lat=df_final["latitude"],
			lon=df_final["longitude"],
			mode="markers",
			marker=dict(size=12, color="blue"),
			customdata=df_final[
				["name", "altitude", "distance_km", "snow_depth_cm",
				 "weather_desc", "temp_c_str"]
			],
			hovertemplate=(
				"<b>%{customdata[0]}</b><br>"
				"Altitude: %{customdata[1]} m<br>"
				"Distance: %{customdata[2]} km<br>"
				"<b>❄️ Snow: %{customdata[3]} cm</b><br>"
				"Weather: %{customdata[4]}<br>"
				"🌡 Temp: %{customdata[5]}<extra></extra>"
			),
		)
	)

	fig.add_trace(
		go.Scattermap(
			lat=[lat],
			lon=[lon],
			mode="markers",
			marker=dict(size=24, color="white", opacity=0.8),
			hoverinfo="skip",
		)
	)
	fig.add_trace(
		go.Scattermap(
			lat=[lat],
			lon=[lon],
			mode="markers",
			marker=dict(size=12, color="red"),
			text=["Search Center"],
			hoverinfo="text",
		)
	)

	fig.update_layout(
		map=dict(style="open-street-map", center={"lat": lat, "lon": lon}, zoom=9),
		margin={"r": 0, "t": 40, "l": 0, "b": 0},
		height=1024,
		width=1024,
		showlegend=False,
	)
	return fig


# --- GRADIO UI ---

with gr.Blocks(theme=gr.themes.Soft(), title="Snow Finder") as demo:
	gr.Markdown("# ☃️ Snow Finder for Families")
	gr.Markdown("Find nearby snowy peaks perfect for sledding and snowmen!")

	with gr.Row():
		with gr.Column(scale=1):
			location_search = gr.Textbox(label="Search Location")
			search_location_btn = gr.Button("🔍 Find Location")

			lat_input = gr.Number(value=DEFAULT_LAT, label="Latitude", precision=4)
			lon_input = gr.Number(value=DEFAULT_LON, label="Longitude", precision=4)
			snow_slider = gr.Radio(choices=[1, 2, 3, 4, 5, 6], value=1, label="Min Snow (cm)")
			radius_slider = gr.Radio(choices=[10, 20, 30, 40, 50, 60], value=30, label="Radius (km)")
			altitude_slider = gr.Radio(choices=[100, 200, 300, 400, 500, 600, 700, 800, 900, 1000], value=300, label="Min Altitude (m)")
			search_button = gr.Button("❄️ Find Snow!", variant="primary")
			status_output = gr.Textbox(lines=4, interactive=False)

		with gr.Column(scale=2):
			init_fig = create_map_with_center(DEFAULT_LAT, DEFAULT_LON)
			init_fig.update_layout(title_text="Luxembourg City – Click 'Find Snow!' to start")
			map_plot = gr.Plot(init_fig, label="Map")

	search_location_btn.click(
		fn=geocode_location,
		inputs=[location_search],
		outputs=[lat_input, lon_input, status_output],
	)

	search_button.click(
		fn=find_snowy_peaks,
		inputs=[snow_slider, radius_slider, altitude_slider, lat_input, lon_input],
		outputs=[map_plot, status_output],
	)

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