import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from crop.extract_soup_area import extract_one_image
from density.predict_density import predict_density


REQUIRED_OUTPUT_KEYS = [
    "image_id",
    "source_image_path",
    "crop_status",
    "crop_confidence",
    "crop_image_path",
    "predicted_label",
    "predicted_label_id",
    "density_confidence",
    "density_probabilities",
    "model_versions",
]


def run_pipeline(
    image_path: str | Path,
    crop_output_dir: str | Path = "data/processed_images/soup_density_crop",
    crop_predictor=None,
    density_model_path: str | Path = "models/density/density_classifier_extracted_v1/model.pt",
) -> dict:
    crop_result = extract_one_image(
        image_path,
        crop_output_dir,
        checkpoint_path="models/crop/soup_valid_area_rfdetr_seg_v2/checkpoint_best_regular.pth",
        threshold=0.35,
        low_confidence_threshold=0.35,
        predictor=crop_predictor,
        save_alpha=False,
    )
    density = (
        predict_density(crop_result["soup_crop_path"], density_model_path)
        if crop_result["extract_status"] != 0
        else {"predicted_label": None, "predicted_label_id": None, "confidence": 0.0, "probabilities": {}}
    )
    return {
        "image_id": crop_result["image_id"],
        "source_image_path": crop_result["source_image_path"],
        "crop_status": crop_result["extract_status"],
        "crop_confidence": crop_result["extract_confidence"],
        "crop_image_path": crop_result["soup_crop_path"],
        "predicted_label": density["predicted_label"],
        "predicted_label_id": density["predicted_label_id"],
        "density_confidence": density["confidence"],
        "density_probabilities": density["probabilities"],
        "model_versions": {
            "crop": "soup_valid_area_rfdetr_seg_v2",
            "density": "density_classifier_extracted_v1",
        },
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Run soup density end-to-end inference pipeline.")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", default="reports/pipeline_result.json")
    parser.add_argument("--crop-output-dir", default="data/processed_images/soup_density_crop")
    parser.add_argument("--density-model-path", default="models/density/density_classifier_extracted_v1/model.pt")
    args = parser.parse_args()
    result = run_pipeline(args.input, args.crop_output_dir, density_model_path=args.density_model_path)
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())