"""Format samples from various VLM datasets into unified instruction-following format.

LLaVA-1.5 proved that format prompts are critical:
- VQA: "Answer the question using a single word or phrase."
- MCQ: "Answer with the option's letter from the given choices directly."
- Conversation: Full multi-turn instruction following
"""


def format_for_task(question: str, answer: str, task_type: str) -> dict:
    """Format a QA pair with task-appropriate instruction prompt.

    Args:
        question: The question or instruction
        answer: The expected answer
        task_type: One of "vqa", "mcq", "caption", "conversation", "detection", "alert", "counting", "ocr"

    Returns:
        dict with "instruction", "answer", "task_token"
    """
    task_tokens = {
        "vqa": "<vqa>",
        "mcq": "<vqa>",
        "caption": "<caption>",
        "conversation": "<vqa>",
        "detection": "<detect>",
        "alert": "<alert>",
        "counting": "<count>",
        "ocr": "<ocr>",
        "tracking": "<track>",
        "reasoning": "<reason>",
    }

    task_token = task_tokens.get(task_type, "<vqa>")

    if task_type == "vqa":
        instruction = f"{question}\nAnswer the question using a single word or phrase."
    elif task_type == "mcq":
        instruction = f"{question}\nAnswer with the option's letter from the given choices directly."
    elif task_type == "caption":
        instruction = "Describe this image in detail."
        if question:
            instruction = question
    elif task_type == "detection":
        instruction = f"{question}\nList all objects you can identify." if question else "List all objects visible in this image."
    elif task_type == "alert":
        instruction = f"{question}\nDescribe any anomalies or safety concerns." if question else "Describe any anomalies or safety concerns visible in this image."
    elif task_type == "counting":
        instruction = f"{question}\nProvide the count as a number." if question else "Count the objects in this image."
    elif task_type == "ocr":
        instruction = f"{question}\nRead and transcribe any visible text." if question else "Read and transcribe all visible text in this image."
    elif task_type == "conversation":
        instruction = question
    else:
        instruction = question

    return {
        "instruction": instruction,
        "answer": answer,
        "task_token": task_token,
    }


def classify_dataset_task(dataset_name: str) -> str:
    """Map dataset name to default task type."""
    mapping = {
        "vqav2": "vqa",
        "gqa": "vqa",
        "okvqa": "vqa",
        "aokvqa": "vqa",
        "textvqa": "ocr",
        "ocrvqa": "ocr",
        "textcaps": "caption",
        "llava_instruct": "conversation",
        "visual_genome": "vqa",
        "sharegpt": "conversation",
        "scienceqa": "mcq",
        "cc3m": "caption",
        "sbu": "caption",
        "laion_coco": "caption",
        "coco_detect": "detection",
        "visdrone": "detection",
        "mot": "tracking",
        "ucf_crime": "alert",
        "activitynet": "caption",
        "synthetic_rtsp": "detection",
    }
    return mapping.get(dataset_name, "vqa")