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,
                tid TEXT UNIQUE NOT NULL,
                train_number TEXT NOT NULL,
                days TEXT NOT NULL,
                route TEXT NOT NULL
            );
        ''')
        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 — список словників:
      {
        'tid': str,
        '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:
        train_ids = []
        # Додаємо або оновлюємо trains по tid
        for e in entries:
            tid_val = e['tid']
            tn = e['train_number']
            days = e['days']
            route = e['route']
            con.execute('''
                INSERT INTO trains (tid, train_number, days, route)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(tid) DO UPDATE SET
                    train_number = excluded.train_number,
                    days         = excluded.days,
                    route        = excluded.route
            ''', (tid_val, tn, days, route))
            train_id = con.execute(
                'SELECT id FROM trains WHERE tid = ?',
                (tid_val,)
            ).fetchone()[0]
            train_ids.append(train_id)

        # Вставляємо записи schedules
        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))
                station_id = 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, station_id, arr, dep, today))
        con.commit()

def get_schedule(route: Optional[str] = None, travel_date: Optional[str] = None) -> List[Dict]:
    """
    Повертає розклад поїздів. Якщо вказано route, фільтрує за ним.
    Повертає список:
      [{'train_number': ..., 'route': ..., 'times': [...]}, ...]
    """
    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()
    ]