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,
                route TEXT
            );
        ''')
        # Таблиця станцій
        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]):
    """
    Зберігає повний розклад для заданого напряму.
    direction не використовується напряму — маршрут береться з поля 'route' у entries.
    entries — список словників з полями:
      {
        'train_number': str,
        'days': '1111111',
        'route': str,
        'times': [
          {'station': str, 'arrival': str, 'departure': str}, ...
        ]
      }
    """
    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']
            # Беріть маршрут із поля entry['route']
            route = e.get('route', '')
            # Додаємо або оновлюємо поїзд
            con.execute('''
                INSERT INTO trains (train_number, days, route)
                VALUES (?, ?, ?)
                ON CONFLICT(train_number) DO UPDATE SET
                    days = excluded.days,
                    route = excluded.route
            ''', (train_number, days, route))
            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 = None, travel_date: Optional[str] = None) -> List[Dict]:
    """
    Повертає розклад поїздів за датою.
    direction поки ігнорується; можна додати фільтрацію за route.
    """
    travel_date = travel_date or date.today().isoformat()
    with sqlite3.connect(DB_PATH) as con:
        rows = con.execute('''
            SELECT tr.train_number, tr.route, 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[tuple, List[Dict]] = {}
    for num, route, station, arrival, departure in rows:
        schedule.setdefault((num, route), []).append({
            'station': station,
            'arrival': arrival,
            'departure': departure
        })
    return [
        {'train_number': num, 'route': route, 'times': times}
        for (num, route), times in schedule.items()
    ]