import requests
from bs4 import BeautifulSoup
from typing import List, Dict

BASE_URL = 'https://swrailway.gov.ua/timetable/eltrain/?rid=2'
LOCAL_HTML = 'rozklad.html'

# Індекс днів тижня: понеділок=0, ..., неділя=6
DAY_INDEX = {
    'пн': 0, 'вт': 1, 'ср': 2,
    'чт': 3, 'пт': 4, 'сб': 5,
    'нд': 6
}

def parse_days(text: str) -> str:
    """Перетворює текст днів курсування на бінарний рядок довжиною 7."""
    text = text.lower().strip()
    if 'щоденно' in text:
        return '1111111'
    if text.startswith('крім'):
        days = [d.strip(' .') for d in text.split('крім', 1)[1].split(',')]
        mask = [1]*7
        for d in days:
            idx = DAY_INDEX.get(d)
            if idx is not None:
                mask[idx] = 0
        return ''.join(str(b) for b in mask)
    if text.startswith('по'):
        days = [d.strip(' .') for d in text.split('по', 1)[1].split(',')]
        mask = [0]*7
        for d in days:
            idx = DAY_INDEX.get(d)
            if idx is not None:
                mask[idx] = 1
        return ''.join(str(b) for b in mask)
    return '1111111'

def fetch_schedule(tab: int = 1, use_local: bool = False) -> List[Dict]:
    """
    Парсить вкладку розкладу:
      tab=1 — Київ→Ніжин,
      tab=2 — Ніжин→Київ.
    Повертає список поїздів з полями:
      'train_number', 'days', 'route', 'times' (список словників station/arrival/departure).
    """
    # Завантаження HTML
    if use_local:
        with open(LOCAL_HTML, 'r', encoding='utf-8') as f:
            html = f.read()
    else:
        resp = requests.get(BASE_URL, timeout=10)
        resp.raise_for_status()
        html = resp.text

    soup = BeautifulSoup(html, 'html.parser')
    prefix = f'div#tabs-trains{tab}'

    # Таблиця з розкладом
    times_table = soup.select_one(f'{prefix} table.td_center')
    if not times_table:
        raise RuntimeError(f'Не знайдено таблицю розкладу для tab={tab}')

    # Парсимо маршрути (по одному <td class="course"> на потяг)
    route_tags = times_table.select('td.course')
    routes = [tag.get_text(strip=True) for tag in route_tags]

    # Список станцій (35)
    station_tags = soup.select(
        f'{prefix} table.left tr.on a.et, '
        f'{prefix} table.left tr.onx a.et'
    )
    stations = [a.get_text(strip=True) for a in station_tags]

    # Заголовок з номерами потягів і днями курсування
    trs = times_table.find_all('tr')
    header_row = next(r for r in trs if r.find('td', class_='on_right_t'))
    cells = header_row.find_all('td', class_='on_right_t')

    entries: List[Dict] = []
    for idx, cell in enumerate(cells):
        parts = cell.get_text(separator='|', strip=True).split('|')
        num = parts[0].rstrip(',').strip()
        days = parse_days(parts[1] if len(parts) > 1 else 'щоденно')
        route = routes[idx] if idx < len(routes) else ''
        entries.append({
            'train_number': num,
            'days': days,
            'route': route,
            'times': []
        })

    # Рядки з часами руху
    time_rows = [r for r in trs if r.find('td', class_='q0') or r.find('td', class_='q1')]
    for idx, entry in enumerate(entries):
        base = idx * 3
        for si, row in enumerate(time_rows):
            cells = row.find_all('td')
            arrival = cells[base+1].get_text(strip=True) if base+1 < len(cells) else ''
            departure = cells[base+2].get_text(strip=True) if base+2 < len(cells) else ''
            entry['times'].append({
                'station': stations[si],
                'arrival': arrival,
                'departure': departure
            })

    return entries