src/utils/yaml_loader.py

"""
YAML Frontmatter 기반 지식 로더
================================

이미 구조화된 YAML frontmatter가 있는 마크다운 파일을 직접 로드합니다.
LLM 호출 없이 기존 추출 결과를 활용합니다.
"""

import yaml
import frontmatter
from pathlib import Path
from typing import Dict, Any, List, Optional
from datetime import datetime


def load_extracted_knowledge(file_path: str) -> Dict[str, Any]:
    """
    마크다운 파일의 YAML frontmatter에서 extracted_knowledge를 로드합니다.
    
    두 가지 YAML 구조를 지원:
    1. extracted_knowledge 래퍼 안에 중첩된 구조 (KOR, Terms)
    2. 최상위 레벨에 credit_cards 등이 직접 있는 구조 (USA)
    
    Args:
        file_path: 마크다운 파일 경로
        
    Returns:
        extracted_knowledge 딕셔너리 또는 전체 frontmatter
    """
    # 지원하는 지식 유형 키들
    knowledge_keys = [
        # 기존 키들
        'credit_cards', 'membership_tiers', 'loyalty_program', 
        'subscription_programs', 'points_system', 'milestone_program',
        'best_price_guarantee', 'point_exclusions', 'general_policies',
        'common_card_features', 'hilton_honors_references', 'card_comparison_summary',
        'facts',
        # 호텔 프로퍼티 관련 키들 (Pullman, Fairmont, Hotel Naru 등)
        'hotel_properties', 'hotel_facilities', 'room_types', 'tier_implementations',
        'room_common_amenities', 'loyalty_program_features', 'pricing_analysis',
        'ratings', 'nearby_attractions', 'channel_implementations', 'member_rates',
        'dining_venues', 'room_service', 'policies', 'pros_cons', 'hotel_brands',
        'benefits', 'promotion', 'exclusions', 'terms_and_conditions', 'points_policy'
    ]
    
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"파일을 찾을 수 없습니다: {file_path}")
    
    with open(path, 'r', encoding='utf-8') as f:
        post = frontmatter.load(f)
    
    metadata = post.metadata
    
    # extracted_knowledge가 있으면 사용
    if 'extracted_knowledge' in metadata:
        ek = metadata['extracted_knowledge']
    else:
        # 최상위 레벨에서 지식 키들 수집
        ek = {}
        for key in knowledge_keys:
            if key in metadata:
                ek[key] = metadata[key]
    
    # 지식 데이터가 있는지 확인
    if ek:
        # 체인 결정 (다양한 소스에서 시도)
        chain = None
        
        # 1. identity에서 확인
        identity = metadata.get('identity', {})
        chain = identity.get('chain')
        
        # 2. document_reference 또는 document_ref에서 확인
        if not chain:
            doc_ref = metadata.get('document_reference', metadata.get('document_ref', {}))
            if isinstance(doc_ref, dict):
                if 'identity' in doc_ref:
                    chain = doc_ref['identity'].get('chain')
                else:
                    chain = doc_ref.get('chain')
        
        # 3. extracted_knowledge 내부에서 확인
        if not chain:
            if 'loyalty_program' in ek and ek['loyalty_program']:
                chain = ek['loyalty_program'].get('chain')
            elif 'points_system' in ek and ek['points_system']:
                chain = ek['points_system'].get('chain')
            elif 'credit_cards' in ek and ek['credit_cards']:
                chain = ek['credit_cards'][0].get('chain')
            elif 'membership_tiers' in ek and ek['membership_tiers']:
                chain = ek['membership_tiers'][0].get('chain')
            # 4. hotel_properties에서 체인 확인 (Pullman, Hotel Naru 등)
            elif 'hotel_properties' in ek and ek['hotel_properties']:
                chain = ek['hotel_properties'][0].get('chain')
            # 5. tier_implementations에서 체인 확인
            elif 'tier_implementations' in ek and ek['tier_implementations']:
                chain = ek['tier_implementations'][0].get('chain')
        
        # identity에 chain 추가
        if chain and not identity.get('chain'):
            identity['chain'] = chain
        
        return {
            'file_path': str(file_path),
            'identity': identity,
            'source': metadata.get('source', {}),
            'version': metadata.get('version', {}),
            'extracted_knowledge': ek
        }
    else:
        return {
            'file_path': str(file_path),
            'metadata': metadata,
            'has_extracted_knowledge': False
        }


def load_all_from_directory(
    directory_path: str,
    pattern: str = "**/*.md"
) -> List[Dict[str, Any]]:
    """
    디렉토리 내 모든 마크다운 파일에서 extracted_knowledge를 로드합니다.
    
    Args:
        directory_path: 디렉토리 경로
        pattern: 파일 패턴 (기본값: "**/*.md")
        
    Returns:
        [{'file_path': ..., 'extracted_knowledge': ...}, ...] 리스트
    """
    directory = Path(directory_path)
    if not directory.exists():
        raise FileNotFoundError(f"디렉토리를 찾을 수 없습니다: {directory_path}")
    
    results = []
    files = list(directory.glob(pattern))
    
    print(f"📂 {len(files)}개 파일 발견: {directory_path}")
    
    for file_path in files:
        try:
            data = load_extracted_knowledge(str(file_path))
            if data.get('has_extracted_knowledge', True):
                results.append(data)
                print(f"  ✅ {file_path.name}")
            else:
                print(f"  ⚠️  {file_path.name} (extracted_knowledge 없음)")
        except Exception as e:
            print(f"  ❌ {file_path.name}: {e}")
            continue
    
    return results


def get_summary(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    로드된 데이터의 요약 정보를 반환합니다.
    """
    identity = data.get('identity', {})
    ek = data.get('extracted_knowledge', {})
    
    summary = {
        'file_path': data.get('file_path'),
        'chain': identity.get('chain'),
        'program_name': identity.get('program_name'),
        'doc_type': identity.get('doc_type'),
        'extraction_timestamp': ek.get('extraction_timestamp'),
        'extractor_model': ek.get('extractor_model'),
    }
    
    # 콘텐츠 요약
    if 'subscription_programs' in ek:
        summary['subscription_programs_count'] = len(ek['subscription_programs'])
    if 'membership_tiers' in ek:
        summary['membership_tiers_count'] = len(ek['membership_tiers'])
    if 'loyalty_program' in ek:
        summary['has_loyalty_program'] = True
    if 'credit_cards' in ek:
        summary['credit_cards_count'] = len(ek['credit_cards'])
    if 'benefits' in ek:
        summary['benefits_count'] = len(ek['benefits'])
    
    return summary


if __name__ == "__main__":
    import json
    
    print("🧪 YAML Frontmatter 로더 테스트")
    print("=" * 60)
    
    # 디렉토리 로드 테스트
    results = load_all_from_directory("data/raw/Hotel")
    
    print(f"\n📊 로드 결과: {len(results)}개 파일")
    print("=" * 60)
    
    for data in results:
        summary = get_summary(data)
        print(f"\n📄 {Path(summary['file_path']).name}")
        print(f"   체인: {summary.get('chain')}")
        print(f"   프로그램: {summary.get('program_name')}")
        print(f"   문서타입: {summary.get('doc_type')}")
        print(f"   추출 모델: {summary.get('extractor_model')}")
        
        if summary.get('subscription_programs_count'):
            print(f"   구독 프로그램: {summary['subscription_programs_count']}개")
        if summary.get('membership_tiers_count'):
            print(f"   멤버십 등급: {summary['membership_tiers_count']}개")
        if summary.get('credit_cards_count'):
            print(f"   신용카드: {summary['credit_cards_count']}개")