# filepath: src/qa_pairs_generation/paths/process_single_file_visibility.py
import json, os, numpy as np
from typing import Dict, Optional, List, Tuple
from pathfinder_visibility import find_path_visibility, _get_label

SEED = 42

CHOICE_OPTIONS = {
    "kitchens": ["stove", "fridge", "sink", "dishwasher", "door", "window", "table"],
    "living_rooms": ["door", "window", "sofa", "loveseat", "armchair", "coffee_table", "side_table", "tv_stand", "television", "bookshelf", "fireplace"],
    "bedrooms": ["door", "window", "bed", "dresser", "wardrobe", "desk", "bookshelf", "chair"],
}


def _eligible_labels(room_objects: List[Dict], room_type: str) -> List[str]:
    pool = CHOICE_OPTIONS.get(room_type, [])
    labels = []
    for o in room_objects:
        lab = _get_label(o).lower()
        if any(tok in lab for tok in pool):
            labels.append(_get_label(o))
    # fallback: all labels
    return labels or [_get_label(o) for o in room_objects]


def process_single_file(file_path: str, file: str, room_type: str) -> Optional[Dict]:
    np.random.seed(SEED)
    layout_id = file.replace(".json", "").replace("real_", "").replace("room_", "")
    with open(file_path, "r") as f:
        data = json.load(f)

    objects = data.get("objects", [])
    # drop rugs from candidates
    labels = [_get_label(o) for o in objects if "rug" not in _get_label(o).lower()]

    start_label = None
    goal_label = None

    # up to 3 attempts to get a path
    for _ in range(3):
        start_label, goal_label = map(str, np.random.choice(labels, 2, replace=False))

        # try both directions (sometimes one way is blocked if start/goal center lands inside buffer)
        for a, b in [(start_label, goal_label), (goal_label, start_label)]:
            path, room, obs, sp, gp = find_path_visibility(data, a, b, clearance_cm=15)  # set to 0 for no clearance
            if path:
                start_label, goal_label = a, b
                break
        if path:
            break

    path = path if path else []

    return {
        "layout_id": layout_id,
        "room_type": room_type,
        "reference_object": start_label,
        "target_object": goal_label,
        "path": path,
        "answer": path,  # alias
    }


# in your runner
from functools import partial
from src.qa_pairs_generation.utils import generate_qa_pairs_with_subsampling, save_and_info


def main_visibility(
    input_dir: str = "data/real_data",
    output_csv: str = "qa_pair_generation_real_data/{parent_folder_name}/{parent_folder_name}_qa_synth.csv",
    enable_subsampling: bool = True,
    bedrooms_count: int = 80,
    living_rooms_count: int = 80,
    kitchens_count: int = 40,
):
    parent = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
    output_csv = output_csv.format(parent_folder_name=parent)

    subsample = {"bedrooms": bedrooms_count, "living_rooms": living_rooms_count, "kitchens": kitchens_count} if enable_subsampling else None

    qa_pairs = generate_qa_pairs_with_subsampling(
        input_dir=input_dir, process_single_file=process_single_file, SEED=42, subsample_config=subsample  # <- signature matches your utils
    )
    save_and_info(qa_pairs, output_csv)
