"""
ClimaIQ Gemma Layer — Bilingual Explanation Engine
Converts ClimaIQ model output into plain language for farmers and loan officers.

Inference modes:
  - cloud: Google AI Studio / GenAI API (GEMMA_API_KEY, model GEMMA_MODEL_CLOUD)
  - ollama: local Ollama HTTP API at OLLAMA_BASE_URL (default http://127.0.0.1:11434)
Author: Krishna Dahale
"""

import json
import os
import urllib.error
import urllib.request
from typing import Optional, Tuple

from google import genai

# ─── Client Setup ──────────────────────────────────────────────────────────────

def get_client(api_key: str = None):
    key = api_key or os.environ.get("GEMMA_API_KEY")
    if not key:
        raise ValueError("API key not found. Set GEMMA_API_KEY environment variable.")
    return genai.Client(api_key=key)


GEMMA_MODEL_CLOUD = "gemma-4-26b-a4b-it"
# Backwards compatibility
GEMMA_MODEL = GEMMA_MODEL_CLOUD

DEFAULT_OLLAMA_BASE = os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/")
DEFAULT_OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "gemma3:4b")
CROP_WATER_CLASS = {
    "Rice": "very high",
    "Sugarcane": "very high",
    "Cotton": "moderate",
    "Wheat": "low",
    "Millets": "very low",
}


def _ollama_installed_name_matches(installed: str, requested: str) -> bool:
    """True if Ollama tags name matches user model (e.g. gemma3:4b vs gemma3:4b:latest)."""
    installed = (installed or "").strip()
    requested = (requested or "").strip()
    if not requested:
        return True
    if installed == requested:
        return True
    if installed.startswith(requested + ":"):
        return True
    return False


def ollama_verify_connection(
    base_url: Optional[str] = None,
    model: Optional[str] = None,
    timeout: int = 12,
) -> Tuple[bool, str]:
    """
    GET /api/tags — confirms Ollama is reachable and the chosen model is installed.
    Returns (success, human-readable message). No LLM inference run.
    """
    base = (base_url or DEFAULT_OLLAMA_BASE).rstrip("/")
    want = (model or DEFAULT_OLLAMA_MODEL).strip()
    url = f"{base}/api/tags"
    try:
        with urllib.request.urlopen(url, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")[:300] if e.fp else ""
        return False, f"Ollama returned HTTP {e.code} at {url}. {body}"
    except Exception as e:
        return False, f"Cannot reach Ollama at {base}. Start the server (e.g. ollama serve), then try again. ({e})"

    models = data.get("models") or []
    names = [m.get("name", "") for m in models if m.get("name")]
    if not names:
        hint = f"Run: ollama pull {want}" if want else "Run: ollama pull <model>"
        return False, f"Ollama responded but no models are installed yet. {hint}"

    if want:
        hits = [n for n in names if _ollama_installed_name_matches(n, want)]
        if not hits:
            preview = ", ".join(names[:10])
            suffix = " …" if len(names) > 10 else ""
            return (
                False,
                f"Ollama is running, but model '{want}' is not installed. "
                f"You have: {preview}{suffix}. Run: ollama pull {want}",
            )
        return True, f"Connected locally at {base}. Model '{hits[0]}' is ready, so you can use offline narratives."

    return True, f"Connected at {base}. {len(names)} model(s) installed."


def ollama_generate(
    prompt: str,
    base_url: Optional[str] = None,
    model: Optional[str] = None,
    timeout: int = 300,
) -> str:
    """
    Call Ollama's /api/chat (non-streaming). Works offline once Ollama is running.
    """
    base = (base_url or DEFAULT_OLLAMA_BASE).rstrip("/")
    m = model or DEFAULT_OLLAMA_MODEL
    url = f"{base}/api/chat"
    payload = json.dumps(
        {
            "model": m,
            "messages": [{"role": "user", "content": prompt}],
            "stream": False,
        }
    ).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace") if e.fp else ""
        raise RuntimeError(f"Ollama HTTP {e.code}: {body[:500]}") from e
    except urllib.error.URLError as e:
        raise RuntimeError(
            f"Cannot reach Ollama at {base}. Is `ollama serve` running? ({e.reason})"
        ) from e
    msg = data.get("message") or {}
    text = (msg.get("content") or "").strip()
    if not text and data.get("error"):
        raise RuntimeError(str(data["error"]))
    return text


def generate_explanation_text(
    prompt: str,
    *,
    inference_mode: str = "cloud",
    cloud_client=None,
    ollama_base: Optional[str] = None,
    ollama_model: Optional[str] = None,
) -> str:
    """Single entry: cloud (Google GenAI) or local (Ollama)."""
    if inference_mode == "ollama":
        return ollama_generate(prompt, base_url=ollama_base, model=ollama_model)
    if cloud_client is None:
        raise ValueError("Cloud mode requires a Google GenAI client and GEMMA_API_KEY.")
    response = cloud_client.models.generate_content(
        model=GEMMA_MODEL_CLOUD,
        contents=prompt,
    )
    return response.text

# ─── Prompt Templates ──────────────────────────────────────────────────────────

def _build_farmer_prompt(farmer_data: dict, result: dict, language: str) -> str:
    drivers = "\n".join([f"- {d['display_name']}" for d in result["top_risk_drivers"]])
    crop = str(farmer_data["crop_type"])
    water_class = CROP_WATER_CLASS.get(crop, "moderate")

    if language == "Hindi":
        return f"""
आप एक सहानुभूतिपूर्ण कृषि ऋण सलाहकार हैं जो किसानों को सरल हिंदी में समझाते हैं।

किसान की जानकारी:
- फसल: {farmer_data['crop_type']}
- राज्य: {farmer_data['state']}
- वार्षिक आय: ₹{farmer_data['annual_income_lakhs']} लाख
- मांगा गया ऋण: ₹{farmer_data['loan_amount_lakhs']} लाख
- इस मौसम में बारिश की कमी: {abs(farmer_data['rainfall_deficit_pct'])}%
- सूखे की गंभीरता (SPI): {farmer_data['spi']}
- लगातार सूखे के वर्ष: {farmer_data['consecutive_drought_years']}

ClimaIQ परिणाम:
- क्रेडिट स्कोर: {result['credit_score']} (300–850 में से)
- जोखिम श्रेणी: {result['risk_band']}
- डिफ़ॉल्ट संभावना: {result['default_probability']}%
- मुख्य जोखिम कारण:
{drivers}

निम्नलिखित प्रारूप में 4-5 वाक्यों में उत्तर दें:
1. किसान को सरल भाषा में बताएं कि उनका स्कोर क्यों कम/अधिक है
2. मौसम और फसल का जोखिम पर क्या असर हुआ, यह समझाएं
3. किसान 3 व्यावहारिक कदम क्या उठा सकते हैं (बीमा, वैकल्पिक फसल, सिंचाई आदि)
महत्वपूर्ण तथ्य: इस इंजन में {crop} की पानी-आधारित श्रेणी "{water_class}" मानी जाती है। पानी-आवश्यकता के बारे में इससे उल्टा दावा न करें।
केवल हिंदी में उत्तर दें। जटिल शब्दों से बचें।
"""
    else:
        return f"""
You are a helpful agricultural credit advisor explaining a loan assessment to a farmer in simple English.

Farmer Profile:
- Crop: {farmer_data['crop_type']}
- State: {farmer_data['state']}
- Annual Income: ₹{farmer_data['annual_income_lakhs']} lakhs
- Loan Requested: ₹{farmer_data['loan_amount_lakhs']} lakhs
- Rainfall Deficit this season: {abs(farmer_data['rainfall_deficit_pct'])}%
- Drought Severity (SPI): {farmer_data['spi']}
- Consecutive Drought Years: {farmer_data['consecutive_drought_years']}

ClimaIQ Result:
- Credit Score: {result['credit_score']} out of 850
- Risk Band: {result['risk_band']}
- Default Probability: {result['default_probability']}%
- Key Risk Drivers:
{drivers}

In 4-5 simple sentences:
1. Explain why their score is what it is
2. Explain how the weather and crop type has affected their risk
3. Give 3 practical steps they can take (insurance, alternate crops, irrigation, etc.)
Important fact guardrail: in this engine, {crop} is classified as "{water_class}" water-intensity. Do not claim the opposite.
Use simple, non-technical language. Be empathetic and constructive.
"""


def _build_officer_prompt(farmer_data: dict, result: dict, language: str) -> str:
    drivers = "\n".join([f"- {d['display_name']}: {d['direction']}" for d in result["top_risk_drivers"]])

    if language == "Hindi":
        return f"""
आप एक वरिष्ठ कृषि ऋण अधिकारी हैं। नीचे दी गई जानकारी के आधार पर एक संक्षिप्त ऋण मूल्यांकन रिपोर्ट हिंदी में तैयार करें।

किसान प्रोफाइल:
- फसल: {farmer_data['crop_type']} | राज्य: {farmer_data['state']}
- आयु: {farmer_data['age']} वर्ष | भूमि: {farmer_data['land_size_acres']} एकड़
- वार्षिक आय: ₹{farmer_data['annual_income_lakhs']} लाख
- ऋण राशि: ₹{farmer_data['loan_amount_lakhs']} लाख
- पिछले डिफ़ॉल्ट: {'हाँ' if farmer_data['previous_defaults'] else 'नहीं'}
- बारिश की कमी: {abs(farmer_data['rainfall_deficit_pct'])}%
- SPI: {farmer_data['spi']} | लगातार सूखा: {farmer_data['consecutive_drought_years']} वर्ष

ClimaIQ स्कोर: {result['credit_score']} | जोखिम: {result['risk_band']} | डिफ़ॉल्ट संभावना: {result['default_probability']}%
सुझाई गई कार्रवाई: {result['recommended_action']}

मुख्य जोखिम कारण:
{drivers}

निम्नलिखित शीर्षक ठीक इसी रूप में अलग पंक्ति पर लिखें, फिर एक खाली पंक्ति, फिर अनुच्छेद:
1. ऋण मूल्यांकन सारांश
2. जलवायु जोखिम विश्लेषण
3. सुझाई गई कार्रवाई और शर्तें
4. निगरानी बिंदु
शीर्षकों पर मार्कडाउन बोल्ड न लगाएँ। प्रत्येक खंड में लगभग 2 वाक्य रखें।
"""
    else:
        return f"""
You are a senior agricultural loan officer. Prepare a concise loan assessment report based on the ClimaIQ analysis below.

Farmer Profile:
- Crop: {farmer_data['crop_type']} | State: {farmer_data['state']}
- Age: {farmer_data['age']} | Land: {farmer_data['land_size_acres']} acres
- Annual Income: ₹{farmer_data['annual_income_lakhs']} lakhs
- Loan Requested: ₹{farmer_data['loan_amount_lakhs']} lakhs
- Previous Defaults: {'Yes' if farmer_data['previous_defaults'] else 'No'}
- Rainfall Deficit: {abs(farmer_data['rainfall_deficit_pct'])}%
- SPI: {farmer_data['spi']} | Consecutive Drought Years: {farmer_data['consecutive_drought_years']}

ClimaIQ Score: {result['credit_score']} | Risk Band: {result['risk_band']} | Default Probability: {result['default_probability']}%
Recommended Action: {result['recommended_action']}

Top Risk Drivers:
{drivers}

Structure your report using these exact numbered headings, each on its own line, followed by a blank line, then the paragraph(s):
1. Assessment Summary
2. Climate Risk Analysis
3. Recommended Action with conditions
4. Monitoring flags if loan is approved
Do not use markdown bold for the numbers. Keep each heading text exactly as shown (you may add detail after the blank line only).
Be professional and specific. Reference the ClimaIQ score in your reasoning.
"""


def _build_portfolio_prompt(stress_results: list, language: str) -> str:
    scenario_text = "\n".join([
        f"- {r['scenario']}: {r['avg_default_pct']}% avg default, ₹{r['total_loss_lakhs']}L estimated loss"
        for r in stress_results
    ])

    if language == "Hindi":
        return f"""
आप एक पोर्टफोलियो जोखिम प्रबंधक हैं। निम्नलिखित जलवायु तनाव परीक्षण परिणामों के आधार पर हिंदी में एक संक्षिप्त जोखिम रिपोर्ट तैयार करें।

तनाव परीक्षण परिणाम ({stress_results[0]['portfolio_size']} ऋणों का पोर्टफोलियो):
{scenario_text}

निम्नलिखित पर रिपोर्ट दें:
1. सामान्य बनाम सूखे की स्थिति में जोखिम वृद्धि
2. सबसे गंभीर परिदृश्य का प्रभाव
3. पूंजी बफर और प्रावधान की सिफारिश
3-4 वाक्यों में, स्पष्ट और व्यावसायिक भाषा में।
"""
    else:
        return f"""
You are a portfolio risk manager at an agricultural lending institution.
Based on the ClimaIQ stress test results below, write a concise risk narrative for senior management.

Stress Test Results (Portfolio of {stress_results[0]['portfolio_size']} loans):
{scenario_text}

Cover:
1. Risk escalation from normal to stressed conditions
2. The non-linear jump — why losses multiply rather than add
3. Capital buffer and provisioning recommendation
4. Which drought scenario should trigger portfolio rebalancing
Write in 4-5 sentences. Be specific about the numbers. Professional tone.
"""


# ─── Main Explanation Functions ────────────────────────────────────────────────

def explain_for_farmer(
    farmer_data: dict,
    result: dict,
    language: str = "Hindi",
    client=None,
    inference_mode: str = "cloud",
    ollama_base: Optional[str] = None,
    ollama_model: Optional[str] = None,
) -> str:
    """Generate farmer-facing explanation in Hindi or English."""
    prompt = _build_farmer_prompt(farmer_data, result, language)
    return generate_explanation_text(
        prompt,
        inference_mode=inference_mode,
        cloud_client=client,
        ollama_base=ollama_base,
        ollama_model=ollama_model,
    )


def explain_for_officer(
    farmer_data: dict,
    result: dict,
    language: str = "English",
    client=None,
    inference_mode: str = "cloud",
    ollama_base: Optional[str] = None,
    ollama_model: Optional[str] = None,
) -> str:
    """Generate loan officer assessment report in Hindi or English."""
    prompt = _build_officer_prompt(farmer_data, result, language)
    return generate_explanation_text(
        prompt,
        inference_mode=inference_mode,
        cloud_client=client,
        ollama_base=ollama_base,
        ollama_model=ollama_model,
    )


def explain_portfolio_stress(
    stress_results: list,
    language: str = "English",
    client=None,
    inference_mode: str = "cloud",
    ollama_base: Optional[str] = None,
    ollama_model: Optional[str] = None,
) -> str:
    """Generate portfolio-level stress test narrative."""
    prompt = _build_portfolio_prompt(stress_results, language)
    return generate_explanation_text(
        prompt,
        inference_mode=inference_mode,
        cloud_client=client,
        ollama_base=ollama_base,
        ollama_model=ollama_model,
    )


# ─── Quick Test ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import os
    from climaiq_engine import load_model, predict_single, run_stress_test

    API_KEY = input("Enter your Gemma API key: ").strip()
    client = get_client(API_KEY)

    model, scaler = load_model()

    sample_farmer = {
        "age": 42,
        "land_size_acres": 3.5,
        "annual_income_lakhs": 3.0,
        "loan_amount_lakhs": 2.0,
        "previous_defaults": 0,
        "crop_type": "Cotton",
        "state": "Maharashtra",
        "rainfall_deficit_pct": -35.0,
        "spi": -1.8,
        "consecutive_drought_years": 1
    }

    result = predict_single(sample_farmer, model, scaler)

    print("\n" + "="*60)
    print("FARMER EXPLANATION (Hindi)")
    print("="*60)
    print(explain_for_farmer(sample_farmer, result, language="Hindi", client=client))

    print("\n" + "="*60)
    print("OFFICER REPORT (English)")
    print("="*60)
    print(explain_for_officer(sample_farmer, result, language="English", client=client))

    print("\n" + "="*60)
    print("PORTFOLIO STRESS TEST NARRATIVE")
    print("="*60)
    portfolio = [sample_farmer] * 50
    stress = run_stress_test(model, scaler, portfolio)
    print(explain_portfolio_stress(stress, language="English", client=client))