import requests
from SRT import SRT
from SRT.errors import SRTLoginError, SRTNetFunnelError, SRTResponseError
from SRT.passenger import Adult, Child, Senior
from SRT.seat_type import SeatType as SRTSeatType

from koreantrain.base import Passenger, Reservation, SeatType, Train, TrainService
from koreantrain.errors import LoginError, NetworkError, ReservationError, SoldOutError

_PASSENGER_MAP = {"adult": Adult, "child": Child, "senior": Senior}

_SEAT_TYPE_MAP = {
    SeatType.GENERAL_FIRST: SRTSeatType.GENERAL_FIRST,
    SeatType.GENERAL_ONLY: SRTSeatType.GENERAL_ONLY,
    SeatType.SPECIAL_FIRST: SRTSeatType.SPECIAL_FIRST,
    SeatType.SPECIAL_ONLY: SRTSeatType.SPECIAL_ONLY,
}


def _convert_passengers(passengers):
    if not passengers:
        return None
    result = []
    for p in passengers:
        cls = _PASSENGER_MAP.get(p.type)
        if cls:
            result.append(cls(p.count))
    return result or None


def _to_train(raw):
    return Train(
        train_name=raw.train_name,
        train_number=raw.train_number,
        dep_station=raw.dep_station_name,
        arr_station=raw.arr_station_name,
        dep_date=raw.dep_date,
        dep_time=raw.dep_time,
        arr_date=raw.arr_date,
        arr_time=raw.arr_time,
        general_available="예약가능" in raw.general_seat_state,
        special_available="예약가능" in raw.special_seat_state,
        _raw=raw,
    )


def _to_reservation(raw):
    return Reservation(
        id=raw.reservation_number,
        total_cost=raw.total_cost,
        seat_count=raw.seat_count,
        payment_date=raw.payment_date,
        payment_time=raw.payment_time,
        paid=raw.paid,
        train_name=raw.train_name,
        train_number=raw.train_number,
        dep_station=raw.dep_station_name,
        arr_station=raw.arr_station_name,
        dep_date=raw.dep_date,
        dep_time=raw.dep_time,
        arr_time=raw.arr_time,
        _raw=raw,
    )


class SRTService(TrainService):
    def __init__(self, id, pw):
        try:
            self._srt = SRT(id, pw)
        except SRTLoginError as e:
            raise LoginError(str(e)) from e
        except (SRTNetFunnelError, requests.exceptions.ConnectionError) as e:
            raise NetworkError(repr(e)) from e

    def search(self, dep, arr, date=None, time=None, *, time_limit=None, passengers=None):
        try:
            trains = self._srt.search_train(dep, arr, date, time, time_limit=time_limit)
            return [_to_train(t) for t in trains]
        except SRTResponseError as e:
            if "조회 결과" in str(e):
                return []
            raise ReservationError(str(e)) from e
        except SRTLoginError as e:
            raise LoginError(str(e)) from e
        except (SRTNetFunnelError, requests.exceptions.ConnectionError) as e:
            raise NetworkError(repr(e)) from e

    def reserve(self, train, *, passengers=None, seat_type=SeatType.GENERAL_FIRST):
        try:
            raw = self._srt.reserve(
                train._raw,
                passengers=_convert_passengers(passengers),
                special_seat=_SEAT_TYPE_MAP.get(seat_type, SRTSeatType.GENERAL_FIRST),
            )
            return _to_reservation(raw)
        except SRTResponseError as e:
            if "잔여석없음" in str(e):
                raise SoldOutError(str(e)) from e
            raise ReservationError(str(e)) from e
        except SRTLoginError as e:
            raise LoginError(str(e)) from e
        except (SRTNetFunnelError, requests.exceptions.ConnectionError) as e:
            raise NetworkError(repr(e)) from e

    def reservations(self):
        try:
            return [_to_reservation(r) for r in self._srt.get_reservations()]
        except SRTLoginError as e:
            raise LoginError(str(e)) from e
        except (SRTNetFunnelError, requests.exceptions.ConnectionError) as e:
            raise NetworkError(repr(e)) from e

    def cancel(self, reservation):
        try:
            return self._srt.cancel(reservation._raw)
        except SRTResponseError as e:
            raise ReservationError(str(e)) from e
        except SRTLoginError as e:
            raise LoginError(str(e)) from e
        except (SRTNetFunnelError, requests.exceptions.ConnectionError) as e:
            raise NetworkError(repr(e)) from e
