# -*- coding: utf-8 -*-
"""Step 2: 말뭉치(01~12월)에서 어절 빈도 목록 구축

입력: corpus_dir/SC_YYYYMM.parquet (12개)
출력: output/{corpus_type}/step2_eojeol_type_freq.csv
"""

from __future__ import annotations

import re
from collections import Counter

import polars as pl
from tqdm.auto import tqdm

from wordextractor.config import PipelineConfig

_KO = re.compile(r"[가-힣]+")


def extract_words_from_parquet(path: str, text_columns: list[str]) -> Counter:
    df = pl.read_parquet(path, columns=text_columns)
    counter: Counter = Counter()

    for col in text_columns:
        for text in df[col].cast(pl.Utf8).fill_null("").to_list():
            counter.update(_KO.findall(text))

    return counter


def run(cfg: PipelineConfig) -> None:
    out_dir = cfg.ensure_output_dir()

    total: Counter = Counter()
    for ym in tqdm(cfg.months, desc="월별 말뭉치 처리"):
        path = cfg.corpus_dir / f"SC_{ym}.parquet"
        if not path.exists():
            print(f"  [스킵] {path} 없음")
            continue

        c = extract_words_from_parquet(str(path), cfg.corpus_text_columns)
        total.update(c)
        print(f"  {ym}: 고유 어절 {len(c):,}개, 누적 총 {sum(total.values()):,}개")

    print(f"\n전체 고유 어절: {len(total):,}개")
    print(f"전체 총 어절 수: {sum(total.values()):,}개")

    merged = pl.DataFrame({"어절": list(total.keys()), "빈도": list(total.values())}).sort(
        "빈도", descending=True
    )

    out_path = out_dir / "step2_eojeol_type_freq.csv"
    merged.to_pandas().to_csv(out_path, index=False, encoding="utf-8-sig")
    print(f"\n어절 빈도 저장: {out_path}")
