import pandas as pd
from pathlib import Path

def convert_unified_duty(
    input_path,
    output_path="data/unified_timetable.parquet",
    line=1,
    diagram_type="weekday"
):
    # 1. 엑셀 로드 - 모든 시트 검사 (기본은 첫 번째 시트)
    raw = pd.read_excel(input_path, sheet_name=0, header=None)
    
    # 2. 헤더 및 데이터 위치 탐색
    # '기지' 컬럼이 있는 행을 헤더 행으로 판단
    header_row_idx = None
    for i in range(15):
        if "기지" in raw.iloc[i].astype(str).tolist():
            header_row_idx = i
            break
    
    if header_row_idx is None:
        raise ValueError("표의 헤더(기지)를 찾을 수 없습니다.")

    # 3. 데이터프레임 구조화
    header = raw.iloc[header_row_idx].astype(str).str.strip().tolist()
    df = raw.iloc[header_row_idx + 1:].reset_index(drop=True)
    
    # 중복 컬럼 이름을 고유하게 변경 (예: 열번, 열번_1, 열번_2...)
    new_cols = []
    counts = {}
    for col in header:
        col_name = col if col != "nan" else "unnamed"
        if col_name in counts:
            counts[col_name] += 1
            new_cols.append(f"{col_name}_{counts[col_name]}")
        else:
            counts[col_name] = 0
            new_cols.append(col_name)
    
    df.columns = new_cols

    # 4. 필수 컬럼 및 역 컬럼 자동 식별
    # 사용자가 말한 '열(1열, 2열...)'은 무조건 0번 컬럼
    track_col = new_cols[0]
    direction_col = next((c for c in new_cols if "상하" in c), None)
    run_col = next((c for c in new_cols if "입출고" in c), None)
    duty_col = next((c for c in new_cols if "DIA" in c or "다이아" in c), None)
    train_col = next((c for c in new_cols if "열번" in c), None)
    depot_col = next((c for c in new_cols if "기지" in c), None)

    # 역 컬럼 정의: '기지' 이후부터 끝까지 중, 메타 성격이 아닌 모든 컬럼
    # (단, 마지막 '열번_N'은 중복이므로 제외할 수도 있으나 안전하게 포함 후 나중에 불필요하면 제거)
    meta_keywords = ["상하", "입출고", "DIA", "다이아", "열번", "기지", "track_no", "unnamed"]
    station_cols = []
    col_found_start = False
    for i, col in enumerate(new_cols):
        # 기지 컬럼 이후부터 역 데이터 시작
        if col == depot_col:
            col_found_start = True
            continue
        if col_found_start:
            # 숫자로만 된 '열번_N' 등은 제외하고 실제 역 이름인 것만 수집
            is_meta = any(k in col for k in meta_keywords)
            if not is_meta:
                station_cols.append(col)

    print(f"📡 탐색된 역 ({len(station_cols)}개): {station_cols}")

    # 5. 데이터 클렌징
    # 열차번호는 필수 (원본 첫 번째 '열번' 컬럼 기준)
    df["train_number"] = pd.to_numeric(df[train_col], errors="coerce")
    df = df.dropna(subset=["train_number"])

    # 메타 데이터 변환
    direction_map = {"상": "up", "하": "down"}
    def classify_run_type(text):
        if pd.isna(text) or text == "nan": return "regular"
        text = str(text)
        if "출" in text: return "depot_out"
        elif "입" in text: return "depot_in"
        else: return "regular"

    df["direction"] = df[direction_col].map(direction_map).fillna("unknown")
    df["run_type"] = df[run_col].apply(classify_run_type)
    df["depot_detail"] = df[run_col].astype(str).replace("nan", "")
    df["duty_id"] = pd.to_numeric(df[duty_col], errors="coerce").fillna(0).astype(int)
    df["depot"] = df[depot_col].astype(str).replace("nan", "")
    df["track_no"] = pd.to_numeric(df[track_col], errors="coerce").fillna(0).astype(int)
    
    df["line"] = line
    df["diagram_type"] = diagram_type

    # 6. 세로형 변환 (Melt) - 모든 역 정보를 펼침
    long_df = df.melt(
        id_vars=[
            "line", "diagram_type", "duty_id", "train_number", 
            "direction", "run_type", "depot", "depot_detail", "track_no"
        ],
        value_vars=station_cols,
        var_name="station",
        value_name="time"
    )

    # 7. 후처리
    # 시간 없는 데이터 제거
    long_df = long_df.dropna(subset=["time"])
    
    def format_time(t):
        try:
            ts = pd.to_datetime(str(t))
            return ts.strftime("%H:%M:%S")
        except:
            return None

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

    # 역 순번 부여
    station_order = {s: i + 1 for i, s in enumerate(station_cols)}
    long_df["seq"] = long_df["station"].map(station_order)

    # 정렬
    long_df = long_df.sort_values(by=["duty_id", "train_number", "seq"]).reset_index(drop=True)

    # 최종 컬럼
    final_df = long_df[[
        "line", "diagram_type", "duty_id", "train_number", 
        "direction", "run_type", "depot", "depot_detail", "track_no", 
        "station", "seq", "time"
    ]]

    # 타입 고정
    final_df = final_df.astype({
        "line": "int8",
        "diagram_type": "string",
        "duty_id": "int16",
        "train_number": "int32",
        "direction": "string",
        "run_type": "string",
        "depot": "string",
        "depot_detail": "string",
        "track_no": "int16",
        "station": "string",
        "seq": "int16",
        "time": "string"
    })

    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    final_df.to_parquet(output_path, index=False, engine="pyarrow", compression="zstd")

    print(f"✅ 전체 데이터 파싱 완료: {output_path}")
    print(f"   - 총 데이터 건수: {len(final_df)} 건")
    print(f"   - 추출된 열차 대수: {df['train_number'].nunique()} 대")

if __name__ == "__main__":
    convert_unified_duty(
        input_path="sp_weekday.xlsx",
        output_path="data/unified_timetable.parquet",
        line=1,
        diagram_type="weekday"
    )