# -*- coding: utf-8 -*-
"""
chunk_and_summary_v2.py

기능:
- 정리중 폴더 아래의 .md 파일들을 대상으로
  1) domain_prompt.txt 기반 GPT 의미 청킹
  2) 청크별 요약 생성
  3) 청크 내 이미지 참조 보존
  4) JSON 저장 (원문+청크+요약+이미지)
  5) RAG용 *_chunks.json 저장

전제:
- extract_1_v2.py 실행 후 .md 파일들이 존재할 것
- step1_domainprompt.py 실행 후 domain_prompt.txt가 존재할 것
"""

import os
import sys
import json
import re
from pathlib import Path
from datetime import datetime

from openai import OpenAI
from api_config import API_KEYS

# ===== 경로 =====
DATA_ROOT   = Path(r"D:\for python\geulbeot-light\geulbeot-light\00.test\hwpx\out")
OUTPUT_ROOT = Path(r"D:\for python\geulbeot-light\geulbeot-light\00.test\hwpx\out\out")   # 출력 위치

TEXT_DIR       = OUTPUT_ROOT / "text"
JSON_DIR       = OUTPUT_ROOT / "json"
RAG_DIR        = OUTPUT_ROOT / "rag"
CONTEXT_DIR    = OUTPUT_ROOT / "context"
LOG_DIR        = OUTPUT_ROOT / "logs"

for d in [TEXT_DIR, JSON_DIR, RAG_DIR, CONTEXT_DIR, LOG_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# ===== OpenAI 설정 =====
OPENAI_API_KEY = API_KEYS.get('GPT_API_KEY', '')
GPT_MODEL      = "gpt-5-2025-08-07"

client = OpenAI(api_key=OPENAI_API_KEY)

# ===== 스킵할 폴더 =====
SKIP_DIR_NAMES = {"System Volume Information", "$RECYCLE.BIN", ".git", "__pycache__", "output"}

# ===== 이미지 참조 패턴 =====
IMAGE_PATTERN = re.compile(r'!\[([^\]]*)\]\(([^)]+)\)')


def log(msg: str):
    print(msg, flush=True)
    with (LOG_DIR / "chunk_and_summary_log.txt").open("a", encoding="utf-8") as f:
        f.write(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}\n")


def load_domain_prompt() -> str:
    p = CONTEXT_DIR / "domain_prompt.txt"
    if not p.exists():
        log(f"domain_prompt.txt가 없습니다: {p}")
        log("먼저 step1_domainprompt.py를 실행해야 합니다.")
        sys.exit(1)
    return p.read_text(encoding="utf-8", errors="ignore").strip()


def safe_rel(p: Path) -> str:
    """DATA_ROOT 기준 상대 경로 반환"""
    try:
        return str(p.relative_to(DATA_ROOT))
    except Exception:
        return str(p)


def extract_text_md(p: Path) -> str:
    """마크다운 파일 텍스트 읽기"""
    try:
        return p.read_text(encoding="utf-8", errors="ignore")
    except Exception:
        return p.read_text(encoding="cp949", errors="ignore")


def find_images_in_text(text: str) -> list:
    """텍스트에서 이미지 참조 찾기"""
    matches = IMAGE_PATTERN.findall(text)
    return [{"alt": m[0], "path": m[1]} for m in matches]


def semantic_chunk(domain_prompt: str, text: str, source_name: str):
    """GPT 기반 의미 청킹"""
    if not text.strip():
        return []

    # 텍스트가 너무 짧으면 그냥 하나의 청크로
    if len(text) < 500:
        return [{
            "title": "전체 내용",
            "keywords": "",
            "content": text
        }]

    user_prompt = f"""
아래 문서를 의미 단위(문단/항목/섹션 등)로 분리하고,
각 청크는 title / keywords / content 를 포함한 JSON 배열로 출력하라.

규칙:
1. 추측 금지, 문서 내용 기반으로만 분리
2. 이미지 참조(![...](...))는 관련 텍스트와 같은 청크에 포함
3. 각 청크는 최소 100자 이상
4. keywords는 쉼표로 구분된 핵심 키워드 3~5개

문서:
{text[:12000]}

JSON 배열만 출력하라. 다른 설명 없이.
"""

    try:
        resp = client.chat.completions.create(
            model=GPT_MODEL,
            messages=[
                {"role": "system", "content": domain_prompt + "\n\n너는 의미 기반 청킹 전문가이다. JSON 배열만 출력한다."},
                {"role": "user", "content": user_prompt},
            ],
        )
        data = resp.choices[0].message.content.strip()
        
        # JSON 파싱 시도
        # ```json ... ``` 형식 처리
        if "```json" in data:
            data = data.split("```json")[1].split("```")[0].strip()
        elif "```" in data:
            data = data.split("```")[1].split("```")[0].strip()
        
        if data.startswith("["):
            return json.loads(data)
            
    except json.JSONDecodeError as e:
        log(f"[WARN] JSON 파싱 실패 ({source_name}): {e}")
    except Exception as e:
        log(f"[WARN] semantic_chunk API 실패 ({source_name}): {e}")

    # fallback: 페이지/섹션 기반 분리
    log(f"[INFO] Fallback 청킹 적용: {source_name}")
    return fallback_chunk(text)


def fallback_chunk(text: str) -> list:
    """GPT 실패 시 대체 청킹 (페이지/섹션 기반)"""
    chunks = []
    
    # 페이지 구분자로 분리 시도
    if "## Page " in text:
        pages = re.split(r'\n## Page \d+\n', text)
        for i, page_content in enumerate(pages):
            if page_content.strip():
                chunks.append({
                    "title": f"Page {i+1}",
                    "keywords": "",
                    "content": page_content.strip()
                })
    else:
        # 빈 줄 2개 이상으로 분리
        sections = re.split(r'\n{3,}', text)
        for i, section in enumerate(sections):
            if section.strip() and len(section.strip()) > 50:
                chunks.append({
                    "title": f"섹션 {i+1}",
                    "keywords": "",
                    "content": section.strip()
                })
    
    # 청크가 없으면 전체를 하나로
    if not chunks:
        chunks.append({
            "title": "전체 내용",
            "keywords": "",
            "content": text.strip()
        })
    
    return chunks


def summary_chunk(domain_prompt: str, text: str, limit: int = 300) -> str:
    """청크 요약 생성"""
    if not text.strip():
        return ""
    
    # 이미지 참조 제거 후 요약 (텍스트만)
    text_only = IMAGE_PATTERN.sub('', text).strip()
    
    if len(text_only) < 100:
        return text_only
    
    prompt = f"""
아래 텍스트를 {limit}자 이내로 사실 기반으로 요약하라.
추측 금지, 고유명사와 수치는 보존.

{text_only[:8000]}
"""
    try:
        resp = client.chat.completions.create(
            model=GPT_MODEL,
            messages=[
                {"role": "system", "content": domain_prompt + "\n\n너는 사실만 요약하는 전문가이다."},
                {"role": "user", "content": prompt},
            ],
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        log(f"[WARN] summary 실패: {e}")
        return text_only[:limit]


def save_chunk_files(src: Path, text: str, domain_prompt: str) -> int:
    """
    의미 청킹 → 요약 → JSON 저장
    
    Returns:
        생성된 청크 수
    """
    stem = src.stem
    folder_ctx = safe_rel(src.parent)
    
    # 원문 저장
    (TEXT_DIR / f"{stem}_text.txt").write_text(text, encoding="utf-8", errors="ignore")
    
    # 의미 청킹
    chunks = semantic_chunk(domain_prompt, text, src.name)
    
    if not chunks:
        log(f"[WARN] 청크 없음: {src.name}")
        return 0
    
    rag_items = []
    
    for idx, ch in enumerate(chunks, start=1):
        content = ch.get("content", "")
        
        # 요약 생성
        summ = summary_chunk(domain_prompt, content, 300)
        
        # 이 청크에 포함된 이미지 찾기
        images_in_chunk = find_images_in_text(content)
        
        rag_items.append({
            "source": src.name,
            "source_path": safe_rel(src),
            "chunk": idx,
            "total_chunks": len(chunks),
            "title": ch.get("title", ""),
            "keywords": ch.get("keywords", ""),
            "text": content,
            "summary": summ,
            "folder_context": folder_ctx,
            "images": images_in_chunk,
            "has_images": len(images_in_chunk) > 0
        })
    
    # JSON 저장
    (JSON_DIR / f"{stem}.json").write_text(
        json.dumps(rag_items, ensure_ascii=False, indent=2),
        encoding="utf-8"
    )
    
    # RAG용 JSON 저장
    (RAG_DIR / f"{stem}_chunks.json").write_text(
        json.dumps(rag_items, ensure_ascii=False, indent=2),
        encoding="utf-8"
    )
    
    return len(chunks)


def main():
    log("=" * 60)
    log("청킹/요약 파이프라인 시작")
    log(f"데이터 폴더: {DATA_ROOT}")
    log(f"출력 폴더: {OUTPUT_ROOT}")
    log("=" * 60)
    
    # 도메인 프롬프트 로드
    domain_prompt = load_domain_prompt()
    log(f"도메인 프롬프트 로드 완료 ({len(domain_prompt)}자)")
    
    # 통계
    stats = {"docs": 0, "chunks": 0, "images": 0, "errors": 0}
    
    # .md 파일 찾기
    md_files = []
    for root, dirs, files in os.walk(DATA_ROOT):
        dirs[:] = [d for d in dirs if d not in SKIP_DIR_NAMES and not d.startswith(".")]
        for fname in files:
            if fname.lower().endswith(".md"):
                md_files.append(Path(root) / fname)
    
    log(f"\n총 {len(md_files)}개 .md 파일 발견\n")
    
    for idx, fpath in enumerate(md_files, 1):
        try:
            rel_path = safe_rel(fpath)
            log(f"[{idx}/{len(md_files)}] {rel_path}")
            
            # 텍스트 읽기
            text = extract_text_md(fpath)
            
            if not text.strip():
                log(f"    ⚠ 빈 파일, 스킵")
                continue
            
            # 이미지 개수 확인
            images = find_images_in_text(text)
            stats["images"] += len(images)
            
            # 청킹 및 저장
            chunk_count = save_chunk_files(fpath, text, domain_prompt)
            
            stats["docs"] += 1
            stats["chunks"] += chunk_count
            
            log(f"    ✓ {chunk_count}개 청크, {len(images)}개 이미지")
            
        except Exception as e:
            stats["errors"] += 1
            log(f"    ✗ 오류: {e}")
    
    # 전체 통계 저장
    summary = {
        "processed_at": datetime.now().isoformat(),
        "data_root": str(DATA_ROOT),
        "output_root": str(OUTPUT_ROOT),
        "statistics": stats
    }
    
    (LOG_DIR / "chunk_summary_stats.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2),
        encoding="utf-8"
    )
    
    # 결과 출력
    log("\n" + "=" * 60)
    log("청킹/요약 완료!")
    log("=" * 60)
    log(f"처리된 문서: {stats['docs']}개")
    log(f"생성된 청크: {stats['chunks']}개")
    log(f"포함된 이미지: {stats['images']}개")
    log(f"오류: {stats['errors']}개")
    log(f"\n결과 저장 위치:")
    log(f"  - 원문: {TEXT_DIR}")
    log(f"  - JSON: {JSON_DIR}")
    log(f"  - RAG: {RAG_DIR}")


if __name__ == "__main__":
    main()