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 NOT NULL,
                days TEXT NOT NULL,
                route TEXT NOT NULL,
                UNIQUE(train_number, route)
            );
        ''')
        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(entries: List[Dict]):
    """
    Зберігає повний розклад без видалення попередніх записів.
    entries — список словників з полями:
      'train_number', 'days', 'route', 'times'.
    """
    today = date.today().isoformat()
    with sqlite3.connect(DB_PATH) as con:
        train_ids = []
        # Додаємо або оновлюємо поїзди
        for e in entries:
            tn = e['train_number']
            days = e['days']
            route = e.get('route', '')
            con.execute('''
                INSERT INTO trains (train_number, days, route)
                VALUES (?, ?, ?)
                ON CONFLICT(train_number, route) DO UPDATE SET days = excluded.days
            ''', (tn, days, route))
            tid = con.execute(
                'SELECT id FROM trains WHERE train_number = ? AND route = ?',
                (tn, route)
            ).fetchone()[0]
            train_ids.append(tid)

        # Додаємо або оновлюємо станції та розклад
        for idx, e in enumerate(entries):
            tid = train_ids[idx]
            for t in e['times']:
                st = t['station']
                km = t.get('km')
                con.execute('''
                    INSERT INTO stations (name, km)
                    VALUES (?, ?)
                    ON CONFLICT(name) DO UPDATE SET km = COALESCE(excluded.km, stations.km)
                ''', (st, km))
                sid = con.execute(
                    'SELECT id FROM stations WHERE name = ?', (st,)
                ).fetchone()[0]
                arr = t['arrival']
                dep = t['departure']
                con.execute('''
                    INSERT OR REPLACE INTO schedules
                    (train_id, station_id, arrival_time, departure_time, travel_date)
                    VALUES (?, ?, ?, ?, ?)
                ''', (tid, sid, arr, dep, today))
        con.commit()


def get_schedule(route: Optional[str] = None, travel_date: Optional[str] = None) -> List[Dict]:
    """Повертає розклад поїздів. Якщо вказано route, фільтрує за ним."""
    from datetime import date as _date
    travel_date = travel_date or _date.today().isoformat()
    with sqlite3.connect(DB_PATH) as con:
        if route:
            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 = ? AND tr.route = ?
                ORDER BY tr.train_number, st.id
            ''', (travel_date, route)).fetchall()
        else:
            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, rt, station, arrival, departure in rows:
        schedule.setdefault((num, rt), []).append({'station': station, 'arrival': arrival, 'departure': departure})
    return [{'train_number': num, 'route': rt, 'times': times} for (num, rt), times in schedule.items()]