import json
import math
import heapq
import sys

# --- KONFIGURACJA ---
BASE_PATH = "/var/lib/pterodactyl/volumes/cc7ccc03-03bd-4bde-8488-3057d3803420/squaremap/web"
RAILS_FILE = f"{BASE_PATH}/rails.json"
STATIONS_FILE = f"{BASE_PATH}/stations.json"
OUTPUT_FILE = f"{BASE_PATH}/final_map_data.json"

# Tolerancja łączenia torów (w kratkach). Jeśli tory są blisko, uznajemy je za połączone.
CONNECTION_TOLERANCE = 0.5 
# Jak daleko stacja może być od toru, żeby ją przypiąć
SNAP_DISTANCE = 64

def dist_sq(p1, p2):
    return (p1[0] - p2[0])**2 + (p1[1] - p2[1])**2

def point_to_key(pt):
    """Zaokrągla punkt, żeby połączyć bliskie węzły"""
    # Zaokrąglamy do 0.5, żeby złapać drobne przesunięcia
    return (round(pt[0] * 2) / 2, round(pt[1] * 2) / 2)

def build_graph(rail_paths):
    graph = {}
    nodes = set()
    
    print("Budowanie grafu połączeń...")
    
    # rail_paths to teraz lista list punktów: [[p1, p2, p3...], [p1, p2...]]
    for path in rail_paths:
        # Traktujemy każdy mały odcinek łuku jako krawędź grafu
        for i in range(len(path) - 1):
            p1 = tuple(path[i])
            p2 = tuple(path[i+1])
            
            # Klucze do grafu (zaokrąglone)
            k1 = point_to_key(p1)
            k2 = point_to_key(p2)
            
            if k1 not in graph: graph[k1] = []
            if k2 not in graph: graph[k2] = []
            
            d = math.sqrt(dist_sq(p1, p2))
            
            # Dodajemy krawędź. Zapamiętujemy dokładne współrzędne sąsiada
            graph[k1].append((k2, d, p2)) 
            graph[k2].append((k1, d, p1))
            
            nodes.add(k1)
            nodes.add(k2)
            
    print(f"Graf gotowy. Węzłów: {len(nodes)}")
    return graph, list(nodes)

def find_nearest_node_key(target, nodes):
    """Znajduje najbliższy węzeł grafu dla danego punktu"""
    best_key = None
    min_dist = SNAP_DISTANCE**2
    
    target_tuple = (target[0], target[1])
    
    # Optymalizacja: można by użyć k-d tree, ale przy tej ilości węzłów brute-force jeszcze ujdzie
    # lub przeszukać tylko subset, ale tutaj robimy brute-force z limitem.
    
    for k in nodes:
        # Szybka eliminacja
        if abs(k[0] - target[0]) > SNAP_DISTANCE or abs(k[1] - target[1]) > SNAP_DISTANCE:
            continue
            
        d = dist_sq(k, target_tuple)
        if d < min_dist:
            min_dist = d
            best_key = k
            
    return best_key

def a_star(graph, start_key, goal_key):
    frontier = []
    heapq.heappush(frontier, (0, start_key))
    came_from = {start_key: None}
    cost_so_far = {start_key: 0}
    # Przechowujemy prawdziwe punkty do rysowania
    real_points = {} 
    
    while frontier:
        current_key = heapq.heappop(frontier)[1]
        
        if current_key == goal_key:
            break
        
        if current_key not in graph: continue

        for next_key, distance, real_pt in graph[current_key]:
            new_cost = cost_so_far[current_key] + distance
            if next_key not in cost_so_far or new_cost < cost_so_far[next_key]:
                cost_so_far[next_key] = new_cost
                priority = new_cost + math.sqrt(dist_sq(goal_key, next_key))
                heapq.heappush(frontier, (priority, next_key))
                came_from[next_key] = current_key
                real_points[next_key] = real_pt # Zapamiętujemy, jak dokładnie wyglądał ten punkt
                
    if goal_key not in came_from:
        return None
        
    path = []
    curr = goal_key
    while curr != start_key:
        # Dodajemy zapamiętany prawdziwy punkt (żeby nie był zaokrąglony)
        pt = real_points.get(curr, curr)
        path.append(list(pt))
        curr = came_from[curr]
    
    # Dodajemy start
    path.append(list(start_key))
    path.reverse()
    return path

def main():
    print("--- Start Generowania Tras (v2 - Łuki i Tolerancja) ---")
    
    try:
        with open(RAILS_FILE, 'r') as f: rails = json.load(f)
        with open(STATIONS_FILE, 'r') as f: data = json.load(f)
    except Exception as e:
        print(f"Błąd: {e}")
        return

    graph, nodes = build_graph(rails)
    
    new_routes = []
    total = len(data.get('routes', []))
    success_count = 0
    
    print(f"Przetwarzanie {total} tras...")
    
    for i, route in enumerate(data.get('routes', [])):
        original_path = route['path']
        detailed_path = []
        
        if not original_path: continue
        detailed_path.append(original_path[0])
        
        for j in range(len(original_path) - 1):
            p_start = original_path[j]
            p_end = original_path[j+1]
            
            # 1. Znajdź wejście na tory
            k_start = find_nearest_node_key(p_start, nodes)
            k_end = find_nearest_node_key(p_end, nodes)
            
            segment_points = None
            
            if k_start and k_end:
                # 2. Szukaj drogi
                segment_points = a_star(graph, k_start, k_end)
            
            if segment_points:
                # Znaleziono drogę po torach!
                detailed_path.extend(segment_points)
            else:
                # Brak połączenia - rysujemy prostą linię (fallback)
                # Ale dodajemy mały log, żebyś wiedział
                # print(f"  [!] Brak torów między punktami trasy {route['name']}")
                detailed_path.append(p_end)
        
        if len(detailed_path) > len(original_path):
            success_count += 1
            
        route['path'] = detailed_path
        new_routes.append(route)
        
        if (i+1) % 10 == 0:
            print(f"-> {i+1}/{total}...")

    output_data = {"stations": data['stations'], "routes": new_routes}
    
    with open(OUTPUT_FILE, 'w') as f:
        json.dump(output_data, f)
        
    print(f"SUKCES! Zapisano mapę. {success_count}/{total} tras znalazło tory.")

if __name__ == "__main__":
    main()