"""
Tutorial 5: Water + LLMs — Practical Applications
Smart Bhujal — Water + AI Workshop
Duration: 20 minutes | Level: Beginner

Use ChatGPT/Gemini/Claude APIs to build water-sector tools:
  1. Water Quality Report Generator
  2. Borewell Site Advisor
  3. Farmer Irrigation Chatbot
  4. Policy Document Summarizer

Run: python tutorial.py
Needs: OPENAI_API_KEY or GOOGLE_API_KEY in environment
"""

import os
import json

# ── Check which LLM is available ──
OPENAI_KEY = os.environ.get("OPENAI_API_KEY")
GOOGLE_KEY = os.environ.get("GOOGLE_API_KEY")
ANTHROPIC_KEY = os.environ.get("ANTHROPIC_API_KEY")

def get_llm_client():
    """Try OpenAI → Gemini → Claude. Use whichever is available."""
    if OPENAI_KEY:
        from openai import OpenAI
        client = OpenAI()
        def ask(system, user):
            r = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[{"role": "system", "content": system},
                          {"role": "user", "content": user}],
                temperature=0.3
            )
            return r.choices[0].message.content
        print("Using: OpenAI GPT-4o-mini")
        return ask

    elif GOOGLE_KEY:
        import google.generativeai as genai
        genai.configure(api_key=GOOGLE_KEY)
        model = genai.GenerativeModel("gemini-2.0-flash")
        def ask(system, user):
            r = model.generate_content(f"{system}\n\n{user}")
            return r.text
        print("Using: Google Gemini 2.0 Flash")
        return ask

    elif ANTHROPIC_KEY:
        from anthropic import Anthropic
        client = Anthropic()
        def ask(system, user):
            r = client.messages.create(
                model="claude-haiku-4-5-20251001",
                max_tokens=1024,
                system=system,
                messages=[{"role": "user", "content": user}]
            )
            return r.content[0].text
        print("Using: Anthropic Claude Haiku")
        return ask

    else:
        print("=" * 60)
        print("No API key found!")
        print("Set one of these environment variables:")
        print("  export OPENAI_API_KEY=sk-...")
        print("  export GOOGLE_API_KEY=AI...")
        print("  export ANTHROPIC_API_KEY=sk-ant-...")
        print()
        print("Running in DEMO MODE (showing prompts only)")
        print("=" * 60)
        return None


def demo_mode(system, user):
    """Show what would be sent to the LLM."""
    print(f"\n  [SYSTEM PROMPT]")
    print(f"  {system[:200]}...")
    print(f"\n  [USER INPUT]")
    print(f"  {user[:300]}...")
    print(f"\n  [LLM would generate response here]")


# ═══════════════════════════════════════════════════════════════
# APP 1: Water Quality Report Generator
# ═══════════════════════════════════════════════════════════════
def app1_water_quality_report(ask):
    print("\n" + "=" * 60)
    print("APP 1: Water Quality Report Generator")
    print("=" * 60)

    system = """You are a water quality expert. Given lab test results,
generate a concise water quality assessment report with:
1. Parameter-by-parameter comparison with BIS IS:10500 standards
2. Overall verdict (Safe / Unsafe / Needs Treatment)
3. Health risks if any parameter exceeds limits
4. Recommended treatment methods
Keep it under 300 words. Use clear, non-technical language."""

    user = """Generate a water quality report for these lab results:

Sample ID: WQ-2026-042
Location: Anantapur District, Andhra Pradesh
Date: March 2026
Source: Open borewell, 60m depth

Results:
- pH: 7.8
- TDS: 1200 mg/L
- Fluoride: 1.8 mg/L
- Nitrate: 52 mg/L
- Iron: 0.4 mg/L
- Hardness: 450 mg/L
- Chloride: 280 mg/L
- E. coli: Absent

BIS IS:10500 limits for reference:
pH: 6.5-8.5, TDS: <500, Fluoride: <1.0, Nitrate: <45,
Iron: <0.3, Hardness: <200, Chloride: <250"""

    if ask:
        response = ask(system, user)
        print(response)
    else:
        demo_mode(system, user)


# ═══════════════════════════════════════════════════════════════
# APP 2: Borewell Site Advisor
# ═══════════════════════════════════════════════════════════════
def app2_borewell_advisor(ask):
    print("\n" + "=" * 60)
    print("APP 2: Borewell Site Advisor")
    print("=" * 60)

    system = """You are a hydrogeologist advisor. Given location details,
provide borewell drilling recommendations including:
1. Expected depth range based on geology
2. Expected yield estimate
3. Best drilling season
4. Risks and precautions
5. Estimated cost range
Be practical and specific to the region. Under 250 words."""

    user = """Advise on borewell drilling for this location:

Location: Chittoor District, Andhra Pradesh
Geology: Weathered granite with quartz veins
Elevation: 280m above MSL
Nearby wells: 2 borewells within 500m (depths: 120m, 180m)
Rainfall: 850mm/year
Land use: Agricultural (groundnut + paddy)
Water table (nearby): 15-20m below ground
Existing water source: Rain-fed tank (seasonal)

Purpose: Irrigation for 5 acres"""

    if ask:
        response = ask(system, user)
        print(response)
    else:
        demo_mode(system, user)


# ═══════════════════════════════════════════════════════════════
# APP 3: Farmer Irrigation Chatbot
# ═══════════════════════════════════════════════════════════════
def app3_irrigation_chatbot(ask):
    print("\n" + "=" * 60)
    print("APP 3: Farmer Irrigation Chatbot")
    print("=" * 60)

    system = """You are an agricultural water advisor for Indian farmers.
Answer irrigation questions in simple language.
Give specific, actionable advice based on the crop, soil, and season.
Include water quantity in liters/acre/day when possible.
If relevant, mention government schemes (PMKSY, micro-irrigation subsidies).
Keep answers under 200 words."""

    questions = [
        "I'm growing paddy in Kharif season in Tamil Nadu. How much water does it need per acre per day? When should I irrigate?",
        "I want to switch from flood irrigation to drip for my tomato farm (2 acres). What's the water saving and cost?",
        "My borewell yield dropped from 2 inches to 0.5 inches this summer. What should I do?",
    ]

    for i, q in enumerate(questions, 1):
        print(f"\n--- Question {i} ---")
        print(f"Farmer: {q}")
        print(f"\nAdvisor:")
        if ask:
            response = ask(system, q)
            print(response)
        else:
            demo_mode(system, q)
        print()

        if i < len(questions):
            print("(Next question...)")


# ═══════════════════════════════════════════════════════════════
# APP 4: Policy Document Summarizer
# ═══════════════════════════════════════════════════════════════
def app4_policy_summarizer(ask):
    print("\n" + "=" * 60)
    print("APP 4: Policy Document Summarizer")
    print("=" * 60)

    system = """You are a water policy analyst. Summarize the given policy
text into:
1. Key objectives (3-5 bullet points)
2. Main provisions
3. Implementation mechanism
4. Relevance to groundwater management
Keep it under 250 words. Use bullet points."""

    # Sample policy text (Atal Bhujal Yojana summary)
    user = """Summarize this water policy:

ATAL BHUJAL YOJANA (ABY) — Government of India

Atal Bhujal Yojana is a Central Sector Scheme worth Rs 6,000 crore
with World Bank assistance for sustainable groundwater management.
The scheme covers 8,220 gram panchayats in 229 blocks of 80 districts
across 7 states: Gujarat, Haryana, Karnataka, Madhya Pradesh,
Maharashtra, Rajasthan, and Uttar Pradesh.

The scheme has two components:
Component A - Institutional Strengthening (Rs 1,400 crore): Focuses on
strengthening institutional framework, capacity building, and improving
monitoring and data dissemination through network improvements,
capacity building at state and district levels.

Component B - Incentive Component (Rs 4,600 crore): Provides incentives
to states for achievement of specific outcomes related to improved
groundwater management practices. States receive incentive funds
based on: (i) public disclosure of groundwater data, (ii) preparation
of community-led water security plans, (iii) convergence of investments,
(iv) adoption of demand-side management practices, and (v) improvement
in groundwater conditions.

The scheme promotes demand-side management through community-led
approaches, water budgeting at gram panchayat level, and convergence
with existing schemes like MGNREGA, PMKSY, and watershed programs.
Duration: 2020-2025 (extended to 2026)."""

    if ask:
        response = ask(system, user)
        print(response)
    else:
        demo_mode(system, user)


# ═══════════════════════════════════════════════════════════════
# APP 5: Structured Data Extraction from Field Reports
# ═══════════════════════════════════════════════════════════════
def app5_data_extraction(ask):
    print("\n" + "=" * 60)
    print("APP 5: Extract Structured Data from Field Notes")
    print("=" * 60)

    system = """You are a data extraction assistant. Extract structured
data from field survey notes and return as JSON.
Extract: location, coordinates (if mentioned), well_type, depth_m,
water_level_m, yield_lps, geology, water_quality_issues, date.
If a field is not mentioned, use null. Return only valid JSON."""

    user = """Extract data from these field notes:

"Visited Rampur village (lat 15.42, lon 78.31) on 12 March 2026.
Checked the community borewell drilled in 2019. Depth is 150 meters
in weathered granite. Current water level measured at 22 meters below
ground. Yield has reduced to about 1.5 liters per second from the
original 3 lps. Locals report yellowish tinge in water during summer
months - likely iron contamination. The hand pump near the school
(about 200m away) at 40m depth in laterite is still functioning well
with clear water."

Note: This text mentions TWO wells. Extract both as a JSON array."""

    if ask:
        response = ask(system, user)
        print(response)
        # Try to parse and pretty-print
        try:
            # Find JSON in response
            start = response.find('[')
            end = response.rfind(']') + 1
            if start >= 0 and end > start:
                data = json.loads(response[start:end])
                print("\n--- Parsed as Python objects ---")
                for i, well in enumerate(data):
                    print(f"\nWell {i+1}:")
                    for k, v in well.items():
                        print(f"  {k}: {v}")
        except:
            pass
    else:
        demo_mode(system, user)


# ═══════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════
if __name__ == "__main__":
    print("=" * 60)
    print("  Water + LLM Applications")
    print("  Smart Bhujal Tutorial")
    print("=" * 60)

    ask = get_llm_client()

    apps = [
        ("Water Quality Report Generator", app1_water_quality_report),
        ("Borewell Site Advisor", app2_borewell_advisor),
        ("Farmer Irrigation Chatbot", app3_irrigation_chatbot),
        ("Policy Document Summarizer", app4_policy_summarizer),
        ("Field Notes Data Extractor", app5_data_extraction),
    ]

    print("\nAvailable apps:")
    for i, (name, _) in enumerate(apps, 1):
        print(f"  {i}. {name}")
    print(f"  6. Run all")

    choice = input("\nWhich app? (1-6): ").strip()

    if choice == "6":
        for name, func in apps:
            func(ask)
    elif choice in "12345":
        apps[int(choice) - 1][1](ask)
    else:
        print("Running all apps...")
        for name, func in apps:
            func(ask)

    print("\n" + "=" * 60)
    print("Done! Try modifying the prompts for your own use cases.")
    print("=" * 60)