"""MDX ↔ Figma frame matching — 공통 유틸리티"""
import os
import re
import yaml
from pathlib import Path

ROOT = Path(r"d:\ad-hoc\kei\design_agent")
BLOCKS_DIR = ROOT / "figma_to_html_agent" / "blocks"
MDX_DIR = ROOT / "samples" / "mdx_batch"
GT_PATH = ROOT / "tests" / "matching" / "ground_truth.yaml"


def load_ground_truth():
    with open(GT_PATH, encoding="utf-8") as f:
        return yaml.safe_load(f)["sections"]


def _strip_figma_meta(text):
    """texts.md의 메타 헤더/구조 라벨 제거 + 타이틀 3회 반복(가중치).
    - '# Frame 1171281XXX — 텍스트 목록' 파일 최상위 제목 제거
    - '> ...' blockquote 메타 제거
    - '## 타이틀', '## 서브헤더', '## 열1', '### 세로 라벨' 구조 라벨 제거
    - '## 타이틀' 아래 bullet 또는 plain 라인을 3회 반복
    """
    lines = text.split("\n")
    cleaned = []
    title_lines = []
    in_title_section = False
    for ln in lines:
        if re.match(r"^#\s+Frame\s+\d+", ln):
            continue
        if ln.startswith("> "):
            continue
        # ## 타이틀 시작
        if re.match(r"^##\s+타이틀", ln):
            in_title_section = True
            continue
        # 다른 ## 나 ### 헤더 → 타이틀 구역 종료 (헤더 자체는 제거)
        if re.match(r"^#{2,}\s", ln):
            in_title_section = False
            continue
        # 타이틀 구역 안의 비어있지 않은 라인은 수집 (bullet이든 plain이든)
        if in_title_section:
            s = ln.strip()
            if s:
                # 앞의 '- ' 제거해서 깔끔하게
                clean_s = re.sub(r"^-\s*", "", s)
                if clean_s:
                    title_lines.append(clean_s)
            cleaned.append(ln)
            continue
        cleaned.append(ln)
    out = "\n".join(cleaned).strip()
    if title_lines:
        boost = "\n".join(title_lines)
        out = boost + "\n" + boost + "\n" + out
    return out


def load_figma_texts(clean_meta=True):
    """32개 프레임의 texts.md 읽기. {frame_id: text}
    clean_meta=True: 메타 헤더/라벨 제거 + 타이틀 가중치 적용 (권장)
    """
    out = {}
    for d in sorted(os.listdir(BLOCKS_DIR)):
        p = BLOCKS_DIR / d / "texts.md"
        if p.is_file():
            with open(p, encoding="utf-8") as f:
                text = f.read()
            if clean_meta:
                text = _strip_figma_meta(text)
            out[d] = text
    return out


def load_mdx_sections():
    """MDX 01~03을 ## 중목차 단위로 분리.
    {section_id: section_text}. section_id는 GT의 id와 매칭."""
    sections = {}
    for fname in ["01.mdx", "02.mdx", "03.mdx"]:
        mdx_num = fname.replace(".mdx", "")
        with open(MDX_DIR / fname, encoding="utf-8") as f:
            content = f.read()
        # frontmatter 제거
        content = re.sub(r"^---.*?---", "", content, flags=re.DOTALL).strip()
        # import 문 제거
        content = re.sub(r"^import .*$", "", content, flags=re.MULTILINE)
        # ## 기준 분리
        parts = re.split(r"(?=^## )", content, flags=re.MULTILINE)
        intro_collected = False
        sec_counter = 0
        for part in parts:
            part = part.strip()
            if not part:
                continue
            if part.startswith("## "):
                sec_counter += 1
                key = f"MDX{mdx_num}-{sec_counter}"
            else:
                key = f"MDX{mdx_num}-intro"
                intro_collected = True
            sections[key] = part
    return sections


def tokenize_simple(text):
    """단순 공백/특수문자 기준 토큰화."""
    # HTML/MDX 태그 제거
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"\{[^}]+\}", " ", text)
    text = re.sub(r"[|*#>\-\[\](){}:;,.!?/\\=\"'`~]", " ", text)
    # 2글자 이상 토큰만
    words = [w for w in text.split() if len(w) >= 2]
    return words


def clean_text(text):
    """기본 정리 (HTML/MDX 제거)."""
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"\{[^}]+\}", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def char_ngrams(text, n):
    """겹치는 n글자 집합. 공백 제거 후."""
    text = re.sub(r"\s+", "", clean_text(text))
    return set(text[i:i + n] for i in range(max(0, len(text) - n + 1)))


def evaluate_ranking(gt_entry, ranked_frames):
    """한 섹션의 평가.
    ranked_frames: [(frame_id, score), ...] top부터 내림차순.
    returns: dict with hit@1, hit@3, mrr
    """
    primary = gt_entry["primary"]
    secondary = gt_entry.get("secondary", []) or []

    if primary is None:
        # null 정답: top-1 점수가 낮으면 OK (즉 "없음"을 잘 맞춤)
        # 여기선 hit@1 = 1 if ranked_frames[0][1] < 0.2 else 0 같은 규칙
        # 복잡하니 일단 null은 "모든 방법이 약한 점수 내야 성공"으로 단순 처리
        # → hit 체크에서 제외, 별도 분석용
        return {"hit@1": None, "hit@3": None, "mrr": None, "gt": None}

    gt_set = {str(primary)} | {str(s) for s in secondary}
    rank_ids = [str(fid) for fid, _ in ranked_frames]

    hit1 = 1 if rank_ids and rank_ids[0] == str(primary) else 0
    hit3 = 1 if any(r in gt_set for r in rank_ids[:3]) else 0

    # MRR: primary의 역순위 (secondary는 MRR에서 제외 — 엄격)
    mrr = 0.0
    for i, r in enumerate(rank_ids, 1):
        if r == str(primary):
            mrr = 1.0 / i
            break

    return {"hit@1": hit1, "hit@3": hit3, "mrr": mrr, "gt": str(primary)}


def run_method(method_name, method_fn, mdx_sections, figma_texts, gt_list):
    """방법을 실행하고 결과 dict 반환.
    method_fn(mdx_text, figma_texts_dict) -> [(frame_id, score), ...] top-k desc
    """
    gt_by_id = {g["id"]: g for g in gt_list}
    results = {}
    for sec_id, sec_text in mdx_sections.items():
        if sec_id not in gt_by_id:
            continue
        ranked = method_fn(sec_text, figma_texts)
        results[sec_id] = {
            "ranked": ranked[:5],
            "eval": evaluate_ranking(gt_by_id[sec_id], ranked),
        }
    return {"method": method_name, "results": results}


def aggregate_metrics(method_result):
    """hit@1, hit@3, mrr 평균 (null GT 제외)"""
    vals = {"hit@1": [], "hit@3": [], "mrr": []}
    for sec_id, r in method_result["results"].items():
        e = r["eval"]
        if e["gt"] is None:
            continue
        for k in vals:
            vals[k].append(e[k])
    return {
        k: (sum(v) / len(v) if v else 0.0) for k, v in vals.items()
    }