import json
from pathlib import Path
import logging
import fire
from time import sleep
import geopandas as gpd


from cropclass.openeo.inference import croptype_map
from cropclass.utils import laea20km_id_to_extent
from cropclass.config import get_job_options, get_processing_options, get_collection_options
import openeo
from openeo.extra.job_management import MultiBackendJobManager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("openeo_classification.cropmap")

terrascope = openeo.connect("openeo.vito.be").authenticate_oidc()
#terrascope_dev = openeo.connect("openeo-dev.vito.be").authenticate_oidc()
#creo = openeo.connect("openeo-dev.creo.vito.be").authenticate_oidc()
#creo_poc = openeo.connect("http://openeo-poc.creo.vito.be/openeo/1.0").authenticate_oidc()


def produce_on_terrascope():
    produce_eu27_croptype_map(year=2021, parallel_jobs=3, input_file="croptype2021_terrascope.geojson",
                              status_file="croptype_ccn_terrascope_run2.csv", output_dir="/home/driesj/data/ccn/terrascope2")


def produce_on_sentinelhub():
    produce_eu27_croptype_map(provider="sentinelhub", year=2021, parallel_jobs=8,
                              input_file="creo_test.geojson", status_file="hrl_vlcc_eu27_2021_shub3.csv")


def produce_on_creodias():
    produce_eu27_croptype_map(year=2021, parallel_jobs=1, input_file="croptype2021.geojson", status_file="croptype_ccn_nonterrascope_all.csv",
                              output_dir="/home/driesj/data/ccn/creodias", backends=["creodias", "sentinelhub"])


def produce_terrascope_test():
    produce_eu27_croptype_map(year=2021, parallel_jobs=3, input_file="terrascope_test.geojson",
                              status_file="eu27_2021_terrascope_klein.csv", output_dir=".", backends=["terrascope"])


def produce_creo_test():
    produce_eu27_croptype_map(year=2021, parallel_jobs=1, input_file="creo_test.geojson",
                              status_file="creo_sample_fix.csv", output_dir="creo_geoid_fix", backends=["creodias"])


def produce_eu27_croptype_map(year=2021, parallel_jobs=20, input_file="terrascope_test.geojson", status_file="eu27_2021_terrascope_klein.csv", output_dir=".", backends=["terrascope"]):
    """
    Script to start and monitor jobs for the EU27 croptype map project in openEO platform CCN.
    The script can use multiple backends, to maximize throughput. Jobs are tracked in a CSV file, upon failure, the script can resume
    processing by pointing to the same csv file. Delete that file to start processing from scratch.

    @param provider: The data provider: terrascope - sentinelhub - creodias
    @param year:  The year for which to generate a cropmap
    @param parallel_jobs:
    @param status_file: The local file where status should be tracked.
    @return:
    """

    with Path(input_file).open('r') as f:
        tiles_to_produce = gpd.GeoDataFrame.from_features(json.load(f))

    logger.info(
        f"Found {len(tiles_to_produce)} tiles to process. Year: {year}")

    class CustomJobManager(MultiBackendJobManager):

        def __init__(self, poll_sleep=60):
            super().__init__(poll_sleep)

        def on_job_error(self, job, row):
            logs = job.logs()
            error_logs = [l for l in logs if l.level.lower() == "error"]
            job_metadata = job.describe_job()

            title = job_metadata['title']
            base_dir = Path(output_dir)

            if len(error_logs) > 0:
                (base_dir / f'job_{title}_errors.json').write_text(
                    json.dumps(error_logs, indent=2))

        def on_job_done(self, job, row):

            base_dir = Path(output_dir)
            job_metadata = job.describe_job()

            results = job.get_results()
            title = job_metadata['title']
            metadata_file = base_dir / f"{title}-results.json"
            metadata_file.write_text(json.dumps(
                results.get_metadata(), indent=2))

            results.download_files(base_dir,include_stac_metadata=False)
            (base_dir /
             f'job_{title}.json').write_text(json.dumps(job_metadata, indent=2))

            logs = job.logs()
            error_logs = [l for l in logs if l.level.lower(
            ) == "error" and "partial writes" not in l.message]
            orfeo_errors = [
                l for l in error_logs if "orfeo" in l.message.lower()]
            if len(orfeo_errors) > 0:
                row["error_reason"] = "orfeo_error"
                logger.error(f"Job had orfeo error {row}")
            if len(error_logs) > 0:
                (base_dir / f'job_{title}_errors.json').write_text(
                    json.dumps(error_logs, indent=2))
            return row

    def run(row, connection_provider, connection, provider):
        job_options = get_job_options(provider)

        EXTENT_20KM = laea20km_id_to_extent(row['name'])

        print(f"submitting job to {provider}")

        with open(f'/vitodata/EEA_HRL_VLCC/data/ref/METEO/{str(year)}/METEO-{row["name"]}-{str(year)}') as meteo:
            METEO_json = json.load(meteo)

        clf_results = croptype_map(EXTENT_20KM, connection,provider, processing_options={'METEO_data': METEO_json})

        job = clf_results.create_job(
            title=row.title,
            description=row.description,
            out_format="GTiff",
            job_options=job_options, overview_method="mode", filename_prefix=f"croptype_openEO_platform_{row['name']}")#colormap=cropclass.openeo.croptype_colors())

        return job

    manager = CustomJobManager()
    if "terrascope" in backends:
        manager.add_backend(
            "terrascope", connection=terrascope, parallel_jobs=parallel_jobs)

    if "sentinelhub" in backends:
        manager.add_backend(
            "sentinelhub", connection=terrascope, parallel_jobs=parallel_jobs)
    if "creodias" in backends:
        manager.add_backend("creodias", connection=creo, parallel_jobs=parallel_jobs)

    manager.run_jobs(
        df=tiles_to_produce,
        start_job=run,
        output_file=Path(status_file)
    )


if __name__ == '__main__':
  fire.Fire(produce_terrascope_test())