# -*- coding: utf-8 -*-
r"""Step 4: 말뭉치 용례 추출 (네이버 뉴스 → 유튜브 2단계)

입력:
  - output/{corpus_type}/step3_candidates.csv
  - corpus parquets
출력:
  - output/{corpus_type}/step4_candidates_with_examples.csv
"""

from __future__ import annotations

import re
from pathlib import Path

import pandas as pd
import polars as pl
from ahocorasick_rs import AhoCorasick
from tqdm.auto import tqdm

from wordextractor.config import PipelineConfig

_EXPLAIN_PAT = re.compile(r"\([가-힣]+\+[가-힣]+\)")
MAX_EXAMPLES = 3


def _extract_context(
    text: str,
    start: int,
    end: int,
    sent_split: re.Pattern,
) -> str:
    sentences = []
    pos = 0
    for m in sent_split.finditer(text):
        seg = text[pos : m.start()]
        if seg.strip():
            sentences.append((pos, m.start(), seg.strip()))
        pos = m.end()
    tail = text[pos:]
    if tail.strip():
        sentences.append((pos, len(text), tail.strip()))

    if not sentences:
        return re.sub(r"[ \t]+", " ", text).strip()

    hit_idx = None
    for i, (s, e, _) in enumerate(sentences):
        if s <= start and end <= e:
            hit_idx = i
            break
    if hit_idx is None:
        for i, (s, e, _) in enumerate(sentences):
            if e > start and s < end:
                hit_idx = i
                break
    if hit_idx is None:
        hit_idx = 0

    lo = max(0, hit_idx - 1)
    hi = min(len(sentences) - 1, hit_idx + 1)
    parts = [re.sub(r"[ \t]+", " ", sentences[i][2]) for i in range(lo, hi + 1)]
    return " ".join(parts)


def _ac_pass(
    target_words: list[str],
    corpus_dir: Path,
    text_columns: list[str],
    sent_split: re.Pattern,
    months: list[str],
    max_n: int = MAX_EXAMPLES,
    sort_normal_by_length: bool = False,
    desc: str = "월별 용례 탐색",
) -> dict[str, dict]:
    ac = AhoCorasick(target_words)
    examples: dict[str, dict] = {w: {"explain": [], "normal": []} for w in target_words}
    done: set[str] = set()

    for ym in tqdm(months, desc=desc):
        path = corpus_dir / f"SC_{ym}.parquet"
        if not path.exists():
            continue

        df = pl.read_parquet(str(path), columns=text_columns)

        for col in text_columns:
            texts = df[col].cast(pl.Utf8).fill_null("").to_list()
            for text in texts:
                hits = ac.find_matches_as_indexes(text)
                if not hits:
                    continue

                for pat_idx, s, e in hits:
                    w = target_words[pat_idx]
                    if w in done:
                        continue

                    ex = examples[w]
                    if len(ex["explain"]) >= max_n:
                        done.add(w)
                        continue
                    if len(ex["explain"]) + len(ex["normal"]) >= max_n * 5:
                        continue

                    context = _extract_context(text, s, e, sent_split)

                    if _EXPLAIN_PAT.search(context):
                        ex["explain"].append(context)
                        if len(ex["explain"]) >= max_n:
                            done.add(w)
                    else:
                        ex["normal"].append(context)

        if len(done) == len(target_words):
            break

    if sort_normal_by_length:
        for w in target_words:
            examples[w]["normal"].sort(key=len, reverse=True)

    return examples


def _search_corpus(
    words: list[str],
    profile_dict: dict,
    months: list[str],
    max_n: int,
    sort_normal_by_length: bool = False,
    label: str = "",
) -> dict[str, dict]:
    corpus_dir = profile_dict["corpus_dir"]
    text_columns = profile_dict["text_columns"]
    sent_split = re.compile(profile_dict["sent_split"])

    print(f"\n  [{label}] AC 전체 매칭 ({len(words):,}개)...")
    ex = _ac_pass(
        words,
        corpus_dir,
        text_columns,
        sent_split,
        months,
        max_n,
        sort_normal_by_length,
        desc=f"[{label}] 월별 용례",
    )

    no_ex = [w for w in words if not ex[w]["explain"] and not ex[w]["normal"]]
    if no_ex:
        print(f"  [{label}] Pass 2 정밀 검색 ({len(no_ex):,}개)...")
        ex2 = _ac_pass(
            no_ex,
            corpus_dir,
            text_columns,
            sent_split,
            months,
            max_n,
            sort_normal_by_length,
            desc=f"[{label}] 정밀 검색",
        )
        for w in no_ex:
            ex[w] = ex2[w]

    has = sum(1 for w in words if ex[w]["explain"] or ex[w]["normal"])
    print(f"  [{label}] 용례 있음: {has:,}개, 없음: {len(words) - has:,}개")
    return ex


def collect_examples(
    words: list[str],
    corpus_profiles: dict,
    months: list[str],
    max_n: int = MAX_EXAMPLES,
) -> dict[str, list[str]]:
    """네이버 뉴스 → 유튜브 순서로 용례 수집 후 우선순위 병합"""

    # Resolve CorpusProfile objects to plain dicts for _search_corpus
    def _to_dict(profile):
        if hasattr(profile, "corpus_dir"):
            return {
                "corpus_dir": profile.corpus_dir,
                "text_columns": profile.text_columns,
                "sent_split": profile.sent_split,
            }
        return profile

    profiles = {k: _to_dict(v) for k, v in corpus_profiles.items()}

    # 1단계: 네이버 뉴스
    news_ex = _search_corpus(
        words,
        profiles["naver_news"],
        months,
        max_n,
        label="네이버 뉴스",
    )

    # 2단계: 유튜브 (뉴스에서 부족한 어절만)
    need_yt = [w for w in words if len(news_ex[w]["explain"]) + len(news_ex[w]["normal"]) < max_n]

    yt_ex: dict[str, dict] = {w: {"explain": [], "normal": []} for w in words}
    if need_yt and "youtube" in profiles:
        yt_raw = _search_corpus(
            need_yt,
            profiles["youtube"],
            months,
            max_n,
            sort_normal_by_length=True,
            label="유튜브",
        )
        for w in need_yt:
            yt_ex[w] = yt_raw[w]

    # 병합: 뉴스 화용 → 뉴스 일반 → 유튜브 화용 → 유튜브 일반
    result = {}
    for w in words:
        merged: list[str] = []
        for bucket in [
            news_ex[w]["explain"],
            news_ex[w]["normal"],
            yt_ex[w]["explain"],
            yt_ex[w]["normal"],
        ]:
            remaining = max_n - len(merged)
            if remaining <= 0:
                break
            merged.extend(bucket[:remaining])
        result[w] = merged

    return result


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

    print("step3 결과 로드...")
    candidates = pd.read_csv(out_dir / "step3_candidates.csv", encoding="utf-8-sig")
    print(f"  전체 후보: {len(candidates):,}개")

    words = candidates["어절"].tolist()

    print("\n용례 추출 중 (네이버 뉴스 → 유튜브)...")
    all_examples = collect_examples(words, cfg.corpus_profiles, cfg.months)

    for i in range(MAX_EXAMPLES):
        candidates[f"용례{i + 1}"] = ""

    no_example = []
    for idx, row in candidates.iterrows():
        w = row["어절"]
        exs = all_examples.get(w, [])
        if not exs:
            no_example.append(w)
        for i, ex in enumerate(exs):
            candidates.at[idx, f"용례{i + 1}"] = ex

    out_path = out_dir / "step4_candidates_with_examples.csv"
    candidates.to_csv(out_path, index=False, encoding="utf-8-sig")
    print(f"\n결과 저장: {out_path}")

    has_any = len(words) - len(no_example)
    has_explain = sum(
        1
        for w in words
        if all_examples.get(w) and any(_EXPLAIN_PAT.search(e) for e in all_examples[w])
    )
    print(f"\n=== 요약 ===")
    print(f"  전체 후보: {len(words):,}개")
    print(f"  용례 있음: {has_any:,}개")
    print(f"  용례 없음: {len(no_example):,}개")
    print(f"  설명 패턴 포함: {has_explain:,}개")

    if no_example:
        print(f"\n  용례 없는 어절: {no_example}")
