app/utils/climate_intelligence.py

import logging
import re
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

try:
    import openmeteo_requests
    import pandas as pd
    import requests_cache
    from retry_requests import retry
except Exception:  # pragma: no cover - optional dependency fallback for offline use
    openmeteo_requests = None
    pd = None
    requests_cache = None
    retry = None


NIGERIA_LOCATION_COORDS: Dict[str, Tuple[float, float]] = {
    "lagos": (6.5244, 3.3792),
    "ikeja": (6.596, 3.342),
    "abuja": (9.0765, 7.3986),
    "kaduna": (10.5222, 7.4383),
    "kano": (12.0022, 8.5920),
    "ibadan": (7.3775, 3.9470),
    "port harcourt": (4.8156, 7.0498),
    "enugu": (6.5244, 7.4954),
    "jos": (9.8965, 8.8583),
    "owerri": (5.4833, 7.0353),
    "katsina": (12.9894, 7.6000),
    "bauchi": (10.3158, 9.8433),
    "gombe": (10.2897, 11.1670),
    "sokoto": (13.0627, 5.2430),
    "minna": (9.6150, 6.5469),
    "akure": (7.2508, 5.1950),
    "ilorin": (8.4966, 4.5421),
    "oyo": (7.8500, 3.9300),
    "niger": (9.6500, 6.3500),
    "kogi": (7.8000, 6.7400),
    "benue": (7.7300, 8.5200),
    "nasarawa": (8.5333, 8.5200),
    "plateau": (9.9200, 8.9000),
    "rivers": (4.8156, 7.0498),
    "ebonyi": (6.2600, 8.1300),
    "anambra": (6.2100, 7.0700),
    "bayelsa": (4.9240, 6.0890),
    "delta": (5.8940, 5.8620),
    "edo": (6.5000, 5.7500),
    "kwara": (8.5000, 4.5500),
    "ogun": (7.0000, 3.3500),
    "osun": (7.7700, 4.5600),
    "ondo": (7.1000, 4.8400),
    "taraba": (7.8700, 9.7800),
    "zamfara": (12.1700, 6.6600),
    "kebbi": (12.4530, 4.1970),
    "yobe": (11.7333, 11.0833),
    "jigawa": (12.2280, 9.9960),
    "adamawa": (9.3260, 12.3950),
    "akwa ibom": (5.0300, 7.9200),
    "cross river": (5.8800, 8.3400),
}

CROP_REQUIREMENTS: Dict[str, Dict[str, str]] = {
    "maize": {
        "ideal": "Moderate rainfall during establishment; avoid prolonged dry spells before tasselling.",
        "warning": "Poor emergence and yield loss if rainfall is delayed or dry spells persist.",
    },
    "rice": {
        "ideal": "Saturated soils or consistent moisture during early growth and tillering.",
        "warning": "Waterlogging and nutrient losses can occur during intense rainfall events.",
    },
    "tomato": {
        "ideal": "Regular moisture, but avoid persistent wet conditions that encourage fungal disease.",
        "warning": "Heat stress and excess humidity raise disease and blossom-drop risk.",
    },
    "cassava": {
        "ideal": "Well-distributed rainfall and adequate soil moisture during establishment.",
        "warning": "Severe drought and heat can suppress rooting and tuber formation.",
    },
    "groundnut": {
        "ideal": "Moisture at germination and flowering; avoid prolonged waterlogging.",
        "warning": "Dry spells during pegging and pod filling reduce yields.",
    },
}


CROP_KEYWORDS: Dict[str, List[str]] = {
    "maize": ["maize", "corn"],
    "rice": ["rice", "paddy"],
    "tomato": ["tomato", "tomatoes"],
    "cassava": ["cassava", "manioc"],
    "groundnut": ["groundnut", "peanut", "peanuts"],
}

GROWTH_STAGE_PATTERNS: Dict[str, List[str]] = {
    "pre-planting": ["before planting", "pre planting", "pre-planting", "planting soon", "ready to plant", "sowing soon"],
    "planting": ["planting", "sowing", "seedbed", "direct sowing"],
    "early-growth": ["seedling", "early growth", "germination", "seedling stage"],
    "flowering": ["flowering", "tasselling", "blooming"],
    "fruiting": ["fruiting", "grain filling", "pod filling", "harvest soon"],
}


def extract_nigeria_location(query: str) -> Tuple[str, Optional[float], Optional[float]]:
    if not query:
        return "", None, None

    q = query.strip().lower()
    ordered_locations = sorted(NIGERIA_LOCATION_COORDS.keys(), key=lambda item: len(item), reverse=True)
    for location in ordered_locations:
        pattern = rf"\b{re.escape(location)}\b"
        if re.search(pattern, q):
            lat, lon = NIGERIA_LOCATION_COORDS[location]
            return location.title(), lat, lon

    return "", None, None


def infer_crop_and_growth_stage(query: str) -> Tuple[Optional[str], Optional[str]]:
    if not query:
        return None, None

    q = query.strip().lower()
    crop = None
    for key, aliases in CROP_KEYWORDS.items():
        if any(alias in q for alias in aliases):
            crop = key
            break

    stage = None
    for key, patterns in GROWTH_STAGE_PATTERNS.items():
        if any(pattern in q for pattern in patterns):
            stage = key
            break

    if crop is None and "maize" in q:
        crop = "maize"
    if stage is None and any(term in q for term in [
        "before planting",
        "pre planting",
        "plant soon",
        "ready to sow",
        "before the rains",
        "before rains",
        "before rain",
    ]):
        stage = "pre-planting"
    if stage is None and any(term in q for term in ["planting", "sowing", "plant maize", "plant rice", "plant cassava"]):
        stage = "planting"
    if stage is None and "should i plant" in q:
        stage = "pre-planting"

    return crop, stage


def _normalize_crop(crop: Optional[str]) -> str:
    if crop is None:
        return "general"
    key = crop.strip().lower().replace(" ", "")
    if key in CROP_REQUIREMENTS:
        return key
    return "general"


def _safe_float(value: Any, default: float = 0.0) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def _get_openmeteo_client():
    if openmeteo_requests is None:
        return None
    try:
        cache_session = requests_cache.CachedSession(".cache", expire_after=3600)
        retry_session = retry(cache_session, retries=5, backoff_factor=0.2)
        return openmeteo_requests.Client(session=retry_session)
    except Exception:
        return None


def _forecast_from_openmeteo(latitude: float, longitude: float, days: int = 7) -> Dict[str, Any]:
    client = _get_openmeteo_client()
    if client is None:
        raise RuntimeError("Open-Meteo client is unavailable")

    url = "https://api.open-meteo.com/v1/forecast"
    params = {
        "latitude": latitude,
        "longitude": longitude,
        "current": [
            "temperature_2m",
            "relative_humidity_2m",
            "apparent_temperature",
            "precipitation",
            "rain",
            "weather_code",
            "wind_speed_10m",
        ],
        "daily": [
            "weather_code",
            "temperature_2m_max",
            "temperature_2m_min",
            "precipitation_sum",
            "rain_sum",
            "daylight_duration",
        ],
        "forecast_days": max(1, min(int(days), 10)),
        "timezone": "auto",
    }

    responses = client.weather_api(url, params=params)
    response = responses[0]

    curr = response.Current()
    daily = response.Daily()

    current = {
        "temperature_c": _safe_float(curr.Variables(0).Value(), 0.0),
        "apparent_temperature_c": _safe_float(curr.Variables(2).Value(), 0.0),
        "relative_humidity_percent": _safe_float(curr.Variables(1).Value(), 0.0),
        "precipitation_mm": _safe_float(curr.Variables(3).Value(), 0.0),
        "rain_mm": _safe_float(curr.Variables(4).Value(), 0.0),
        "wind_speed_kph": _safe_float(curr.Variables(6).Value(), 0.0),
        "weather_code": int(curr.Variables(5).Value()),
    }

    daily_dates = pd.date_range(
        start=pd.to_datetime(daily.Time(), unit="s", utc=True),
        end=pd.to_datetime(daily.TimeEnd(), unit="s", utc=True),
        freq=pd.Timedelta(seconds=daily.Interval()),
        inclusive="left",
    ) if pd is not None else []

    max_temp = daily.Variables(1).ValuesAsNumpy() if daily.Variables(1) else []
    min_temp = daily.Variables(2).ValuesAsNumpy() if daily.Variables(2) else []
    precip_sum = daily.Variables(3).ValuesAsNumpy() if daily.Variables(3) else []
    rain_sum = daily.Variables(4).ValuesAsNumpy() if daily.Variables(4) else []

    forecast_days = []
    for index, date_value in enumerate(daily_dates):
        forecast_days.append(
            {
                "date": date_value.strftime("%Y-%m-%d"),
                "temperature_max_c": _safe_float(max_temp[index], 0.0),
                "temperature_min_c": _safe_float(min_temp[index], 0.0),
                "precipitation_mm": _safe_float(precip_sum[index], 0.0),
                "rain_mm": _safe_float(rain_sum[index], 0.0),
            }
        )

    return {
        "location": {
            "latitude": latitude,
            "longitude": longitude,
            "country": "Nigeria",
            "source": "Open-Meteo forecast",
        },
        "current_weather": current,
        "forecast": forecast_days,
    }


def _historical_from_openmeteo(latitude: float, longitude: float, days: int = 30) -> Dict[str, Any]:
    client = _get_openmeteo_client()
    if client is None:
        raise RuntimeError("Open-Meteo client is unavailable")

    end_date = datetime.utcnow().date()
    start_date = end_date - timedelta(days=max(1, int(days)))
    url = "https://archive-api.open-meteo.com/v1/archive"
    params = {
        "latitude": latitude,
        "longitude": longitude,
        "start_date": start_date.isoformat(),
        "end_date": end_date.isoformat(),
        "daily": ["temperature_2m_max", "temperature_2m_min", "precipitation_sum"],
        "timezone": "auto",
    }

    responses = client.weather_api(url, params=params)
    response = responses[0]
    daily = response.Daily()

    if pd is None:
        return {"recent_rainfall_mm": 0.0, "average_max_temp_c": 0.0, "average_min_temp_c": 0.0}

    dates = pd.date_range(
        start=pd.to_datetime(daily.Time(), unit="s", utc=True),
        end=pd.to_datetime(daily.TimeEnd(), unit="s", utc=True),
        freq=pd.Timedelta(seconds=daily.Interval()),
        inclusive="left",
    )
    max_temp = daily.Variables(0).ValuesAsNumpy()
    min_temp = daily.Variables(1).ValuesAsNumpy()
    precip = daily.Variables(2).ValuesAsNumpy()

    values = []
    for idx, date_value in enumerate(dates):
        values.append(
            {
                "date": date_value.strftime("%Y-%m-%d"),
                "temperature_max_c": _safe_float(max_temp[idx], 0.0),
                "temperature_min_c": _safe_float(min_temp[idx], 0.0),
                "precipitation_mm": _safe_float(precip[idx], 0.0),
            }
        )

    total_rain = sum(item["precipitation_mm"] for item in values)
    avg_max = sum(item["temperature_max_c"] for item in values) / max(len(values), 1)
    avg_min = sum(item["temperature_min_c"] for item in values) / max(len(values), 1)

    return {
        "period_days": len(values),
        "recent_rainfall_mm": round(total_rain, 1),
        "average_max_temp_c": round(avg_max, 1),
        "average_min_temp_c": round(avg_min, 1),
        "daily_history": values,
    }


def _fallback_weather(latitude: float, longitude: float, days: int = 7) -> Dict[str, Any]:
    current = {
        "temperature_c": 29.0,
        "apparent_temperature_c": 30.0,
        "relative_humidity_percent": 72.0,
        "precipitation_mm": 0.8,
        "rain_mm": 0.8,
        "wind_speed_kph": 12.0,
        "weather_code": 1,
    }
    forecast = [
        {
            "date": (datetime.now(timezone.utc) + timedelta(days=i)).strftime("%Y-%m-%d"),
            "temperature_max_c": 31.0 + i * 0.25,
            "temperature_min_c": 23.0 + i * 0.15,
            "precipitation_mm": 8.0 if i % 2 == 0 else 3.0,
            "rain_mm": 8.0 if i % 2 == 0 else 3.0,
        }
        for i in range(max(1, int(days)))
    ]
    return {
        "location": {
            "latitude": latitude,
            "longitude": longitude,
            "country": "Nigeria",
            "source": "Fallback climate heuristic",
        },
        "current_weather": current,
        "forecast": forecast,
    }


def _assess_climate_risk(
    crop: Optional[str],
    growth_stage: Optional[str],
    current_weather: Dict[str, Any],
    forecast: List[Dict[str, Any]],
    history: Dict[str, Any],
) -> Tuple[str, int, str]:
    risk_score = 0
    reasons: List[str] = []

    temp = _safe_float(current_weather.get("temperature_c"), 28.0)
    humidity = _safe_float(current_weather.get("relative_humidity_percent"), 70.0)
    rain_7d = sum(_safe_float(item.get("precipitation_mm"), 0.0) for item in forecast[:7])
    recent_rain = _safe_float(history.get("recent_rainfall_mm"), 0.0)

    if recent_rain < 40:
        risk_score += 2
        reasons.append("Recent rainfall remains below the moisture needed for establishment and early growth.")
    if temp >= 35:
        risk_score += 2
        reasons.append("High temperature signals heat stress risk for crops under active growth.")
    if humidity >= 80 and crop in {"tomato", "maize"}:
        risk_score += 1
        reasons.append("High humidity raises disease pressure risk in humid conditions.")
    if rain_7d < 20 and growth_stage in {"pre-planting", "early-growth", "seedling", "planting"}:
        risk_score += 2
        reasons.append("Projected rainfall remains too low to support reliable planting or establishment.")
    if any(_safe_float(item.get("precipitation_mm"), 0.0) > 35 for item in forecast[:3]):
        risk_score += 1
        reasons.append("Short intense rainfall events may raise runoff and waterlogging concerns.")

    if crop in {"maize", "rice", "groundnut"} and growth_stage in {"pre-planting", "planting"}:
        risk_score += 1

    if risk_score <= 2:
        level = "low"
    elif risk_score <= 5:
        level = "medium"
    else:
        level = "high"

    summary = " | ".join(reasons) if reasons else "No major immediate climate anomaly detected under the available forecast and history."
    return level, risk_score, summary


def _recommendations_for_crop(crop: Optional[str], growth_stage: Optional[str], risk_level: str) -> List[str]:
    normalized_crop = _normalize_crop(crop)
    stage = (growth_stage or "general").strip().lower()

    if normalized_crop == "maize":
        if stage in {"pre-planting", "planting"}:
            return [
                "Delay planting until the soil has received enough cumulative rainfall for germination.",
                "Use conservation tillage or mulch to hold soil moisture and reduce evaporation.",
                "Consider a shorter-season or drought-tolerant variety if the onset remains late.",
            ]
        return [
            "Monitor moisture stress during early vegetative growth.",
            "Apply mulch and avoid late irrigation to preserve soil moisture.",
            "Use field scouting for leaf curl, fungal disease, and nutrient stress.",
        ]

    if normalized_crop == "rice":
        return [
            "Keep flooded or saturated paddy conditions stable during early establishment.",
            "If rainfall is intense, improve drainage to avoid standing water and root stress.",
            "Check bund integrity and ensure a reliable water control plan is in place.",
        ]

    if normalized_crop == "tomato":
        return [
            "Protect plants from heat stress by irrigating during early morning and mulching around roots.",
            "Improve spacing and airflow to reduce disease pressure under humid conditions.",
            "Avoid excess foliage wetness and monitor for blight after rainfall events.",
        ]

    if risk_level == "high":
        return [
            "Prioritize soil moisture conservation and immediate field checks.",
            "Delay non-urgent field operations until conditions are more stable.",
            "Consult a local extension officer if there is persistent drought or flood stress.",
        ]

    return [
        "Use the current conditions to guide irrigation, planting, and crop protection timing.",
        "Keep a short field checklist for soil moisture, leaf stress, and pest activity.",
        "Plan the next field decision around the next 5–7 days of weather risk.",
    ]


def build_climate_context_for_query(query: str, days: int = 7) -> str:
    """Build climate advisory context for a user question, using location/crop/stage extraction."""
    if not query:
        return ""

    location_name, latitude, longitude = extract_nigeria_location(query)
    if latitude is None or longitude is None:
        return ""

    crop, growth_stage = infer_crop_and_growth_stage(query)
    climate_context = build_climate_intelligence_context(
        latitude=latitude,
        longitude=longitude,
        crop=crop,
        growth_stage=growth_stage,
        days=days,
    )

    return (
        f"Location: {location_name or 'Nigeria'}\n"
        f"Crop: {climate_context.get('crop', crop or 'not specified')}\n"
        f"Growth stage: {climate_context.get('growth_stage', growth_stage or 'not specified')}\n"
        f"Climate risk: {climate_context.get('climate_risk', 'unknown')}\n"
        f"Summary: {climate_context.get('risk_reason', 'No major risk summary available')}\n"
        f"Recommended actions: {'; '.join(climate_context.get('recommendations', []))}"
    )


def build_climate_intelligence_context(
    latitude: float,
    longitude: float,
    crop: Optional[str] = None,
    growth_stage: Optional[str] = None,
    days: int = 7,
) -> Dict[str, Any]:
    """
    Build a structured climate-risk context for use in the farmer advisory prompt.
    This layer combines current weather, short-term forecast, recent historical climate,
    crop agronomic requirements, and a simple risk assessment before the LLM is asked to give advice.
    """
    try:
        climate = _forecast_from_openmeteo(latitude, longitude, days=days)
    except Exception as exc:
        logger.warning("Open-Meteo forecast failed; using fallback climate heuristic: %s", exc)
        climate = _fallback_weather(latitude, longitude, days=days)

    try:
        historical = _historical_from_openmeteo(latitude, longitude, days=30)
    except Exception as exc:
        logger.warning("Open-Meteo historical climate failed; using fallback history: %s", exc)
        historical = {
            "period_days": 30,
            "recent_rainfall_mm": 42.0,
            "average_max_temp_c": 30.5,
            "average_min_temp_c": 23.5,
            "daily_history": [],
        }

    crop_name = crop or "general"
    crop_key = _normalize_crop(crop_name)
    crop_requirements = CROP_REQUIREMENTS.get(crop_key, {
        "ideal": "Use locally recommended practices that match the crop and the season.",
        "warning": "Climate variability can affect emergence, growth rate, and yield.",
    })

    risk_level, risk_score, risk_reason = _assess_climate_risk(
        crop=crop_name,
        growth_stage=growth_stage,
        current_weather=climate.get("current_weather", {}),
        forecast=climate.get("forecast", []),
        history=historical,
    )

    recommendation_list = _recommendations_for_crop(crop_name, growth_stage, risk_level)

    summary = {
        "location": {
            "latitude": latitude,
            "longitude": longitude,
            "country": "Nigeria",
            "source": climate.get("location", {}).get("source", "Open-Meteo"),
        },
        "crop": crop_name,
        "growth_stage": growth_stage or "not specified",
        "current_weather": climate.get("current_weather", {}),
        "forecast": climate.get("forecast", [])[:days],
        "historical_climate": {
            "period_days": historical.get("period_days", 30),
            "recent_rainfall_mm": historical.get("recent_rainfall_mm", 0.0),
            "average_max_temp_c": historical.get("average_max_temp_c", 0.0),
            "average_min_temp_c": historical.get("average_min_temp_c", 0.0),
        },
        "crop_requirements": crop_requirements,
        "climate_risk": risk_level,
        "risk_score": risk_score,
        "risk_reason": risk_reason,
        "recommendations": recommendation_list,
    }

    return summary


def build_climate_prompt_context(
    latitude: float,
    longitude: float,
    crop: Optional[str] = None,
    growth_stage: Optional[str] = None,
    days: int = 7,
) -> str:
    context = build_climate_intelligence_context(latitude, longitude, crop=crop, growth_stage=growth_stage, days=days)

    current = context["current_weather"]
    forecast = context["forecast"]
    hist = context["historical_climate"]
    risk_reason = context["risk_reason"]
    recommendations = "\n- ".join(context["recommendations"])

    forecast_block = []
    for item in forecast[:5]:
        forecast_block.append(
            f"{item.get('date')}: max {item.get('temperature_max_c')}°C, min {item.get('temperature_min_c')}°C, rain {item.get('precipitation_mm')} mm"
        )

    summary = (
        "CLIMATE INTELLIGENCE CONTEXT\n"
        f"Location: Nigeria (lat={latitude}, lon={longitude})\n"
        f"Crop: {context['crop']}\n"
        f"Growth stage: {context['growth_stage']}\n"
        "Current conditions:\n"
        f"- Temperature: {current.get('temperature_c')}°C\n"
        f"- Humidity: {current.get('relative_humidity_percent')}%\n"
        f"- Wind: {current.get('wind_speed_kph')} kph\n"
        f"- Rain: {current.get('rain_mm')} mm\n"
        "Recent climate history:\n"
        f"- 30-day rainfall: {hist.get('recent_rainfall_mm')} mm\n"
        f"- Average max temp: {hist.get('average_max_temp_c')}°C\n"
        f"- Average min temp: {hist.get('average_min_temp_c')}°C\n"
        "Short-term forecast:\n"
        + ("\n".join(f"- {line}" for line in forecast_block) if forecast_block else "- Forecast unavailable")
        + "\n"
        + f"Climate risk: {context['climate_risk']}\n"
        + f"Risk rationale: {risk_reason}\n"
        + "Crop agronomic requirement:\n"
        + f"- {context['crop_requirements'].get('ideal', 'Use crop-specific agronomic guidance')}\n"
        + "Recommended actions:\n"
        + f"- {recommendations}\n"
    )

    return summary