# -*- coding: utf-8 -*-
"""wordextractor public API

사용 예:
    from wordextractor import (
        extract_patterns,
        extract_matched_eojeol,
        search_corpus_example,
        search_naver_first_occurrence,
    )
"""

from __future__ import annotations

import re
from collections import Counter, defaultdict
from statistics import mean, stdev
from typing import Union

import pandas as pd
from tqdm.auto import tqdm

from wordextractor.utils.korean import keep_korean, is_korean_only, KO_PATTERN


# ═══════════════════════════════════════════════════════════
# 1) extract_patterns
# ═══════════════════════════════════════════════════════════


def extract_patterns(
    words: Union[list[str], pd.Series],
    *,
    min_syllable: int = 3,
    ngram_rules: dict[int, int] | None = None,
    ngram_default: int = 4,
    min_freq: int = 2,
    tail_only: bool = True,
) -> pd.DataFrame:
    """기등재 혼성어 목록에서 N-Gram 패턴을 추출한다.

    혼성어의 뒷부분(어미 위치)에서 반복적으로 나타나는 n-gram을 찾아
    새로운 혼성어 후보를 탐색할 패턴으로 활용한다.
    예: "호캉스", "펫캉스" → 패턴 "캉스" (빈도 2)

    Parameters
    ----------
    words : list[str] or pd.Series
        혼성어 어형 목록 (예: ["호캉스", "펫캉스", "브렉시트"])
    min_syllable : int
        최소 음절 수 (기본 3). 이보다 짧은 단어는 무시.
    ngram_rules : dict[int, int] | None
        음절 수 → n-gram 크기 매핑 (기본 {3: 2, 4: 3}).
        예: 3음절 단어에서 2-gram, 4음절 단어에서 3-gram 추출.
    ngram_default : int
        ngram_rules에 없는 음절 수의 기본 n값 (기본 4)
    min_freq : int
        최소 패턴 빈도 (기본 2). 이 빈도 이상 등장한 패턴만 반환.
    tail_only : bool
        True이면 어미 위치(위칫값 평균 == 1.0) 패턴만 반환 (기본 True).
        혼성어는 보통 뒷부분이 공유되므로 tail_only=True 권장.

    Returns
    -------
    pd.DataFrame
        columns: 패턴, 패턴 길이, 빈도, 패턴 위칫값 평균, 패턴 위칫값의 표준편차, 혼성어 목록

    Examples
    --------
    >>> from wordextractor import extract_patterns
    >>> patterns = extract_patterns(["호캉스", "펫캉스", "알캉스", "브렉시트"])
    >>> print(patterns[["패턴", "빈도"]])
       패턴  빈도
    0  캉스     3
    """
    if ngram_rules is None:
        ngram_rules = {3: 2, 4: 3}

    if isinstance(words, pd.Series):
        words = words.dropna().astype(str).tolist()

    positions: dict[str, list[float]] = defaultdict(list)
    sources: dict[str, set[str]] = defaultdict(set)

    for headword in tqdm(words, desc="n-gram 패턴 추출"):
        clean = keep_korean(headword)
        if not clean:
            continue
        syl_len = len(clean)
        if syl_len < min_syllable:
            continue
        n = ngram_rules.get(syl_len, ngram_default)
        if syl_len < n:
            continue

        span = syl_len - n
        for i in range(span + 1):
            pat = clean[i : i + n]
            positions[pat].append(i / span if span > 0 else 0.0)
            sources[pat].add(clean)

    rows = []
    for pat in positions:
        src_set = sources[pat]
        freq = len(src_set)
        pos_list = positions[pat]
        avg_pos = mean(pos_list)
        std_pos = stdev(pos_list) if len(pos_list) >= 2 else 0.0
        src = "; ".join(sorted(src_set))
        rows.append((pat, len(pat), freq, round(avg_pos, 4), round(std_pos, 4), src))

    df = (
        pd.DataFrame(
            rows,
            columns=["패턴", "패턴 길이", "빈도", "패턴 위칫값 평균", "패턴 위칫값의 표준편차", "혼성어 목록"],
        )
        .sort_values(["패턴 길이", "빈도"], ascending=[True, False])
        .reset_index(drop=True)
    )

    # 필터링
    mask = df["빈도"] >= min_freq
    if tail_only:
        mask = mask & (df["패턴 위칫값 평균"] == 1.0)

    return df[mask].reset_index(drop=True)


# ═══════════════════════════════════════════════════════════
# 2) extract_matched_eojeol
# ═══════════════════════════════════════════════════════════


def extract_matched_eojeol(
    corpus: Union[list[str], pd.Series, pd.DataFrame],
    patterns: Union[list[str], pd.DataFrame],
    *,
    text_columns: list[str] | None = None,
    min_syllable: int = 3,
    korean_only: bool = True,
    exclude_urimalsam: bool = True,
    exclude_input_words: bool = True,
) -> pd.DataFrame:
    """말뭉치에서 패턴에 매칭되는 어절과 빈도를 추출한다.

    extract_patterns()로 얻은 패턴을 사용하여 말뭉치에서 해당 패턴을
    포함하는 어절을 찾고, 빈도와 함께 반환한다.

    Parameters
    ----------
    corpus : list[str] | pd.Series | pd.DataFrame
        말뭉치. 문장(문자열) 리스트, Series, 또는 DataFrame.
        DataFrame인 경우 text_columns 지정 필요.
    patterns : list[str] | pd.DataFrame
        패턴 리스트, 또는 extract_patterns()의 반환 DataFrame ("패턴" 컬럼 사용)
    text_columns : list[str] | None
        corpus가 DataFrame일 때 텍스트 컬럼명 (기본: 전체 str 컬럼)
    min_syllable : int
        최소 음절 수 (기본 3)
    korean_only : bool
        한글 전용 어절만 (기본 True)
    exclude_urimalsam : bool
        True이면 내장 우리말샘 표제어(약 97만개)에 포함된 어절을 자동 제외 (기본 True).
    exclude_input_words : bool
        True이면 패턴 추출에 사용된 입력 혼성어 목록을 자동 제외 (기본 True).
        patterns가 extract_patterns()의 반환 DataFrame일 때만 동작.

    Returns
    -------
    pd.DataFrame
        columns: 어절, 빈도, 매칭 패턴

    Examples
    --------
    >>> from wordextractor import extract_patterns, extract_matched_eojeol
    >>> patterns = extract_patterns(blend_words)
    >>> sentences = ["요즘 슬세권이 대세다", "맥세권 아파트 인기"]
    >>> matched = extract_matched_eojeol(sentences, patterns)

    >>> # 우리말샘 제외 끄기
    >>> matched = extract_matched_eojeol(sentences, patterns, exclude_urimalsam=False)

    >>> # 입력 혼성어 제외 끄기
    >>> matched = extract_matched_eojeol(sentences, patterns, exclude_input_words=False)
    """
    from ahocorasick_rs import AhoCorasick

    # patterns 처리
    if isinstance(patterns, pd.DataFrame):
        pat_list = patterns["패턴"].tolist()
    else:
        pat_list = list(patterns)

    # 우리말샘 내장 사전 로드
    urimalsam_set: frozenset[str] = frozenset()
    if exclude_urimalsam:
        from wordextractor.data import load_urimalsam_headwords
        urimalsam_set = load_urimalsam_headwords()
        print(f"  내장 우리말샘 표제어 로드: {len(urimalsam_set):,}개")

    # 입력 혼성어 목록 추출
    input_words: set[str] = set()
    if exclude_input_words and isinstance(patterns, pd.DataFrame) and "혼성어 목록" in patterns.columns:
        for s in patterns["혼성어 목록"].dropna():
            for w in str(s).split(";"):
                w = w.strip()
                if w:
                    input_words.add(w)
        if input_words:
            print(f"  입력 혼성어 제외 목록: {len(input_words):,}개")

    # corpus → 문장 리스트
    sentences: list[str] = _corpus_to_sentences(corpus, text_columns)

    # 어절 빈도 집계
    ko_re = re.compile(r"[가-힣]+")
    word_counter: Counter = Counter()
    for sent in sentences:
        word_counter.update(ko_re.findall(sent))

    # Aho-Corasick 매칭
    ac = AhoCorasick(pat_list)
    results: dict[str, set[str]] = {}
    n_excluded_input = 0
    n_excluded_urimalsam = 0

    for word in tqdm(word_counter, desc="패턴 매칭"):
        if korean_only and not is_korean_only(word):
            continue
        if len(word) < min_syllable:
            continue

        # 우리말샘 등재어 제외
        if urimalsam_set and word in urimalsam_set:
            n_excluded_urimalsam += 1
            continue

        # 입력 혼성어 제외
        if input_words and word in input_words:
            n_excluded_input += 1
            continue

        hits = ac.find_matches_as_indexes(word)
        if hits:
            matched = sorted(set(pat_list[idx] for idx, _s, _e in hits))
            results[word] = set(matched)

    rows = []
    for word, pats in results.items():
        rows.append((word, word_counter[word], "; ".join(sorted(pats))))

    df = pd.DataFrame(rows, columns=["어절", "빈도", "매칭 패턴"])
    df = df.sort_values("빈도", ascending=False).reset_index(drop=True)

    print(f"  매칭 어절: {len(df):,}개")
    if urimalsam_set:
        print(f"  우리말샘 제외: {n_excluded_urimalsam:,}개")
    if input_words:
        print(f"  입력 혼성어 제외: {n_excluded_input:,}개")

    return df


def load_dictionary(
    path: str,
    *,
    column: str = "어휘",
    min_length: int = 2,
    glob_pattern: str | None = None,
) -> set[str]:
    """사전 파일(xlsx/xls)에서 한글 표제어를 로드한다.

    Parameters
    ----------
    path : str
        파일 경로 또는 디렉토리 경로
    column : str
        표제어 컬럼명 (기본 "어휘")
    min_length : int
        최소 글자 수 (기본 2)
    glob_pattern : str | None
        디렉토리일 때 파일 패턴 (기본 "*.xls")

    Returns
    -------
    set[str]
        한글 전용 표제어 집합

    Examples
    --------
    >>> # 우리말샘 디렉토리
    >>> urimal = load_dictionary("./우리말샘_xls/", column="어휘")

    >>> # 혼성어 목록 파일
    >>> blends = load_dictionary("wordlist.xlsx", column="혼성어(색인표제어)")

    >>> # extract_matched_eojeol에서 사용
    >>> matched = extract_matched_eojeol(
    ...     corpus, patterns,
    ...     exclude_words=blends,
    ...     exclude_if_contains=urimal,
    ... )
    """
    from pathlib import Path

    p = Path(path)
    headwords: set[str] = set()

    if p.is_dir():
        pattern = glob_pattern or "*.xls"
        files = sorted(p.glob(pattern))
        print(f"  사전 파일: {len(files)}개 ({p})")

        engine = "xlrd" if pattern.endswith(".xls") else None
        for fp in tqdm(files, desc="사전 로드"):
            raw_df = pd.read_excel(fp, usecols=[column], engine=engine)
            for raw in raw_df[column].dropna().astype(str):
                clean = keep_korean(raw)
                if len(clean) >= min_length:
                    headwords.add(clean)
    else:
        raw_df = pd.read_excel(p, usecols=[column])
        for raw in raw_df[column].dropna().astype(str):
            clean = keep_korean(raw)
            if len(clean) >= min_length:
                headwords.add(clean)

    print(f"  표제어 로드 완료: {len(headwords):,}개")
    return headwords


# ═══════════════════════════════════════════════════════════
# 3) search_corpus_example
# ═══════════════════════════════════════════════════════════

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


def search_corpus_example(
    words: list[str],
    corpus: Union[list[str], pd.Series, pd.DataFrame],
    *,
    text_columns: list[str] | None = None,
    max_examples: int = 3,
) -> pd.DataFrame:
    """후보 어절의 말뭉치 용례를 검색한다.

    각 어절이 실제 말뭉치에서 어떤 문맥으로 사용되었는지 확인하기 위한 함수.
    "(한글+한글)" 형태의 화용표지(설명 구문)가 포함된 용례를 우선 반환한다.

    Parameters
    ----------
    words : list[str]
        용례를 검색할 어절 목록
    corpus : list[str] | pd.Series | pd.DataFrame
        말뭉치 (문장 리스트, Series, 또는 DataFrame)
    text_columns : list[str] | None
        corpus가 DataFrame일 때 텍스트 컬럼명
    max_examples : int
        어절당 최대 용례 수 (기본 3)

    Returns
    -------
    pd.DataFrame
        columns: 어절, 용례1, 용례2, ..., 용례N, 화용표지포함

    Examples
    --------
    >>> from wordextractor import search_corpus_example
    >>> sentences = ["슬세권 아파트가 인기다", "슬세권(슬리퍼+역세권) 열풍"]
    >>> result = search_corpus_example(["슬세권"], sentences)
    >>> print(result.columns.tolist())
    ['어절', '용례1', '용례2', '용례3', '화용표지포함']
    """
    from ahocorasick_rs import AhoCorasick

    sentences = _corpus_to_sentences(corpus, text_columns)

    ac = AhoCorasick(words)
    examples: dict[str, dict] = {w: {"explain": [], "normal": []} for w in words}
    done: set[str] = set()

    for sent in tqdm(sentences, desc="용례 탐색"):
        hits = ac.find_matches_as_indexes(sent)
        if not hits:
            continue

        for pat_idx, _s, _e in hits:
            w = words[pat_idx]
            if w in done:
                continue

            ex = examples[w]
            if len(ex["explain"]) >= max_examples:
                done.add(w)
                continue

            context = sent.strip()

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

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

    # 결과 조합: explain 우선 → normal
    rows = []
    for w in words:
        ex = examples[w]
        merged: list[str] = []
        for bucket in [ex["explain"], ex["normal"]]:
            remaining = max_examples - len(merged)
            if remaining <= 0:
                break
            merged.extend(bucket[:remaining])

        has_explain = any(_EXPLAIN_PAT.search(e) for e in merged) if merged else False
        row = {"어절": w}
        for i in range(max_examples):
            row[f"용례{i + 1}"] = merged[i] if i < len(merged) else ""
        row["화용표지포함"] = has_explain
        rows.append(row)

    return pd.DataFrame(rows)


# ═══════════════════════════════════════════════════════════
# 4) search_naver_first_occurrence
# ═══════════════════════════════════════════════════════════


def _is_colab() -> bool:
    try:
        import google.colab  # noqa: F401
        return True
    except ImportError:
        return False


def _setup_colab_chrome():
    """Colab 환경에서 Chromium + ChromeDriver 자동 설치"""
    import subprocess
    import shutil

    if shutil.which("chromium-browser") or shutil.which("chromium"):
        return  # 이미 설치됨

    print("Colab 환경 감지 — Chromium 설치 중...")
    subprocess.run(
        ["apt-get", "update", "-qq"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    subprocess.run(
        ["apt-get", "install", "-y", "-qq", "chromium-browser"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    print("  Chromium 설치 완료")


def _create_driver_auto():
    """환경에 맞는 Selenium WebDriver 생성"""
    try:
        from selenium import webdriver
        from selenium.webdriver.chrome.options import Options
        from selenium.webdriver.chrome.service import Service
    except ImportError:
        raise ImportError(
            "search_naver_first_occurrence requires selenium. "
            "Install with: pip install wordextractor[naver]"
        )

    import subprocess

    if _is_colab():
        _setup_colab_chrome()

    options = Options()
    options.add_argument("--headless")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    options.add_argument("--disable-blink-features=AutomationControlled")
    options.add_argument("--disable-gpu")
    options.add_argument("--log-level=3")
    options.add_argument("--disable-extensions")
    options.add_argument(
        "user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
    )

    if _is_colab():
        options.binary_location = "/usr/bin/chromium-browser"

    service = Service(log_output=subprocess.DEVNULL)
    return webdriver.Chrome(service=service, options=options)


_DATE_RE = re.compile(r"\d{4}\.\d{2}\.\d{2}\.?")


def _search_single_word(driver, word: str) -> dict:
    """단일 어절 네이버 뉴스 검색"""
    from urllib.parse import quote
    from bs4 import BeautifulSoup

    search_q = quote(f'"{word}"')
    url = (
        "https://search.naver.com/search.naver?"
        f"ssc=tab.news.all&query={search_q}&sm=tab_opt&sort=2"
        "&photo=0&field=0&pd=0&qdt=1&related=0&mynews=0"
        "&office_type=0&office_section_code=0&news_office_checked="
        "&nso=so%3Ar%2Cp%3Aall&is_sug_officeid=0&office_category=0&service_area=0"
    )

    result = {
        "어절": word,
        "검색URL": url,
        "검색상태": "",
        "네이버뉴스_최초일자": "",
        "네이버뉴스_언론사": "",
        "기타뉴스_최초일자": "",
        "기타뉴스_언론사": "",
    }

    import time

    try:
        driver.get(url)
        time.sleep(0.5)

        soup = BeautifulSoup(driver.page_source, "html.parser")

        not_found = soup.select(".not_found02")
        if any("검색결과가없습니다" in el.get_text().replace(" ", "") for el in not_found):
            result["검색상태"] = "검색 결과 없음"
            return result

        profiles = soup.select("div.sds-comps-profile")
        if not profiles:
            result["검색상태"] = "파싱 실패"
            return result

        result["검색상태"] = "성공"

        for profile in profiles:
            is_naver = any("네이버뉴스" in a.get_text() for a in profile.select("a"))

            press = ""
            press_el = profile.select_one(".sds-comps-profile-info-title-text span")
            if press_el:
                press = press_el.get_text(strip=True)

            # 날짜 추출
            date = ""
            for st in profile.select(".sds-comps-profile-info-subtext"):
                txt = st.get_text(strip=True)
                if _DATE_RE.fullmatch(txt):
                    date = txt
                    break
                for span in st.select("span"):
                    txt = span.get_text(strip=True)
                    if _DATE_RE.fullmatch(txt):
                        date = txt
                        break
                if date:
                    break
            if not date:
                m = _DATE_RE.search(profile.get_text())
                if m:
                    date = m.group()

            if is_naver and not result["네이버뉴스_최초일자"]:
                result["네이버뉴스_최초일자"] = date
                result["네이버뉴스_언론사"] = press
            elif not is_naver and not result["기타뉴스_최초일자"]:
                result["기타뉴스_최초일자"] = date
                result["기타뉴스_언론사"] = press

            if result["네이버뉴스_최초일자"] and result["기타뉴스_최초일자"]:
                break

    except Exception as e:
        result["검색상태"] = f"실패: {e}"

    return result


def search_naver_first_occurrence(
    words: Union[list[str], pd.Series],
    *,
    delay: float = 1.0,
    n_workers: int | None = None,
    restart_every: int = 50,
    verbose: bool = True,
) -> pd.DataFrame:
    """어절 목록의 네이버 뉴스 최초 출현일을 검색한다.

    Selenium으로 네이버 뉴스를 날짜순 정렬 검색하여 각 어절의
    최초 출현일과 언론사 정보를 수집한다.
    Google Colab과 PC 환경을 자동 감지하여 Chrome/Chromium을 설정한다.

    - Colab: Chromium을 자동 설치하고 단일 워커로 실행
    - PC: 병렬 워커 지원 (기본 4개)

    Parameters
    ----------
    words : list[str] | pd.Series
        검색할 어절 목록
    delay : float
        요청 간 대기 시간(초) (기본 1.0).
        너무 짧으면 차단될 수 있으므로 1.0 이상 권장.
    n_workers : int | None
        병렬 워커 수. None이면 자동 결정 (Colab: 1, PC: 4).
    restart_every : int
        N건마다 드라이버 재시작 (기본 50).
        메모리 누수 방지용.
    verbose : bool
        진행 상황 출력 (기본 True)

    Returns
    -------
    pd.DataFrame
        columns: 어절, 검색URL, 검색상태, 네이버뉴스_최초일자,
                 네이버뉴스_언론사, 기타뉴스_최초일자, 기타뉴스_언론사

    Examples
    --------
    >>> from wordextractor import search_naver_first_occurrence
    >>> result = search_naver_first_occurrence(["슬세권", "맥세권"])
    >>> print(result[["어절", "네이버뉴스_최초일자"]])
    """
    import time

    if isinstance(words, pd.Series):
        words = words.dropna().astype(str).tolist()

    words = list(dict.fromkeys(words))  # 중복 제거, 순서 유지

    if n_workers is None:
        n_workers = 1 if _is_colab() else 4

    if verbose:
        print(f"네이버 뉴스 최초 출현일 검색")
        print(f"  대상: {len(words):,}개")
        print(f"  워커: {n_workers}개")
        print(f"  딜레이: {delay}초")
        env = "Colab" if _is_colab() else "PC"
        print(f"  환경: {env}")

    # ── 단일 워커 (Colab 또는 n_workers==1) ──
    if n_workers <= 1:
        driver = _create_driver_auto()
        results = []
        request_count = 0

        try:
            for word in tqdm(words, desc="네이버 검색", disable=not verbose):
                result = _search_single_word(driver, word)
                results.append(result)
                request_count += 1

                if restart_every > 0 and request_count >= restart_every:
                    driver.quit()
                    time.sleep(2)
                    driver = _create_driver_auto()
                    request_count = 0

                time.sleep(delay)
        finally:
            driver.quit()

        return pd.DataFrame(results)

    # ── 멀티 워커 (PC) ──
    from multiprocessing import Process, Manager

    manager = Manager()
    results_dict = manager.dict()

    chunk_size = max(1, len(words) // n_workers)
    chunks = []
    for i in range(n_workers):
        start = i * chunk_size
        end = start + chunk_size if i < n_workers - 1 else len(words)
        chunk = words[start:end]
        if chunk:
            chunks.append(chunk)

    def _worker(worker_id, chunk, shared_dict, delay_, restart_every_):
        driver = _create_driver_auto()
        request_count = 0
        try:
            for word in tqdm(chunk, desc=f"Worker {worker_id}", position=worker_id):
                result = _search_single_word(driver, word)
                shared_dict[word] = result
                request_count += 1

                if restart_every_ > 0 and request_count >= restart_every_:
                    driver.quit()
                    time.sleep(2)
                    driver = _create_driver_auto()
                    request_count = 0

                time.sleep(delay_)
        finally:
            driver.quit()

    processes = []
    for i, chunk in enumerate(chunks):
        p = Process(target=_worker, args=(i, chunk, results_dict, delay, restart_every))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    # 순서 유지
    results = []
    for w in words:
        if w in results_dict:
            results.append(dict(results_dict[w]))
        else:
            results.append({"어절": w, "검색상태": "미검색"})

    df = pd.DataFrame(results)

    if verbose:
        success = (df["검색상태"] == "성공").sum()
        no_result = (df["검색상태"] == "검색 결과 없음").sum()
        has_naver = (df["네이버뉴스_최초일자"] != "").sum() if "네이버뉴스_최초일자" in df.columns else 0
        print(f"\n완료: 성공 {success:,} | 결과없음 {no_result:,} | 네이버뉴스 {has_naver:,}")

    return df


# ═══════════════════════════════════════════════════════════
# 내부 유틸
# ═══════════════════════════════════════════════════════════


def _corpus_to_sentences(
    corpus: Union[list[str], pd.Series, pd.DataFrame],
    text_columns: list[str] | None = None,
) -> list[str]:
    """다양한 입력을 문장 리스트로 변환"""
    if isinstance(corpus, pd.DataFrame):
        if text_columns is None:
            text_columns = [c for c in corpus.columns if corpus[c].dtype == object]
        sentences = []
        for col in text_columns:
            sentences.extend(corpus[col].dropna().astype(str).tolist())
        return sentences
    elif isinstance(corpus, pd.Series):
        return corpus.dropna().astype(str).tolist()
    else:
        return list(corpus)
