#!/usr/bin/env python3
"""
벡터 임베딩 스크립트
- DEVONthink 문서 UUID로 텍스트 추출
- GPU 서버(nomic-embed-text)로 임베딩 생성
- ChromaDB에 저장
"""

import os
import sys
import requests
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from pkm_utils import setup_logger, load_credentials, run_applescript_inline

logger = setup_logger("embed")

# ChromaDB 저장 경로
CHROMA_DIR = Path.home() / ".local" / "share" / "pkm" / "chromadb"
CHROMA_DIR.mkdir(parents=True, exist_ok=True)


def get_document_text(uuid: str) -> tuple[str, str]:
    """DEVONthink에서 UUID로 문서 텍스트 + 제목 추출"""
    script = f'''
    tell application id "DNtp"
        set theRecord to get record with uuid "{uuid}"
        set docText to plain text of theRecord
        set docTitle to name of theRecord
        return docTitle & "|||" & docText
    end tell
    '''
    result = run_applescript_inline(script)
    parts = result.split("|||", 1)
    title = parts[0] if len(parts) > 0 else ""
    text = parts[1] if len(parts) > 1 else ""
    return title, text


def get_embedding(text: str, gpu_server_ip: str) -> list[float] | None:
    """GPU 서버의 nomic-embed-text로 임베딩 생성"""
    url = f"http://{gpu_server_ip}:11434/api/embeddings"
    try:
        resp = requests.post(url, json={
            "model": "nomic-embed-text",
            "prompt": text[:8000]  # 토큰 제한
        }, timeout=60)
        resp.raise_for_status()
        return resp.json().get("embedding")
    except Exception as e:
        logger.error(f"임베딩 생성 실패: {e}")
        return None


def store_in_chromadb(doc_id: str, title: str, text: str, embedding: list[float]):
    """ChromaDB에 저장"""
    import chromadb
    client = chromadb.PersistentClient(path=str(CHROMA_DIR))
    collection = client.get_or_create_collection(
        name="pkm_documents",
        metadata={"hnsw:space": "cosine"}
    )
    collection.upsert(
        ids=[doc_id],
        embeddings=[embedding],
        documents=[text[:2000]],
        metadatas=[{"title": title, "source": "devonthink"}]
    )
    logger.info(f"ChromaDB 저장: {doc_id} ({title[:30]})")


def run(uuid: str):
    """단일 문서 임베딩 처리"""
    logger.info(f"임베딩 처리 시작: {uuid}")

    creds = load_credentials()
    gpu_ip = creds.get("GPU_SERVER_IP")
    if not gpu_ip:
        logger.warning("GPU_SERVER_IP 미설정 — 임베딩 건너뜀")
        return

    try:
        title, text = get_document_text(uuid)
        if not text or len(text) < 10:
            logger.warning(f"텍스트 부족 [{uuid}]: {len(text)}자")
            return

        embedding = get_embedding(text, gpu_ip)
        if embedding:
            store_in_chromadb(uuid, title, text, embedding)
            logger.info(f"임베딩 완료: {uuid}")
        else:
            logger.error(f"임베딩 실패: {uuid}")

    except Exception as e:
        logger.error(f"임베딩 처리 에러 [{uuid}]: {e}", exc_info=True)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("사용법: python3 embed_to_chroma.py <DEVONthink_UUID>")
        sys.exit(1)
    run(sys.argv[1])