src/db_manager.py

import duckdb # type: ignore
import os # type: ignore
import sys # type: ignore
from pathlib import Path # type: ignore
from dotenv import load_dotenv # type: ignore


# Add parent directory to path to allow imports from api/
sys.path.insert(0, str(Path(__file__).parent.parent))

from api.update_static import GTFSSyncManager
from api.bus_cache import AsyncBusCache # type: ignore

# Configuration - always save DB in src/ directory
DB_PATH = str(Path(__file__).parent / "ttc_gtfs.duckdb")
STATIC_DIR = str(Path(__file__).parent.parent / "static")

def init_db():
    sync_mgr = GTFSSyncManager()
    remote_data = sync_mgr.get_remote_metadata()
    
    # Connect to existing or new DB
    con = duckdb.connect(sync_mgr.DB_PATH)
    
    # 1. Setup metadata tracking
    con.execute("CREATE TABLE IF NOT EXISTS sync_metadata (key VARCHAR PRIMARY KEY, value VARCHAR)")
    local_update = con.execute("SELECT value FROM sync_metadata WHERE key = 'last_modified'").fetchone()

    # 2. Check if we need to sync based on API metadata
    should_sync = False
    if not local_update or (remote_data and remote_data["updated_at"] > local_update[0]):
        should_sync = True
    
    if should_sync and remote_data:
        print(f"--- Data Stale. Remote: {remote_data['updated_at']} | Local: {local_update[0] if local_update else 'None'} ---")
        con.close() # Close to allow file deletion
        sync_mgr.perform_full_sync(remote_data["url"])
        
        # Reconnect and finalize metadata
        con = duckdb.connect(sync_mgr.DB_PATH)
        con.execute("CREATE TABLE IF NOT EXISTS sync_metadata (key VARCHAR PRIMARY KEY, value VARCHAR)")
        con.execute("INSERT OR REPLACE INTO sync_metadata VALUES ('last_modified', ?)", [remote_data["updated_at"]])

    # 3. Standard Import Loop (runs if DB was nuked or is missing tables)
    tables = [t[0] for t in con.execute("SHOW TABLES").fetchall()]
    if all(t in tables for t in ["routes", "trips", "stops", "stop_times", "shapes"]):
        return con

    print("--- Initializing/Updating DuckDB: Importing CSVs ---")
    
    # Updated files list to include the missing shapes
    files = ["routes.txt", "trips.txt", "stops.txt", "stop_times.txt", "shapes.txt"]
    
    for f in files:
        file_path = Path(STATIC_DIR) / f
        table_name = f.replace(".txt", "")
        
        if file_path.exists():
            print(f"Importing {f} into table '{table_name}'...")
            abs_file_path = str(file_path.resolve())
            
            # Use 'CREATE OR REPLACE' to overwrite existing tables without crashing
            con.execute(f"CREATE OR REPLACE TABLE {table_name} AS SELECT * FROM read_csv_auto('{abs_file_path}')")
        else:
            print(f"Error: {file_path} not found!")
        
    # Add this inside init_db in db_manager.py after the file import loop
    print("--- Creating Indexes for Performance ---")
    # This speeds up the /api/shapes/{shape_id} endpoint significantly
    con.execute("CREATE INDEX IF NOT EXISTS idx_shape_id ON shapes (shape_id)")

    # While you're at it, indexing trip_id in stop_times speeds up your arrival logic
    con.execute("CREATE INDEX IF NOT EXISTS idx_stop_times_trip_id ON stop_times (trip_id)")

    print("--- Database Import Complete ---")
    return con

async def test_data_integrity(con):
    """
    Runs a test join to confirm that a trip ID can be linked to a route name and stop list.
    Uses AsyncBusCache to get a real trip_id from the live API.
    """
    print("--- Running Integrity Test ---")
    try:
        # Get trip_id from live API using AsyncBusCache
        cache = AsyncBusCache(ttl=20)
        vehicles = await cache.get_data()
        
        if not vehicles:
            print("No vehicles available from API, falling back to database trip_id")
            sample_trip = con.execute("SELECT trip_id FROM trips LIMIT 1").fetchone()[0]
        else:
            # Extract trip_id from the first vehicle
            # We need to get the raw GTFS data to access trip_id
            import httpx # type: ignore
            from google.transit import gtfs_realtime_pb2 # type: ignore
            
            load_dotenv()
            gtfs_rt_url = os.getenv("GTFS_RT_URL")
            if not gtfs_rt_url:
                raise ValueError("GTFS_RT_URL is not set")
            
            async with httpx.AsyncClient() as client:
                response = await client.get(gtfs_rt_url, timeout=10)
                response.raise_for_status()
            
            feed = gtfs_realtime_pb2.FeedMessage()
            feed.ParseFromString(response.content)
            
            # Get trip_id from first vehicle entity
            sample_trip = None
            for entity in feed.entity:
                if entity.HasField('vehicle') and entity.vehicle.trip.trip_id:
                    sample_trip = entity.vehicle.trip.trip_id
                    break
            
            if not sample_trip:
                print("No trip_id found in API response, falling back to database")
                sample_trip = con.execute("SELECT trip_id FROM trips LIMIT 1").fetchone()[0]
            else:
                print(f"Using trip_id from live API: {sample_trip}")
        
        # First, get the total count
        count_query = f"""
            SELECT COUNT(*) 
            FROM trips t
            JOIN stop_times st ON t.trip_id = st.trip_id
            WHERE t.trip_id = '{sample_trip}'
        """
        total_count = con.execute(count_query).fetchone()[0]
        
        # Determine sample size - show all if <= 20, otherwise show first 20
        sample_size = min(20, total_count) if total_count > 20 else total_count
        
        query = f"""
            SELECT 
                r.route_short_name, 
                t.trip_headsign, 
                st.stop_sequence, 
                s.stop_name
            FROM trips t
            JOIN routes r ON t.route_id = r.route_id
            JOIN stop_times st ON t.trip_id = st.trip_id
            JOIN stops s ON st.stop_id = s.stop_id
            WHERE t.trip_id = '{sample_trip}'
            ORDER BY st.stop_sequence
            LIMIT {sample_size};
        """
        results = con.execute(query).fetchall()
        
        print(f"\nSuccessfully joined data for Trip ID: {sample_trip}")
        print(f"Total stops in trip: {total_count}")
        if total_count > sample_size:
            print(f"Showing first {sample_size} stops (sample):\n")
        else:
            print(f"Showing all {total_count} stops:\n")
        
        print(f"{'Route':<8} {'Headsign':<30} {'Stop #':<8} {'Stop Name':<50}")
        print("-" * 100)
        for res in results:
            route = res[0] or "N/A"
            headsign = (res[1] or "N/A")[:28]  # Truncate if too long
            stop_seq = res[2]
            stop_name = (res[3] or "N/A")[:48]  # Truncate if too long
            print(f"{route:<8} {headsign:<30} {stop_seq:<8} {stop_name:<50}")
            
    except Exception as e:
        print(f"Integrity test failed: {e}")

if __name__ == "__main__":
    import asyncio # type: ignore
    db_con = init_db()
    asyncio.run(test_data_integrity(db_con))