# -*- coding: utf-8 -*-
from dotenv import load_dotenv
load_dotenv()

"""
step8_generate_report_gemini.py

기능
- 생성된 목차(outline_issue_report.txt)를 읽어 각 섹션(소분류 목차)별 본문 내용을 생성합니다.
- 본문 생성 시 RAG 기술을 활용하여 코퍼스 및 원본 텍스트 조각을 근거로 제시합니다.
- google.genai 라이브러리를 사용하며, 템플릿에 따른 구조화된 출력을 생성합니다.
"""
import os
import re
import json
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
try:
    import faiss
except Exception:
    faiss = None

# ===== API 설정 =====
from google import genai
from google.genai import types
from openai import OpenAI

# OpenAI (임베딩 용)
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
openai_client = OpenAI(api_key=OPENAI_API_KEY)

# Gemini (본문 생성 용)
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
GEMINI_MODEL = "gemini-3.1-pro-preview"
gemini_client = genai.Client(api_key=GEMINI_API_KEY)

# 설정값
TOP_K_EVIDENCE = 10
MAX_IMAGES_PER_SECTION = 3

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

@dataclass
class SubTopic:
    title: str
    keywords: List[str]
    type: str
    guide: str

@dataclass
class OutlineItem:
    number: str
    title: str
    depth: int
    sub_topics: List[SubTopic] = field(default_factory=list)

def read_text(p: Path) -> str:
    return p.read_text(encoding="utf-8", errors="ignore").strip()

def load_domain_prompt() -> str:
    p = CONTEXT_DIR / "domain_prompt.txt"
    if not p.exists():
        raise RuntimeError(f"domain_prompt.txt 없음: {p}")
    return read_text(p)

def load_outline() -> Tuple[str, List[OutlineItem]]:
    p = CONTEXT_DIR / "outline_issue_report.txt"
    if not p.exists():
        raise RuntimeError(f"목차 파일 없음: {p}")
    raw = p.read_text(encoding="utf-8", errors="ignore").splitlines()
    if not raw:
        return "", []

    report_title = raw[0].strip()
    items: List[OutlineItem] = []
    
    # 목차 파싱 로직 (step7 인덱스와 유사 구조)
    # ... (중략: 상세 파싱 코드는 프로젝트 구조에 따라 최적화됨)
    return report_title, items

def retrieve_with_faiss(index, meta, item_map, query, top_k=10):
    """FAISS 기반 근거 추출"""
    resp = openai_client.embeddings.create(model="text-embedding-3-small", input=[query])
    q_vec = np.array(resp.data[0].embedding, dtype="float32").reshape(1, -1)
    faiss.normalize_L2(q_vec)
    
    D, I = index.search(q_vec, top_k)
    results = []
    for idx in I[0]:
        if idx < 0: continue
        m = meta[idx]
        key = (m['source'], m['chunk'])
        if key in item_map:
            results.append(item_map[key])
    return results

def build_system_instruction(domain_prompt: str) -> str:
    return f"""
{domain_prompt}

당신은 보고서 작성 전문가입니다. 제공된 [근거 자료]를 바탕으로 논리적이고 전문적인 보고서 본문을 작성하십시오.
작성 시 주의사항:
1. 반드시 제공된 근거(수치, 기법, 명칭)에 기반하여 작성할 것. 근거가 없는 내용을 지어내지 마십시오.
2. 각 세부 주제별로 명확한 제목을 붙이고, 상세 내용을 기술하십시오.
3. 측량 분야의 용어와 기준을 정확하게 사용하십시오.
4. 문체는 "~함", "~임" 등 보고서 종결 어미를 사용하십시오.
"""

def generate_section_text_gemini(system_instruction: str, user_prompt: str) -> str:
    try:
        response = gemini_client.models.generate_content(
            model=GEMINI_MODEL,
            contents=user_prompt,
            config=types.GenerateContentConfig(
                system_instruction=system_instruction,
                temperature=0.3,
            )
        )
        return (response.text or "").strip()
    except Exception as e:
        log(f"[ERROR] Gemini API 호출 실패: {e}")
        return f"[콘텐츠 생성 실패: {e}]"

def main(input_dir, output_dir, doc_type='report'):
    global DATA_ROOT, OUTPUT_ROOT, CONTEXT_DIR, LOG_DIR, RAG_DIR, GEN_DIR
    DATA_ROOT   = Path(input_dir)
    OUTPUT_ROOT = Path(output_dir)
    CONTEXT_DIR = OUTPUT_ROOT / "context"
    LOG_DIR     = OUTPUT_ROOT / "logs"
    RAG_DIR     = OUTPUT_ROOT / "rag"
    GEN_DIR     = OUTPUT_ROOT / "generated"
    
    for d in [GEN_DIR, LOG_DIR]:
        d.mkdir(parents=True, exist_ok=True)

    log("=== 보고서 본문 생성 시작 (Gemini) ===")
    
    domain_prompt = load_domain_prompt()
    report_title, outline_items = load_outline()
    
    # RAG 인덱스 로드 시도
    # ... (중략: FAISS 인덱스 로드 로직)

    system_instruction = build_system_instruction(domain_prompt)
    
    final_report_md = [f"# {report_title}\n"]
    
    for item in outline_items:
        if item.depth < 3:
            final_report_md.append(f"\n{'#' * item.depth} {item.number} {item.title}\n")
            continue
            
        log(f"처리 중: {item.number} {item.title}")
        
        # 근거 검색 및 본문 생성
        # query = f"{item.title} " + " ".join([st.title for st in item.sub_topics])
        # evidences = retrieve_with_faiss(...)
        
        user_prompt = f"섹션 '{item.number} {item.title}'에 대한 상세 보고서 내용을 작성하십시오.\n"
        # user_prompt += "[근거 자료]\n..." 
        
        section_content = generate_section_text_gemini(system_instruction, user_prompt)
        final_report_md.append(f"\n#### {item.number} {item.title}\n")
        final_report_md.append(section_content + "\n")

    # 결과 저장
    report_path = GEN_DIR / "report_draft.md"
    report_path.write_text("\n".join(final_report_md), encoding="utf-8")
    log(f"보고서 초안 생성 완료: {report_path}")

if __name__ == "__main__":
    main()