import os
from pathlib import Path

# Allow GDAL/OGR to handle large GeoJSON objects
os.environ["OGR_GEOJSON_MAX_OBJ_SIZE"] = "0"

# Source directory (where this repo lives)
SRCDIR = Path(workflow.basedir).parent

# Load configuration
configfile: SRCDIR / "config/config.yaml"

# Data directory from environment variable
DATADIR = Path(os.environ.get("DATADIR", "/data"))

# Taxa list from config
TAXA = config["taxa"]

# Scenarios from config
RADIUS = config["radius"]

# =============================================================================
# Utility Functions for Checkpoint-Based Expansion
# =============================================================================
# These must be defined before the includes so they are available when
# the rule modules are parsed.


def get_species_ids_for_taxa_scenario(wildcards):
    checkpoints.extract_species_data.get().output  # blocks until report.csv exists
    geojson_dir = DATADIR / "species-info" / wildcards.taxa / "current"
    return [p.stem.split('_')[1] for p in sorted(geojson_dir.glob("*.geojson"))]


def get_all_aoh_metadata_for_taxa_scenario(wildcards):
    """
    Return paths to all AOH metadata JSON files for a taxa and scenario.
    """
    species_ids = get_species_ids_for_taxa_scenario(wildcards)
    return [
        DATADIR / "aohs" / wildcards.radius / wildcards.taxa / f"aoh_{sid}_all.json"
        for sid in species_ids
    ]


# =============================================================================
# Elevation layers (precious — only rebuild if explicitly deleted)
# =============================================================================

rule download_elevation:
    """
    Download raw elevation DEM from Zenodo.
    """
    output:
        elevation=DATADIR / "raw_elevation.tif",
    params:
        zenodo_id=config["zenodo"]["elevation"]["zenodo_id"],
        filename=config["zenodo"]["elevation"]["filename"],
    log:
        DATADIR / "logs" / "download_elevation.log",
    shell:
        """
        reclaimer zenodo --zenodo_id {params.zenodo_id} \
                         --filename "{params.filename}" \
                         --output {output.elevation} \
                         2>&1 | tee {log}
        """

rule reproject_elevation:
    """
    Reproject the global elevation map
    """
    input:
        elevation=DATADIR / "raw_elevation.tif",
        mapbiomass=DATADIR / "mapbiomass9" / "MB9_2023_IUCN_habitat.tif",
    output:
        sentinel=DATADIR / "reprojected_elevation.tif"
    log:
        DATADIR / "logs" / "reproject_elevation.log",
    script:
        str(SRCDIR / "scripts" / "reproject_elevation_map.py")

# =============================================================================
# Species Data Extraction (Checkpoint)
# =============================================================================

checkpoint extract_species_data:
    """
    Extract species data from PostgreSQL database.

    This is a checkpoint because the number of output GeoJSON files is only
    known after extraction.

    Environment variables required:
        DB_HOST, DB_NAME, DB_USER, DB_PASSWORD
    """
    input:
        mapbiomass=DATADIR / "mapbiomass9" / "MB9_2023_IUCN_habitat.tif",
    output:
        # The report.csv is the known output; GeoJSON files are dynamic
        report=DATADIR / "species-info" / "report.csv",
    params:
        projection=config["projection"],
        output_dir=lambda wildcards: DATADIR / "species-info",
    resources:
        # Serialise DB access scripts - only one extraction script at a time
        # as it will make many concurrent connections internally
        db_connections=1,
    threads: workflow.cores
    script:
        str(SRCDIR / "scripts" / "extract_species_psql.py")



# =============================================================================
# Per-Species AOH Generation
# =============================================================================


def aoh_species_inputs(wildcards):
    """Return inputs for generate_aoh."""
    return {
        "species_data": DATADIR
            / "species-info"
            / wildcards.taxa
            / "current"
            / f"range_{wildcards.species_id}_all.geojson",
        "habitat_map": DATADIR / "mapbiomass9" / "MB9_2023_IUCN_habitat.tif",
        "elevation_map": DATADIR / "reprojected_elevation.tif",
        "crosswalk": DATADIR / "crosswalk.csv",
    }


rule generate_aoh:
    """
    Generate Area of Habitat raster for a single species in a given scenario.

    Parallelizable: run with `snakemake --cores N` to process multiple species
    concurrently. Uses --force-habitat and --pixel-area flags (no mask).

    For current/scenario AOHs: uses current/ era species data.
    For pnv AOHs: uses historic/ era species data.
    """
    input:
        unpack(aoh_species_inputs),
    output:
        metadata=DATADIR / "aohs" / "{radius}" / "{taxa}" / "aoh_{species_id}_all.json",
    params:
        output_dir=lambda wildcards: DATADIR / "aohs" / wildcards.radius / wildcards.taxa,
        radius=lambda wildcards: wildcards.radius,
    log:
        DATADIR / "logs" / "aoh" / "{radius}" / "{taxa}" / "{species_id}.log",
    script:
        str(SRCDIR / "scripts" / "aohedge.py")


# =============================================================================
# Per-Taxa AOH Aggregation (Checkpoint)
# =============================================================================


checkpoint aggregate_aohs_per_taxa:
    """
    Checkpoint that ensures all AOHs for a taxa/scenario are generated.
    Creates a sentinel file when complete.

    This is a checkpoint so downstream rules (delta P) can re-evaluate the DAG
    after AOHs are created.
    """
    input:
        metadata=get_all_aoh_metadata_for_taxa_scenario,
    output:
        sentinel=DATADIR / "aohs" / "{radius}" / "{taxa}" / ".complete",
    shell:
        """
        echo "Generated $(echo {input.metadata} | wc -w) AOHs for {wildcards.taxa}/{wildcards.radius}"
        touch {output.sentinel}
        """


# =============================================================================
# Collate AOH Data (per scenario)
# =============================================================================


rule collate_aoh_data:
    """
    Collate metadata from all AOH JSON files for a scenario into a single CSV.

    Used by validation (current scenario) and for downstream analysis.
    """
    input:
        sentinels=lambda wildcards: expand(
            str(DATADIR / "aohs" / "{radius}" / "{taxa}" / ".complete"),
            radius=wildcards.radius,
            taxa=TAXA,
        ),
    output:
        collated=DATADIR / "aohs" / "radius_{radius}.csv",
    params:
        aoh_results_dir=lambda wildcards: DATADIR / "aohs" / wildcards.radius,
    log:
        DATADIR / "logs" / "collate_aoh_{radius}.log",
    shell:
        """
        mkdir -p $(dirname {output.collated})
        aoh-collate-data \
            --aoh_results {params.aoh_results_dir} \
            --output {output.collated} \
            2>&1 | tee {log}
        """


rule aohs:
    """Target: generate all AOH rasters and collated summary CSVs."""
    input:
        expand(
            str(DATADIR / "aohs" / "{radius}" / "{taxa}" / ".complete"),
            taxa=TAXA,
            radius=RADIUS,
        ),
        expand(
            str(DATADIR / "aohs" / "radius_{radius}.csv"),
            radius=RADIUS,
        ),

