import argparse
import json

from cockpit_grounding.models.factory import create_grounder

from cockpit_grounding.grounding.predictor import (
    build_grounding_prompt,
    parse_grounding_output,
)

from cockpit_grounding.vision.visualize import (
    visualize_grounding,
)


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--image",
        required=True,
    )

    parser.add_argument(
        "--target",
        required=True,
    )

    parser.add_argument(
        "--output",
        required=True,
    )

    parser.add_argument(
        "--model",
        required=True,
        help="Local model path",
    )

    parser.add_argument(
        "--backend",
        choices=("qwen3vl", "qwen35"),
        default="qwen3vl",
        help="Model backend (default: qwen3vl)",
    )

    args = parser.parse_args()

    # ----------------------------
    # Load model
    # ----------------------------

    grounder = create_grounder(
        backend=args.backend,
        model_path=args.model,
    )

    # ----------------------------
    # Prompt
    # ----------------------------

    prompt = build_grounding_prompt(
        args.target
    )

    # ----------------------------
    # Inference
    # ----------------------------

    raw_output = grounder.generate(
        args.image,
        prompt,
    )

    print()
    print("========== RAW MODEL OUTPUT ==========")
    print(raw_output)

    # ----------------------------
    # Parse bbox
    # ----------------------------

    result = parse_grounding_output(
        raw_output
    )

    # ----------------------------
    # Convert + draw
    # ----------------------------

    pixel_result = visualize_grounding(
        args.image,
        result,
        args.output,
    )

    final_result = {
        "target": args.target,

        "bbox_relative": [
            result.x1,
            result.y1,
            result.x2,
            result.y2,
        ],

        **pixel_result,
    }

    print()
    print("========== FINAL RESULT ==========")

    print(
        json.dumps(
            final_result,
            indent=2,
            ensure_ascii=False,
        )
    )

    print()
    print(
        f"Result image: {args.output}"
    )


if __name__ == "__main__":
    main()