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"

# Maksymalna odległość stacji od toru, żeby "załapało" (w kratkach)
SNAP_DISTANCE = 50 

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

def build_graph(segments):
    """Tworzy sieć połączeń z luźnych kresek"""
    graph = {}
    nodes = set()
    
    print("Budowanie grafu torów...")
    for seg in segments:
        p1 = tuple(seg[0]) # Start odcinka
        p2 = tuple(seg[1]) # Koniec odcinka
        
        # Zaokrąglamy, żeby połączyć punkty, które są blisko siebie (np. 100.0 i 100.001)
        p1 = (round(p1[0], 1), round(p1[1], 1))
        p2 = (round(p2[0], 1), round(p2[1], 1))
        
        if p1 not in graph: graph[p1] = []
        if p2 not in graph: graph[p2] = []
        
        # Obliczamy wagę (długość odcinka)
        d = math.sqrt(dist_sq(p1, p2))
        
        graph[p1].append((p2, d))
        graph[p2].append((p1, d))
        nodes.add(p1)
        nodes.add(p2)
        
    print(f"Graf gotowy. Węzłów: {len(nodes)}")
    return graph, list(nodes)

def find_nearest_node(target, nodes):
    """Znajduje najbliższy kawałek toru dla danego punktu (stacji)"""
    best_node = None
    min_dist = SNAP_DISTANCE**2 # Szukamy w kwadracie promienia
    
    target_tuple = (target[0], target[1])
    
    for node in nodes:
        d = dist_sq(node, target_tuple)
        if d < min_dist:
            min_dist = d
            best_node = node
            
    return best_node

def a_star_search(graph, start, goal):
    """Algorytm szukający drogi w labiryncie torów"""
    frontier = []
    heapq.heappush(frontier, (0, start))
    came_from = {start: None}
    cost_so_far = {start: 0}
    
    while frontier:
        current = heapq.heappop(frontier)[1]
        
        if current == goal:
            break
        
        if current not in graph: continue

        for next_node, distance in graph[current]:
            new_cost = cost_so_far[current] + distance
            if next_node not in cost_so_far or new_cost < cost_so_far[next_node]:
                cost_so_far[next_node] = new_cost
                priority = new_cost + math.sqrt(dist_sq(goal, next_node)) # Heurystyka
                heapq.heappush(frontier, (priority, next_node))
                came_from[next_node] = current
                
    # Odtwarzanie ścieżki
    if goal not in came_from:
        return None # Nie znaleziono drogi (tory są przerwane)
        
    path = []
    curr = goal
    while curr != start:
        path.append(list(curr))
        curr = came_from[curr]
    path.append(list(start))
    path.reverse()
    return path

def main():
    print("--- Start Generowania Szczegółowych Tras ---")
    
    # 1. Wczytaj dane
    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 odczytu plików: {e}")
        return

    # 2. Zbuduj graf
    graph, nodes = build_graph(rails)
    
    # 3. Przetwórz trasy
    new_routes = []
    total_routes = len(data.get('routes', []))
    print(f"Przetwarzanie {total_routes} tras (może to potrwać)...")
    
    for i, route in enumerate(data.get('routes', [])):
        original_path = route['path'] # Lista punktów stacji [A, B, C]
        detailed_path = []
        
        if not original_path: continue
        
        # Dodajemy pierwszy punkt
        detailed_path.append(original_path[0])
        
        # Iterujemy po odcinkach (A->B, B->C, itd.)
        for j in range(len(original_path) - 1):
            start_pt = original_path[j]
            end_pt = original_path[j+1]
            
            # Znajdź najbliższe tory
            rail_start = find_nearest_node(start_pt, nodes)
            rail_end = find_nearest_node(end_pt, nodes)
            
            segment_path = None
            
            if rail_start and rail_end:
                # Szukamy drogi po torach
                segment_path = a_star_search(graph, rail_start, rail_end)
            
            if segment_path:
                # Znaleziono drogę po torach! Dodajemy ją (bez pierwszego punktu, żeby nie dublować)
                detailed_path.extend(segment_path[1:])
            else:
                # Nie znaleziono drogi (brak połączenia) -> Rysuj prostą linię
                detailed_path.append(end_pt)
                
        # Zapisz nową trasę
        route['path'] = detailed_path
        new_routes.append(route)
        
        if i % 5 == 0:
            print(f"-> Przeliczono trasę {i+1}/{total_routes}")

    # 4. Zapisz wynik
    output_data = {
        "stations": data['stations'],
        "routes": new_routes
    }
    
    with open(OUTPUT_FILE, 'w') as f:
        json.dump(output_data, f)
        
    print(f"SUKCES! Zapisano szczegółowe trasy do: {OUTPUT_FILE}")

if __name__ == "__main__":
    main()