# filepath: src/qa_pairs_generation/pathfinder_visibility.py
import os, json, math
from typing import List, Dict, Tuple, Optional, Callable
import numpy as np

from shapely.geometry import Polygon as ShpPolygon, LineString, Point
from shapely.ops import unary_union, linemerge, polygonize
from shapely.validation import make_valid

# =============== Helpers ===============

EPS = 1e-9


def _get_label(o: Dict) -> str:
    return (o.get("label") or o.get("name") or "").strip()


def _points_to_polygon(points: List[Dict[str, float]]) -> Optional[ShpPolygon]:
    if not points:
        return None
    coords = [(p["x"], p["y"]) for p in points]
    if len(coords) >= 3:
        if coords[0] != coords[-1]:
            coords.append(coords[0])
        poly = ShpPolygon(coords)
        if poly.is_valid and not poly.is_empty and poly.area > 0:
            return poly
    return None  # lines/degenerate -> not an area obstacle


def _get_center(o: Dict) -> Tuple[float, float]:
    """Center for any object: polygon centroid if area>0; else midpoint/mean of points."""
    pts = o.get("points")
    if pts:
        poly = _points_to_polygon(pts)
        if poly is not None:
            c = poly.centroid
            return (float(c.x), float(c.y))
        # degenerate (e.g., door/window as a segment)
        xs = [p["x"] for p in pts]
        ys = [p["y"] for p in pts]
        return (sum(xs) / len(xs), sum(ys) / len(ys))
    # fallback for bbox-only old data
    bbox = o.get("bbox")
    if bbox and len(bbox) == 4:
        x1, y1, x2, y2 = bbox
        return ((x1 + x2) / 2.0, (y1 + y2) / 2.0)
    return (0.0, 0.0)


def _walls_to_room_polygon(walls: List[Dict]) -> Optional[ShpPolygon]:
    """Robustly assemble room polygon from unordered/reversed wall segments."""
    segs = []
    for w in walls or []:
        a = (w["start"]["x"], w["start"]["y"])
        b = (w["end"]["x"], w["end"]["y"])
        if a != b:
            segs.append(LineString([a, b]))
    if not segs:
        return None
    merged = linemerge(unary_union(segs))
    faces = list(polygonize(merged))
    if not faces:
        return None
    room = make_valid(max(faces, key=lambda p: p.area))
    if room.is_empty:
        return None
    return room


def _segment_ok_in_room(seg: LineString, room: ShpPolygon) -> bool:
    # allow touching boundary (covers)
    return room.covers(seg)


def _segment_clear_of_obstacles(seg: LineString, obs_union: Optional[ShpPolygon]) -> bool:
    if obs_union is None or obs_union.is_empty:
        return True
    # disallow touching/intruding into buffered obstacles (treat touch as blocked)
    return not obs_union.intersects(seg)


def _round_node(pt: Tuple[float, float], nd=6) -> Tuple[float, float]:
    return (round(pt[0], nd), round(pt[1], nd))


# =============== Visibility Graph ===============


def build_visibility_graph(
    room: ShpPolygon, obstacles: List[ShpPolygon], start_pt: Tuple[float, float], goal_pt: Tuple[float, float], clearance_m: float = 0.15
) -> Dict[Tuple[float, float], Dict[Tuple[float, float], float]]:
    """
    Build visibility graph with clearance by buffering obstacles.
    Nodes: room boundary vertices + obstacle vertices + start + goal.
    Edges: straight segments fully inside room and not intersecting buffered obstacles.
    """
    # Buffer obstacles by clearance (merge once)
    obs_buf = None
    if obstacles:
        if clearance_m > 0:
            obstacles = [o.buffer(clearance_m, join_style=2) for o in obstacles]  # mitered corners
        obs_buf = unary_union(obstacles)
        if obs_buf.is_empty:
            obs_buf = None

    # Candidate nodes
    nodes: List[Tuple[float, float]] = []

    # room vertices
    for x, y in room.exterior.coords[:-1]:
        nodes.append((x, y))
    # obstacle vertices
    for o in obstacles:
        if o is None or o.is_empty:
            continue
        for x, y in o.exterior.coords[:-1]:
            nodes.append((x, y))
    # start/goal
    nodes.append(start_pt)
    nodes.append(goal_pt)

    # Dedup + keep only nodes inside room and off buffered obstacles
    uniq = []
    seen = set()
    for x, y in nodes:
        pt = _round_node((x, y))
        if pt in seen:
            continue
        if not room.covers(Point(pt[0], pt[1])):
            continue
        if obs_buf is not None and obs_buf.intersects(Point(pt[0], pt[1])):
            continue
        seen.add(pt)
        uniq.append(pt)

    nodes = uniq

    # Build adjacency
    graph: Dict[Tuple[float, float], Dict[Tuple[float, float], float]] = {n: {} for n in nodes}
    # For all pairs (O(N^2)); typical N is modest; OK
    for i in range(len(nodes)):
        for j in range(i + 1, len(nodes)):
            a = nodes[i]
            b = nodes[j]
            seg = LineString([a, b])
            # room check
            if not _segment_ok_in_room(seg, room):
                continue
            # obstacle clearance
            if not _segment_clear_of_obstacles(seg, obs_buf):
                continue
            dist = seg.length
            graph[a][b] = dist
            graph[b][a] = dist
    return graph


# =============== Shortest Path (Dijkstra) ===============


def shortest_path_dijkstra(graph: Dict[Tuple[float, float], Dict[Tuple[float, float], float]], start: Tuple[float, float], goal: Tuple[float, float]) -> List[Tuple[float, float]]:
    import heapq

    INF = float("inf")
    dist = {n: INF for n in graph}
    prev = {}
    dist[start] = 0.0
    h = [(0.0, start)]
    seen = set()
    while h:
        d, u = heapq.heappop(h)
        if u in seen:
            continue
        seen.add(u)
        if u == goal:
            break
        for v, w in graph[u].items():
            nd = d + w
            if nd < dist[v]:
                dist[v] = nd
                prev[v] = u
                heapq.heappush(h, (nd, v))
    if goal not in prev and start != goal:
        return []  # no path
    # reconstruct
    path = [goal]
    cur = goal
    while cur != start:
        cur = prev.get(cur)
        if cur is None:
            return []
        path.append(cur)
    path.reverse()
    return path


# =============== Public API used by your generator ===============


def find_path_visibility(room_data: Dict, start_label: str, goal_label: str, clearance_cm: int = 15):
    """
    Returns (path_meters, room_polygon, obstacle_polys, start_pt, goal_pt).
    path_meters is a list of (x,y) in meters. Empty list if none.
    """
    walls = room_data.get("room", {}).get("walls", [])
    room = _walls_to_room_polygon(walls)
    if room is None:
        return [], None, [], None, None

    # Objects (polygons only; lines like windows/doors won’t obstruct)
    obs_polys: List[ShpPolygon] = []
    start_obj = None
    goal_obj = None
    for o in room_data.get("objects", []):
        label = _get_label(o)
        if label == start_label:
            start_obj = o
        if label == goal_label:
            goal_obj = o
        # ignore rugs explicitly if you want
        if "rug" in label:
            continue
        poly = _points_to_polygon(o.get("points", []))
        if poly is not None:
            obs_polys.append(poly)

    if start_obj is None or goal_obj is None:
        return [], room, obs_polys, None, None

    start_pt = _round_node(_get_center(start_obj))
    goal_pt = _round_node(_get_center(goal_obj))

    # Exclude start/goal objects themselves from obstacles (so we can leave/arrive next to them)
    def _same_obj(a, b):
        return a is b

    obs_polys_pruned = []
    for o in room_data.get("objects", []):
        if _same_obj(o, start_obj) or _same_obj(o, goal_obj):
            continue
        poly = _points_to_polygon(o.get("points", []))
        if poly is not None:
            obs_polys_pruned.append(poly)

    graph = build_visibility_graph(room, obs_polys_pruned, start_pt, goal_pt, clearance_m=clearance_cm / 100.0)
    if start_pt not in graph or goal_pt not in graph:
        return [], room, obs_polys_pruned, start_pt, goal_pt

    path = shortest_path_dijkstra(graph, start_pt, goal_pt)
    return path, room, obs_polys_pruned, start_pt, goal_pt
