import sqlite3
from datetime import date
from typing import List, Dict, Optional

DB_PATH = 'schedules.db'

def init_db():
    """Створює таблиці trains, stations та schedules, якщо їх ще нема."""
    with sqlite3.connect(DB_PATH) as con:
        con.execute('''
            CREATE TABLE IF NOT EXISTS trains (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                train_number TEXT UNIQUE NOT NULL,
                days TEXT NOT NULL  -- бінарна маска (Mon→Sun)
            );
        ''')
        con.execute('''
            CREATE TABLE IF NOT EXISTS stations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL,
                km REAL
            );
        ''')
        con.execute('''
            CREATE TABLE IF NOT EXISTS schedules (
                train_id INTEGER NOT NULL,
                station_id INTEGER NOT NULL,
                arrival_time TEXT,
                departure_time TEXT,
                travel_date DATE NOT NULL,
                fetched_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (train_id, station_id, travel_date),
                FOREIGN KEY(train_id) REFERENCES trains(id),
                FOREIGN KEY(station_id) REFERENCES stations(id)
            );
        ''')
        con.commit()

def save_schedule(direction: str, entries: List[Dict]):
    """Зберігає повний розклад для заданого напряму."""
    today = date.today().isoformat()
    with sqlite3.connect(DB_PATH) as con:
        # Видаляємо старі записи за сьогоднішню дату
        con.execute('DELETE FROM schedules WHERE travel_date = ?', (today,))

        for e in entries:
            train_number = e['train_number']
            days = e['days']
            # Додаємо або оновлюємо поїзд
            con.execute('''
                INSERT INTO trains (train_number, days)
                VALUES (?, ?)
                ON CONFLICT(train_number) DO UPDATE SET days = excluded.days
            ''', (train_number, days))
            train_id = con.execute(
                'SELECT id FROM trains WHERE train_number = ?', (train_number,)
            ).fetchone()[0]

            for t in e['times']:
                station = t['station']
                km = t.get('km')  # невідомо, можна None
                # Додаємо або оновлюємо станцію
                con.execute('''
                    INSERT INTO stations (name, km)
                    VALUES (?, ?)
                    ON CONFLICT(name) DO UPDATE SET km = COALESCE(excluded.km, stations.km)
                ''', (station, km))
                station_id = con.execute(
                    'SELECT id FROM stations WHERE name = ?', (station,)
                ).fetchone()[0]

                arrival = t['arrival']
                departure = t['departure']
                # Вставляємо або замінюємо розклад поїзда на станції
                con.execute('''
                    INSERT OR REPLACE INTO schedules
                    (train_id, station_id, arrival_time, departure_time, travel_date)
                    VALUES (?, ?, ?, ?, ?)
                ''', (train_id, station_id, arrival, departure, today))
        con.commit()

def get_schedule(direction: str, travel_date: Optional[str] = None) -> List[Dict]:
    """Повертає розклад поїздів за датою."""
    travel_date = travel_date or date.today().isoformat()
    with sqlite3.connect(DB_PATH) as con:
        rows = con.execute('''
            SELECT tr.train_number, st.name, sc.arrival_time, sc.departure_time
            FROM schedules sc
            JOIN trains tr ON sc.train_id = tr.id
            JOIN stations st ON sc.station_id = st.id
            WHERE sc.travel_date = ?
            ORDER BY tr.train_number, st.id
        ''', (travel_date,)).fetchall()

    schedule: Dict[str, List[Dict]] = {}
    for num, station, arrival, departure in rows:
        schedule.setdefault(num, []).append({
            'station': station,
            'arrival': arrival,
            'departure': departure
        })
    return [{'train_number': num, 'times': times} for num, times in schedule.items()]