import argparse
import json
from pathlib import Path


TASKS = ("semantics", "interactable", "interaction")


def extract_image_id(image_name):
    stem = Path(image_name).stem
    parts = stem.split("_")
    if len(parts) != 2:
        raise ValueError(f"Cannot infer image_id from image name: {image_name}")
    app_id, frame_id = parts
    return int(f"{app_id}{int(frame_id):03d}")


def load_questions(path):
    if path is None:
        return {}

    mapping = {}
    with Path(path).open() as file:
        for line_number, line in enumerate(file, start=1):
            line = line.strip()
            if not line:
                continue
            question = json.loads(line)
            try:
                question_id = str(question["question_id"])
                image_id = question.get("image_id")
                if image_id is None:
                    image_id = extract_image_id(question["image"])
            except KeyError as exc:
                raise ValueError(
                    f"Question file {path} line {line_number} is missing {exc.args[0]!r}"
                ) from exc
            mapping[question_id] = image_id
    return mapping


def normalize_bbox(item):
    bbox = item.get("bbox", item.get("bbox_pixels"))
    if bbox is None:
        raise ValueError(f"Prediction is missing bbox/bbox_pixels: {item}")
    if len(bbox) != 4:
        raise ValueError(f"Prediction bbox must have four values: {item}")
    return bbox


def normalize_score(item):
    return item.get("score", item.get("probability", 1.0))


def normalize_category(category, item, task):
    if task == "interactable":
        return 1

    category_id = item.get("category_id", category)
    if category_id is None:
        raise ValueError(f"Prediction is missing category/category_id for {task}: {item}")
    return category_id


def normalize_image_id(question_id, content, item, questions):
    for source in (item, content):
        if isinstance(source, dict) and "image_id" in source:
            return source["image_id"]
        if isinstance(source, dict) and "image" in source:
            return extract_image_id(source["image"])

    if question_id is not None and str(question_id) in questions:
        return questions[str(question_id)]

    raise ValueError(
        "Prediction is missing image_id/image. Provide --questions for old "
        f"question-keyed prediction files. question_id={question_id!r}"
    )


def iter_old_format(data):
    for question_id, content in data.items():
        if not isinstance(content, dict):
            raise ValueError(f"Prediction for question {question_id!r} must be an object")

        results = content.get("oovd_result")
        if results is None:
            continue
        if not isinstance(results, dict):
            raise ValueError(f"oovd_result for question {question_id!r} must be an object")

        for category, objects in results.items():
            if not isinstance(objects, list):
                raise ValueError(
                    f"oovd_result[{category!r}] for question {question_id!r} must be a list"
                )
            for item in objects:
                if not isinstance(item, dict):
                    raise ValueError(f"Prediction item must be an object: {item!r}")
                yield question_id, content, category, item


def iter_prediction_items(data):
    if isinstance(data, dict):
        yield from iter_old_format(data)
        return

    if not isinstance(data, list):
        raise ValueError("Prediction input must be a list or a question-keyed object")

    for item in data:
        if not isinstance(item, dict):
            raise ValueError(f"Prediction item must be an object: {item!r}")
        yield None, {}, item.get("category_id"), item


def convert_predictions(input_path, task, questions_path=None):
    with Path(input_path).open() as file:
        data = json.load(file)

    questions = load_questions(questions_path)
    output = []
    for question_id, content, category, item in iter_prediction_items(data):
        output.append(
            {
                "image_id": normalize_image_id(question_id, content, item, questions),
                "category_id": normalize_category(category, item, task),
                "bbox": normalize_bbox(item),
                "score": normalize_score(item),
            }
        )
    return output


def output_path_for_task(output_path, task):
    path = Path(output_path)
    if path.suffix:
        return path.with_name(f"{path.stem}_{task}{path.suffix}")
    return path / f"{task}.json"


def write_json(path, data):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w") as file:
        json.dump(data, file, indent=2)
        file.write("\n")


def build_parser():
    parser = argparse.ArgumentParser(
        description="Convert Orienter prediction files to COCO-style result JSON."
    )
    parser.add_argument(
        "--task",
        choices=TASKS + ("all",),
        required=True,
        help="Evaluation task to convert for.",
    )
    parser.add_argument("--input", required=True, help="Prediction JSON path.")
    parser.add_argument(
        "--questions",
        help="Question JSONL path. Required only for old question-keyed inputs without image_id.",
    )
    parser.add_argument("--output", required=True, help="Output JSON path or directory.")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    tasks = TASKS if args.task == "all" else (args.task,)

    for task in tasks:
        converted = convert_predictions(args.input, task, args.questions)
        output_path = (
            output_path_for_task(args.output, task) if args.task == "all" else Path(args.output)
        )
        write_json(output_path, converted)


if __name__ == "__main__":
    main()