from flask import Flask, render_template, request, jsonify, Response
import requests
from bs4 import BeautifulSoup
from flask import stream_with_context
from concurrent.futures import ThreadPoolExecutor, as_completed

app = Flask(__name__)

# Internal mapping of crops to pests (for the form)
CROP_TO_PESTS = {
    "Sorgum": ["FallArmyWorm"],
    "Maize": ["FallArmyWorm"],
    "Rice": ["Blast", "GallMidge", "YSB", "PlantHopper", "BlueBeetle", "BacterialLeafBlight"],
    "Cotton": ["Thrips", "Whitefly", "PinkBollworm", "Jassid", "BollRot", "AmericanBollworm"],
    "Soybean": ["Girdlebeetle", "H.armigera", "Semilooper", "Spodoptera", "StemFLy"],
    "Tur": ["Wilt", "Webbed_Leaves", "Pod_damage"],
    "Sugarcane": ["FallArmyGrub", "WhiteGrub"],
    "Gram": ["H.armigera", "Wilt"]
}

# Fixed year options for the form
YEARS = ["2024-25", "2023-24", "2022-23", "2021-22"]

# Map our internal crop names to the external page's crop values.
CROP_MAPPING = {
    "Cotton": "1",
    "Gram": "4",
    "Maize": "7",
    "Rice": "3",
    "Sorghum": "6",
    "Soybean": "2",
    "Sugarcane": "8",
    "Tur": "5",
    "Sorgum": "6"
}

# Map our internal pest names to external page values per crop.
PEST_MAPPING = {
    "Cotton": {
        "FallArmyWorm": "71"
    },
    "Gram": {
        "H.armigera": "72",
        "Wilt": "73"
    },
    "Maize": {
        "FallArmyWorm": "74"
    },
    "Rice": {
        "Blast": "75",
        "GallMidge": "76",
        "YSB": "77",
        "PlantHopper": "78",
        "BlueBeetle": "79",
        "BacterialLeafBlight": "80"
    },
    "Soybean": {
        "Girdlebeetle": "81",
        "H.armigera": "82",
        "Semilooper": "83",
        "Spodoptera": "84",
        "StemFLy": "85"
    },
    "Tur": {
        "Wilt": "86",
        "Webbed_Leaves": "87",
        "Pod_damage": "88"
    },
    "Sugarcane": {
        "FallArmyGrub": "89",
        "WhiteGrub": "90"
    },
    "Sorgum": {
        "FallArmyWorm": "91"
    }
}

# Parameter codes and labels for the final image URL
PARAMS = {
    "Mint": "Min Temperature",
    "Maxt": "Max Temperature",
    "RH": "Relative Humidity",
    "RF": "Rainfall",
    "PR": "Pest Report"
}


def build_image_url(crop, year, pest, param, week):
    """Build the external ICAR image URL for a given combination."""
    base = f"http://www.icar-crida.res.in:8080/naip/gisimages/{crop}/{year}/{pest}_"
    return f"{base}{param}{week}.jpg"


def image_exists(url):
    """Return True if the remote image URL returns a valid image (HTTP 200 + image content-type)."""
    try:
        resp = requests.head(url, timeout=4, allow_redirects=True)
        if resp.status_code == 200:
            return resp.headers.get("Content-Type", "").startswith("image/")
        # Some servers don't support HEAD — fall back to a streaming GET
        if resp.status_code in (405, 403):
            resp = requests.get(url, timeout=4, stream=True)
            resp.close()
            return resp.status_code == 200 and resp.headers.get("Content-Type", "").startswith("image/")
    except Exception:
        pass
    return False


def week_has_any_data(week, crop, year, pest):
    """Return week string if any param image exists for it, else None."""
    for param_code in PARAMS:
        if image_exists(build_image_url(crop, year, pest, param_code, week)):
            return week
    return None


@app.route('/')
def index():
    crop = request.args.get('crop', '')
    pest = request.args.get('pest', '')
    year = request.args.get('year', '')
    week = request.args.get('week', '')
    param = request.args.get('param', '')

    image_url = ""
    data_available = True

    if crop and pest and year and week and param:
        external_image_url = build_image_url(crop, year, pest, param, week)
        if image_exists(external_image_url):
            image_url = f"/proxy-image?url={external_image_url}"
        else:
            data_available = False

    return render_template('index.html',
                           crops=list(CROP_TO_PESTS.keys()),
                           crop_to_pests=CROP_TO_PESTS,
                           years=YEARS,
                           params=PARAMS,
                           selected_crop=crop,
                           selected_pest=pest,
                           selected_year=year,
                           selected_week=week,
                           selected_param=param,
                           image_url=image_url,
                           data_available=data_available)


@app.route('/fetch_weeks')
def fetch_weeks():
    """
    Return weeks that have at least one available param image.
    Checks are run in parallel to stay well within the gunicorn timeout.
    """
    crop = request.args.get('crop', '')
    pest = request.args.get('pest', '')
    year = request.args.get('year', '')

    if not (crop and pest and year):
        return jsonify({"weeks": []})

    ext_crop = CROP_MAPPING.get(crop, '')
    ext_pest = PEST_MAPPING.get(crop, {}).get(pest, '')

    # Step 1: scrape candidate weeks from upstream (fast — single request)
    candidate_weeks = []
    try:
        payload = {"country": ext_crop, "city": ext_pest, "sowing": year}
        resp = requests.get(
            "http://www.icar-crida.res.in:8080/naip/gismaps.jsp",
            params=payload, timeout=8
        )
        soup = BeautifulSoup(resp.text, 'html.parser')
        candidate_weeks = [
            opt.get('value') for opt in soup.select('select[name="week"] option')
            if opt.get('value') and "Select" not in opt.get('value', '')
        ]
    except Exception:
        pass

    if not candidate_weeks:
        candidate_weeks = [str(i) for i in range(1, 53)]

    # Step 2: check all candidate weeks in parallel (max 20 threads, 4s per request)
    available_weeks = []
    with ThreadPoolExecutor(max_workers=20) as executor:
        futures = {
            executor.submit(week_has_any_data, week, crop, year, pest): week
            for week in candidate_weeks
        }
        found = set()
        for future in as_completed(futures):
            result = future.result()
            if result is not None:
                found.add(result)

    # Preserve original order
    available_weeks = [w for w in candidate_weeks if w in found]
    return jsonify({"weeks": available_weeks})


@app.route('/check_availability')
def check_availability():
    """
    Check which params are available for a given crop/pest/year/week combination.
    All 5 param checks run in parallel.
    Returns a dict of {param_code: bool}.
    """
    crop = request.args.get('crop', '')
    pest = request.args.get('pest', '')
    year = request.args.get('year', '')
    week = request.args.get('week', '')

    if not (crop and pest and year and week):
        return jsonify({"availability": {}})

    def check_param(param_code):
        url = build_image_url(crop, year, pest, param_code, week)
        return param_code, image_exists(url)

    availability = {}
    with ThreadPoolExecutor(max_workers=5) as executor:
        for param_code, exists in executor.map(lambda p: check_param(p), PARAMS.keys()):
            availability[param_code] = exists

    return jsonify({"availability": availability})


@app.route('/proxy-image')
def proxy_image():
    external_url = request.args.get('url')
    if not external_url:
        return "Missing URL", 400

    try:
        resp = requests.get(external_url, timeout=10, stream=True)
        if resp.status_code != 200 or not resp.headers.get('Content-Type', '').startswith('image/'):
            return "Image not available", 404
        return Response(
            stream_with_context(resp.iter_content(chunk_size=1024)),
            mimetype=resp.headers.get('Content-Type', 'image/jpeg')
        )
    except Exception as e:
        return str(e), 500


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)