import pandas as pd
from pathlib import Path


def detect_diagram_type(sheet_name: str) -> str:
    if "평일" in sheet_name:
        return "weekday"
    elif "토요일" in sheet_name:
        return "saturday"
    elif "휴일" in sheet_name:
        return "holiday"
    else:
        return None


def convert_sheet(df: pd.DataFrame, diagram_type: str) -> pd.DataFrame:
    """
    열번 시각표 변환 (최종 안정 버전)
    """

    # 첫 컬럼명을 train_number로 통일
    df = df.rename(columns={df.columns[0]: "train_number"})

    # 역 컬럼 목록
    stations = list(df.columns[1:])

    # 열차번호 숫자 변환
    df["train_number"] = pd.to_numeric(df["train_number"], errors="coerce")

    # 열차번호 없는 행 제거
    df = df.dropna(subset=["train_number"])

    # 시간 없는 열차 제거
    df = df.dropna(subset=stations, how="all")

    # 필요한 컬럼만 유지
    df = df[["train_number"] + stations]

    # 세로형 변환
    long_df = df.melt(
        id_vars=["train_number"],
        value_vars=stations,
        var_name="station",
        value_name="time"
    )

    # 시간 없는 행 제거
    long_df = long_df.dropna(subset=["time"])

    # 시간 문자열 통일
    long_df["time"] = pd.to_datetime(
        long_df["time"], errors="coerce"
    ).dt.strftime("%H:%M:%S")

    long_df = long_df.dropna(subset=["time"])

    # 순번 생성
    station_order = {station: i + 1 for i, station in enumerate(stations)}
    long_df["seq"] = long_df["station"].map(station_order)

    long_df["diagram_type"] = diagram_type

    return long_df.sort_values(["train_number", "seq"])


def convert_excel_to_parquet(xlsx_path: str, output_dir: str):
    xls = pd.ExcelFile(xlsx_path)
    all_data = []

    for sheet in xls.sheet_names:
        diagram_type = detect_diagram_type(sheet)
        if diagram_type is None:
            continue

        print(f"처리중: {sheet} → {diagram_type}")
        df = pd.read_excel(xls, sheet_name=sheet)
        converted = convert_sheet(df, diagram_type)
        all_data.append(converted)

    if not all_data:
        raise ValueError("변환할 시트를 찾지 못했습니다.")

    final_df = pd.concat(all_data, ignore_index=True)

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # partition 저장
    final_df.to_parquet(
        output_path,
        index=False,
        partition_cols=["diagram_type"]
    )

    print("\n변환 완료")
    print(f"저장 위치: {output_path}")
    print(final_df.head())


if __name__ == "__main__":
    convert_excel_to_parquet(
        xlsx_path="timetable.xlsx",
        output_dir="timetable_parquet"
    )