
# weather_capability_v7_full.py
from __future__ import annotations
import os, re, json, requests
import datetime as _dt
from typing import Dict, Any, List, Optional, AsyncGenerator, Tuple

from tvi.solphit.base.logging import SolphitLogger

# Optional LLM (ingialla)
try:
    from tvi.solphit.ingialla.ask import Generator
    _HAS_LLM = True
except Exception:
    _HAS_LLM = False

log = SolphitLogger.get_logger("discera.capabilities.weather")

# --------------------------- HTTP helpers ---------------------------
def _http_json(url: str, *, timeout: float = 20.0) -> Dict[str, Any]:
    r = requests.get(url, timeout=timeout)
    r.raise_for_status()
    return r.json() or {}

# --------------------------- Ollama fallback ---------------------------
def _ollama_base_url() -> str:
    return os.getenv("OLLAMA_BASE_URL", "http://localhost:11434").rstrip("/")

def _ollama_generate(*, prompt: str, model: str = "llama3.2:3b") -> str:
    url = f"{_ollama_base_url()}/api/generate"
    body = {"model": model, "prompt": prompt, "stream": False}
    r = requests.post(url, json=body, timeout=60)
    r.raise_for_status()
    data = r.json() or {}
    text = str(data.get("response", "")).strip()
    if not text:
        msg = data.get("message") or data.get("error") or ""
        raise RuntimeError(f"Ollama /api/generate returned empty response. {msg}")
    return text

# --------------------------- WMO weather codes ---------------------------
WMO = {
    0: "Clear sky",
    1: "Mainly clear", 2: "Partly cloudy", 3: "Overcast",
    45: "Fog", 48: "Depositing rime fog",
    51: "Light drizzle", 53: "Moderate drizzle", 55: "Dense drizzle",
    56: "Light freezing drizzle", 57: "Dense freezing drizzle",
    61: "Slight rain", 63: "Moderate rain", 65: "Heavy rain",
    66: "Light freezing rain", 67: "Heavy freezing rain",
    71: "Slight snow", 73: "Moderate snow", 75: "Heavy snow",
    77: "Snow grains",
    80: "Slight rain showers", 81: "Moderate rain showers", 82: "Violent rain showers",
    85: "Slight snow showers", 86: "Heavy snow showers",
    95: "Thunderstorm", 96: "Thunderstorm with slight hail", 99: "Thunderstorm with heavy hail",
}

def _wmo_label(code: Optional[int]) -> str:
    try:
        return WMO.get(int(code), f"Weather code {code}")
    except Exception:
        return "Unknown"

# --------------------------- Query extraction ---------------------------
ISO_DATE_RE = r"\b(\d{4})-(\d{2})-(\d{2})\b"

_DEF_LOCATION = os.getenv("WEATHER_DEFAULT_LOCATION", "New York")
_DEF_DAYS = int(os.getenv("WEATHER_DEFAULT_DAYS", "3"))
_DEF_MODEL = os.getenv("WEATHER_EXTRACT_MODEL", "llama3.2:3b")

def _extract_json_block(text: str) -> str:
    """Find the first {...} JSON object in text; if none, return text."""
    m = re.search(r"\{[\s\S]*\}", text)
    return m.group(0) if m else text

def _normalize_query_obj(obj: Dict[str, Any]) -> Dict[str, Any]:
    location = str(obj.get("location") or "").strip() or _DEF_LOCATION
    days = int(obj.get("days") or _DEF_DAYS)
    days = max(1, min(days, 16))
    start = obj.get("start")
    if isinstance(start, str):
        start = start.strip() or None
    # allow "on" key for single-day queries (LLM may produce it)
    on = obj.get("on")
    if isinstance(on, str) and on.strip():
        start = on.strip()
        days = 1
    units = str(obj.get("units") or "metric").strip().lower()
    units = "imperial" if units.startswith("imp") or "f" in units else "metric"
    return {"location": location, "days": days, "start": start, "units": units}

def _extract_query_with_regex(q: str) -> Dict[str, Any]:
    q = (q or "").strip()
    # Location after 'in' or 'for'
    loc = None
    m = re.search(r"\b(?:in|for)\s+([A-Za-z0-9\-\.,\s]+)", q)
    if m:
        loc = m.group(1).strip().rstrip(".")
    # Date: ISO or common formats like "January 15, 1942"
    start = None
    m = re.search(ISO_DATE_RE, q)
    if m:
        start = f"{m.group(1)}-{m.group(2)}-{m.group(3)}"
    else:
        m2 = re.search(r"\b([A-Za-z]+)\s+(\d{1,2}),\s*(\d{4})\b", q)
        if m2:
            try:
                dt = _dt.datetime.strptime(m2.group(0), "%B %d, %Y").date()
                start = dt.strftime("%Y-%m-%d")
            except Exception:
                start = None
        else:
            m3 = re.search(r"\b(\d{1,2})\/(\d{1,2})\/(\d{4})\b", q)  # mm/dd/yyyy
            if m3:
                try:
                    dt = _dt.datetime.strptime(m3.group(0), "%m/%d/%Y").date()
                    start = dt.strftime("%Y-%m-%d")
                except Exception:
                    start = None
    # Days
    days = None
    m = re.search(r"next\s+(\d{1,2})\s+days", q, flags=re.I)
    if m:
        days = int(m.group(1))
    else:
        m = re.search(r"(\d{1,2})\s*day(s)?", q, flags=re.I)
        if m:
            days = int(m.group(1))
    # Single-day phrasing
    if re.search(r"\bon\b\s+", q, flags=re.I) and start:
        days = 1
    # Units
    units = "metric"
    if re.search(r"\b(fahrenheit|imperial|°f|deg f|degrees f)\b", q, flags=re.I):
        units = "imperial"
    # Defaults
    location = loc or _DEF_LOCATION
    d = max(1, min(int(days or _DEF_DAYS), 16))
    return {"location": location, "days": d, "start": start, "units": units}

def _extract_query_with_llm(q: str, *, provider: str, model: str) -> Dict[str, Any]:
    """Ask LLM to produce strict JSON: {location, days, start|null, units}."""
    system = (
        "You extract weather query parameters from the user's text. "
        "Return strict JSON with keys: location (string), days (int), start (string or null), units ('metric'|'imperial'). "
        "If the user specifies a single calendar day (e.g., 'on 1942-01-15'), set days=1 and start=<that date>. "
        "No extra keys; return one JSON object only."
    )
    prompt = f"{system}\n\nUSER:\n{q}"
    # Try ingialla
    if _HAS_LLM:
        try:
            gen = Generator(provider=provider, model=model, verbose=False)
            txt = gen.generate(prompt, contexts=[], history=[]) or ""
            obj = json.loads(_extract_json_block(txt))
            return _normalize_query_obj(obj)
        except Exception as e:
            log.warning(f"[weather] LLM extract failed; falling back to regex: {e}")
    # Fallback to Ollama
    try:
        txt = _ollama_generate(prompt=prompt, model=model)
        obj = json.loads(_extract_json_block(txt))
        return _normalize_query_obj(obj)
    except Exception as e:
        log.warning(f"[weather] Ollama extract failed; falling back to regex: {e}")
    return _extract_query_with_regex(q)

# --------------------------- Date classification ---------------------------
def _parse_iso_date(s: str) -> Optional[_dt.date]:
    try:
        return _dt.datetime.strptime(s, "%Y-%m-%d").date()
    except Exception:
        return None

def _is_current_query(q: str) -> bool:
    ql = (q or "").lower()
    return any(k in ql for k in ["right now", "currently", "now", "today"])

def _choose_weather_endpoint(start: Optional[str]) -> str:
    """Return 'forecast' for now/recent, 'archive' for deep historical."""
    if not start:
        return "forecast"
    d = _parse_iso_date(start)
    if not d:
        return "forecast"
    delta = (_dt.date.today() - d).days
    return "archive" if delta > 10 else "forecast"

# --------------------------- Units & formatting ---------------------------
def _fmt_units(units: str) -> Tuple[str, str]:
    # temperature unit, precip unit
    if units == "imperial":
        return ("°F", "in")
    return ("°C", "mm")

def _mm_to_in(mm: float) -> float:
    return round(mm / 25.4, 2)

def _fmt_day_line(date: str, tmin_c: float, tmax_c: float, precip_mm: float, code: int, *, units: str) -> str:
    label = _wmo_label(code)
    t_unit, p_unit = _fmt_units(units)
    if units == "imperial":
        tmin = round(tmin_c * 9 / 5 + 32, 1)
        tmax = round(tmax_c * 9 / 5 + 32, 1)
        precip = _mm_to_in(precip_mm)
    else:
        tmin, tmax, precip = round(tmin_c, 1), round(tmax_c, 1), round(precip_mm, 1)
    return f"- **{date}**: {label}; low **{tmin}{t_unit}**, high **{tmax}{t_unit}**; precip **{precip}{p_unit}**."

def _chunk_text(s: str, max_len: int = 96) -> List[str]:
    out: List[str] = []
    buf = (s or "").strip()
    while buf:
        out.append(buf[:max_len])
        buf = buf[max_len:]
    return out

# --------------------------- Capability ---------------------------
class WeatherCapability:
    name = "weather"

    @staticmethod
    def descriptor() -> Dict[str, Any]:
        return {
            "name": "weather",
            "description": (
                "Answers weather questions with up-to-date multi-day forecasts and current conditions; "
                "supports human-friendly summaries, units (°C/°F, mm/in), and historical dates via archive."
            ),
            "intents": ["weather_forecast", "current_weather", "weather_trend"],
            "triggers": [
                "weather", "forecast", "current", "conditions", "temperature",
                "right now", "now", "today", "ferndale", "mi", "like", "in", "what", "the"
            ],
            "tags": ["weather", "forecast", "trend", "current", "location", "temperature", "precipitation"],
            "excludes": ["climate", "koppen"],
            "examples": [
                "What's the weather in Ferndale, MI right now?",
                "Weather in Detroit on January 15, 1942 (imperial units).",
                "Rain outlook for Amsterdam this week."
            ],
        }

    @staticmethod
    def assets() -> Dict[str, List[str]]:
        css = r"""
.discera .weather-summary { margin-top: 6px; }
.discera .weather-summary .lead { color: #9ca3af; }
"""
        return {"css": [css], "js": []}

    def run_once(self, args: Any, **kwargs) -> Dict[str, Any]:
        q: str = (getattr(args, "question", None) or (args.get("question") if isinstance(args, dict) else "") or "").strip()
        hints: Dict[str, Any] = kwargs.get("hints") or {}

        provider = hints.get("provider") or "ollama"
        model = hints.get("model") or _DEF_MODEL

        concise = False
        if isinstance(args, dict):
            concise = bool(args.get("concise", False))
        if not concise and "concise" in hints:
            concise = bool(hints.get("concise"))

        # 1) Extract structured query
        params = _extract_query_with_llm(q, provider=provider, model=model)
        location, days, start, units = params["location"], params["days"], params["start"], params["units"]
        single_day = bool(start) and days == 1

        # 2) Geocode
        try:
            geo = _http_json(
                f"https://geocoding-api.open-meteo.com/v1/search?name={requests.utils.quote(location)}&count=1"
            )
        except Exception as ex:
            answer = f"Sorry, I couldn't reach the geocoding service ({ex})."
            return {"answer": answer, "contexts": [], "meta": {"capability": self.name, "location": location}}

        if not geo.get("results"):
            answer = f"Sorry, I couldn't find the location **{location}**."
            return {"answer": answer, "contexts": ["https://open-meteo.com/en/docs"], "meta": {"capability": self.name, "location": location}}

        lat = float(geo["results"][0]["latitude"])
        lon = float(geo["results"][0]["longitude"])
        resolved_name = geo["results"][0].get("name") or location
        admin = geo["results"][0].get("admin1") or ""
        country = geo["results"][0].get("country") or ""
        header_loc = f"{resolved_name}" + (f", {admin}" if admin else "") + (f", {country}" if country else "")

        # 3a) Current conditions (lead) for "now" queries
        is_now = _is_current_query(q)
        lead = ""
        if is_now:
            current_url = (
                "https://api.open-meteo.com/v1/forecast"
                f"?latitude={lat}&longitude={lon}"
                "&current=temperature_2m,relative_humidity_2m,precipitation,weathercode,wind_speed_10m"
                "&timezone=auto"
            )
            try:
                current = _http_json(current_url).get("current", {}) or {}
            except Exception:
                current = {}
            if current:
                temp = current.get("temperature_2m")
                wind = current.get("wind_speed_10m")
                rh = current.get("relative_humidity_2m")
                wcode = current.get("weathercode")
                label = _wmo_label(wcode).lower()
                t_unit, _ = _fmt_units(units)
                if units == "imperial" and temp is not None:
                    temp = round(temp * 9/5 + 32, 1)
                if temp is not None and wind is not None and rh is not None:
                    lead = f"Now in {resolved_name}: {label}, {temp}{t_unit}, wind {wind} km/h, RH {rh}%."
                elif temp is not None:
                    lead = f"Now in {resolved_name}: {label}, {temp}{t_unit}."

        # 3b) Forecast or archive daily data
        endpoint = _choose_weather_endpoint(start)
        contexts: List[str] = []
        if endpoint == "forecast":
            base_url = (
                "https://api.open-meteo.com/v1/forecast"
                f"?latitude={lat}&longitude={lon}"
                "&daily=temperature_2m_max,temperature_2m_min,precipitation_sum,weathercode"
                f"&forecast_days={days}&timezone=auto"
            )
            if start:
                base_url += f"&start_date={start}"
        else:
            # archive API needs both start_date and end_date; single day -> same
            end = start or _dt.date.today().strftime("%Y-%m-%d")
            base_url = (
                "https://archive-api.open-meteo.com/v1/archive"
                f"?latitude={lat}&longitude={lon}"
                "&daily=temperature_2m_max,temperature_2m_min,precipitation_sum,weathercode"
                f"&start_date={start}&end_date={end}&timezone=auto"
            )
        try:
            weather = _http_json(base_url)
        except Exception as ex:
            answer = f"Sorry, I couldn't reach the weather service ({ex})."
            return {"answer": answer, "contexts": [base_url, "https://open-meteo.com/en/docs"], "meta": {"capability": self.name, "location": location}}

        daily = weather.get("daily", {})
        dates: List[str] = daily.get("time", []) or []
        tmax: List[float] = daily.get("temperature_2m_max", []) or []
        tmin: List[float] = daily.get("temperature_2m_min", []) or []
        precip: List[float] = daily.get("precipitation_sum", []) or []
        codes: List[int] = daily.get("weathercode", []) or []
        if not dates:
            answer = f"I couldn't find weather data for **{header_loc}**."
            return {"answer": answer, "contexts": [base_url, "https://open-meteo.com/en/docs"], "meta": {"capability": self.name, "location": header_loc}}

        # 4) Build output (concise vs narrative)
        t_unit, p_unit = _fmt_units(units)
        lines = [f"**{header_loc}** · Next **{min(days, len(dates))}** day(s) · Units **{t_unit}/{p_unit}**"]
        # For single-day historical queries, adjust header
        if single_day and start:
            lines[0] = f"**{header_loc}** · **{start}** · Units **{t_unit}/{p_unit}**"

        for i in range(min(days, len(dates))):
            lines.append(_fmt_day_line(dates[i], tmin[i], tmax[i], precip[i], codes[i], units=units))

        if not concise:
            # Weatherman-style narrative
            narrative: List[str] = []
            num_days = min(days, len(dates))
            intro_bits: List[str] = []
            if lead:
                intro_bits.append(lead)
            if single_day:
                intro_bits.append(f"Here’s the detailed weather for {header_loc} on {dates[0]}:")
            else:
                intro_bits.append(f"Here’s your detailed forecast for {header_loc} over the next {num_days} day(s):")

            for i in range(num_days):
                label = _wmo_label(codes[i])
                tminv, tmaxv, precipv = tmin[i], tmax[i], precip[i]
                if units == "imperial":
                    tminv = round(tminv * 9/5 + 32, 1)
                    tmaxv = round(tmaxv * 9/5 + 32, 1)
                    precipv = _mm_to_in(precipv)
                day_desc = (
                    f"On {dates[i]}, expect {label.lower()} with temperatures from "
                    f"{tminv}{t_unit} in the morning to {tmaxv}{t_unit} in the afternoon. "
                    f"Precipitation is around {precipv}{p_unit}."
                )
                if codes[i] in (95, 96, 99):
                    day_desc += " Thunderstorms are possible—stay weather-aware."
                elif codes[i] in (85, 86, 71, 73, 75, 77):
                    day_desc += " Snow is likely; plan travel accordingly."
                narrative.append(day_desc)

            answer = "\n\n".join([b for b in intro_bits if b.strip()]) + "\n\n" + "\n".join(narrative)
        else:
            answer = "\n".join(lines if not lead else [lead, *lines])

        meta = {
            "capability": self.name,
            "location": header_loc,
            "query": {"location": location, "days": days, "start": start, "units": units},
            "lat": lat, "lon": lon,
            "endpoint": endpoint,
            "forecast_days": min(days, len(dates)),
        }
        contexts = [base_url, "https://open-meteo.com/en/docs"]
        return {"answer": answer, "contexts": contexts, "meta": meta}

    async def run_stream(self, args: Any, **kwargs) -> AsyncGenerator[dict, None]:
        result = self.run_once(args, **kwargs)
        text = result.get("answer", "")
        yield {"event": "generation_started", "data": {"capability": self.name}}
        for chunk in _chunk_text(text, max_len=96):
            yield {"event": "token", "data": chunk}
        if (result.get("contexts") or []):
            yield {"event": "context", "data": {"items": result["contexts"]}}
        yield {"event": "generation_done", "data": {}}
        yield {"event": "done", "data": {}}

def register(registry):
    registry.register(WeatherCapability())
