"""
PubMed 논문 검색 테스트
Biopython Entrez를 사용한 의학 논문 검색
"""

import sys
import os

# UTF-8 인코딩 강제 (Windows 한글 깨짐 방지)
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

from Bio import Entrez
from dotenv import load_dotenv

# .env 파일 로드
load_dotenv()

# NCBI Entrez 설정
Entrez.email = os.getenv('PUBMED_EMAIL', 'test@example.com')
api_key = os.getenv('PUBMED_API_KEY')
if api_key:
    Entrez.api_key = api_key
    print(f"[INFO] PubMed API Key 사용 중 (최대 10 req/sec)")
else:
    print(f"[INFO] API Key 없음 (최대 3 req/sec 제한)")


def search_pubmed(query, max_results=5):
    """PubMed에서 논문 검색"""
    try:
        print(f"\n검색어: '{query}'")
        print("-" * 80)

        # 1. 검색 (PMID 목록 가져오기)
        handle = Entrez.esearch(
            db="pubmed",
            term=query,
            retmax=max_results,
            sort="relevance"
        )
        record = Entrez.read(handle)
        handle.close()

        pmids = record["IdList"]
        total_count = int(record["Count"])

        if not pmids:
            print(f"[WARNING] 검색 결과 없음")
            return []

        print(f"[OK] 총 {total_count}건 검색됨, 상위 {len(pmids)}건 조회 중...\n")

        # 2. 논문 상세 정보 가져오기
        handle = Entrez.efetch(
            db="pubmed",
            id=pmids,
            rettype="medline",
            retmode="xml"
        )
        papers = Entrez.read(handle)
        handle.close()

        results = []

        for idx, paper in enumerate(papers['PubmedArticle'], 1):
            try:
                article = paper['MedlineCitation']['Article']

                # PMID
                pmid = str(paper['MedlineCitation']['PMID'])

                # 제목
                title = article.get('ArticleTitle', '(제목 없음)')

                # 초록 (여러 섹션이 있을 수 있음)
                abstract_parts = article.get('Abstract', {}).get('AbstractText', [])
                if abstract_parts:
                    if isinstance(abstract_parts, list):
                        abstract = ' '.join([str(part) for part in abstract_parts])
                    else:
                        abstract = str(abstract_parts)
                else:
                    abstract = "(초록 없음)"

                # 저널
                journal = article.get('Journal', {}).get('Title', '(저널 없음)')

                # 출판 연도
                pub_date = article.get('Journal', {}).get('JournalIssue', {}).get('PubDate', {})
                year = pub_date.get('Year', '(연도 없음)')

                # 저자 (첫 번째 저자만)
                authors = article.get('AuthorList', [])
                if authors:
                    first_author = authors[0]
                    last_name = first_author.get('LastName', '')
                    initials = first_author.get('Initials', '')
                    author_str = f"{last_name} {initials}" if last_name else "(저자 없음)"
                else:
                    author_str = "(저자 없음)"

                result = {
                    'pmid': pmid,
                    'title': title,
                    'abstract': abstract[:500] + '...' if len(abstract) > 500 else abstract,
                    'journal': journal,
                    'year': year,
                    'author': author_str
                }

                results.append(result)

                # 출력
                print(f"[{idx}] PMID: {pmid}")
                print(f"제목: {title}")
                print(f"저자: {author_str}")
                print(f"저널: {journal} ({year})")
                print(f"초록: {result['abstract']}")
                print(f"링크: https://pubmed.ncbi.nlm.nih.gov/{pmid}/")
                print("-" * 80)

            except Exception as e:
                print(f"[ERROR] 논문 파싱 실패: {e}")
                continue

        return results

    except Exception as e:
        print(f"[ERROR] PubMed 검색 실패: {e}")
        return []


def main():
    """메인 실행"""
    print("=" * 80)
    print("PubMed 논문 검색 테스트")
    print("=" * 80)

    # 테스트 1: Statin과 CoQ10 관계
    print("\n[TEST 1] Statin과 CoQ10 근육 부작용 관계")
    results1 = search_pubmed("statin AND coq10 AND muscle", max_results=3)

    # 테스트 2: CoQ10 일반
    print("\n[TEST 2] CoQ10 보충제 효능")
    results2 = search_pubmed("coenzyme q10 supplementation benefits", max_results=3)

    # 테스트 3: 약물 상호작용
    print("\n[TEST 3] Atorvastatin 부작용")
    results3 = search_pubmed("atorvastatin adverse effects", max_results=3)

    print("\n" + "=" * 80)
    print("검색 완료")
    print("=" * 80)
    print(f"총 {len(results1) + len(results2) + len(results3)}개 논문 조회됨")
    print("\n[TIP] GraphRAG에 활용 방법:")
    print("  1. 검색된 PMID를 지식 그래프에 저장")
    print("  2. AI 추천 시 관련 논문 인용")
    print("  3. 예시: 'Statin 복용자에게 CoQ10 추천 (근거: PMID:12345678, 신뢰도: 85%)'")


if __name__ == '__main__':
    main()