from __future__ import annotations

import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))

import joblib
import pandas as pd
from batteryswap_public.utils import iterate_scenarios, load_dataset

from batteryswapai.competition_features import scenario_history_snapshot


def main() -> None:
    dataset_path = Path(os.environ.get("BATTERYSWAP_DATASET_PATH", "/tmp/data"))
    artifact_path = Path(
        os.environ.get("BATTERYSWAP_PLANNER_PATH", "submission_artifacts/planner.joblib")
    )
    splits = os.environ.get("BATTERYSWAP_SPLITS", "public,private").split(",")
    output_path = Path(os.environ.get("BATTERYSWAP_SUBMISSION_PATH", "submission.csv"))

    planner = joblib.load(artifact_path)

    plans = []
    for split in splits:
        locations, timeseries, hidden_eol_times, scenarios = load_dataset(dataset_path / split)

        for scenario, locs, visible_history, _ in iterate_scenarios(
            locations,
            timeseries,
            hidden_eol_times,
            scenarios,
        ):
            snapshot = scenario_history_snapshot(
                visible_history,
                locs,
                scenario["name"],
                scenario["start_time"],
            )

            plan = planner.plan_snapshot(
                snapshot,
                locs,
                scenario["travel_costs"],
                scenario["settings"],
                scenario["start_time"],
            )

            plan["split"] = split
            plan["scenario"] = scenario["name"]
            plans.append(plan)

    submission = pd.concat(plans, ignore_index=True)
    submission.to_csv(output_path, index=False)

    if not output_path.exists():
        raise RuntimeError(f"Submission was not created: {output_path}")


if __name__ == "__main__":
    main()