#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import argparse
from typing import List, Optional

import chromadb
from chromadb.config import Settings

# (선택) 유사도 검색용
try:
    from openai import OpenAI
    _HAS_OPENAI = True
except Exception:
    _HAS_OPENAI = False


def get_chroma_collection(name: str):
    client = chromadb.HttpClient(
        host=os.environ.get("CHROMA_HOST", "localhost"),
        port=int(os.environ.get("CHROMA_PORT", "8001")),
        settings=Settings(chroma_api_impl="rest"),
    )
    return client.get_collection(name)


def preview_items(ids: List[str], docs: Optional[List[str]], metas: Optional[List[dict]], title: str):
    print(f"\n=== {title} (count={len(ids)}) ===")
    for i, _id in enumerate(ids):
        doc_sample = None
        if docs and len(docs) > i:
            d = docs[i] or ""
            doc_sample = d.replace("\n", " ")[:120]
        meta_sample = None
        if metas and len(metas) > i:
            m = metas[i] or {}
            # answer 요약
            ans = (m.get("answer") or "")
            meta_sample = {**m}
            if ans:
                meta_sample["answer"] = (ans.replace("\n", " ")[:120] + ("..." if len(ans) > 120 else ""))
        print(f"- id: { _id }")
        if doc_sample:
            print(f"  doc: {doc_sample}")
        if meta_sample:
            print(f"  meta: {meta_sample}")


def find_by_contains_document(col, needle: str, limit: int = 20):
    # where_document 이용해 본문에서 포함 검색
    res = col.get(
        where_document={"$contains": needle},
        include=["ids", "documents", "metadatas"],
        limit=limit
    )
    ids = res.get("ids", [])
    docs = res.get("documents", [])
    metas = res.get("metadatas", [])
    preview_items(ids, docs, metas, f"where_document CONTAINS '{needle}'")
    return ids


def find_by_contains_metadata(col, key: str, needle: str, limit: int = 20):
    # 예: key="answer", needle="エプソン"
    res = col.get(
        where={key: {"$contains": needle}},
        include=["ids", "documents", "metadatas"],
        limit=limit
    )
    ids = res.get("ids", [])
    docs = res.get("documents", [])
    metas = res.get("metadatas", [])
    preview_items(ids, docs, metas, f"where META[{key}] CONTAINS '{needle}'")
    return ids


def find_by_table(col, table_name: str, limit: int = 1000):
    # saveVector 경로로 넣은 항목은 metadata에 {"table": "Qn"}가 있음
    res = col.get(
        where={"table": {"$eq": table_name}},
        include=["ids", "documents", "metadatas"],
        limit=limit
    )
    ids = res.get("ids", [])
    docs = res.get("documents", [])
    metas = res.get("metadatas", [])
    preview_items(ids, docs, metas, f"where META[table] == '{table_name}'")
    return ids


def find_by_similarity(col, text: str, n_results: int = 10, model: str = "text-embedding-3-large"):
    if not _HAS_OPENAI:
        raise RuntimeError("openai 패키지가 없습니다. `pip install openai` 후 사용하세요.")
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise RuntimeError("환경변수 OPENAI_API_KEY가 필요합니다.")
    client = OpenAI(api_key=api_key)
    emb = client.embeddings.create(model=model, input=text).data[0].embedding
    res = col.query(
        query_embeddings=[emb],
        n_results=n_results,
        include=["ids", "documents", "metadatas", "distances"]
    )
    ids = res.get("ids", [[]])[0]
    docs = res.get("documents", [[]])[0]
    metas = res.get("metadatas", [[]])[0]
    preview_items(ids, docs, metas, f"similarity to '{text}'")
    # 거리도 참고하려면 여기에서 res["distances"][0] 같이 출력 가능
    return ids


def delete_by_ids(col, ids: List[str]):
    if not ids:
        print("삭제할 ID가 없습니다.")
        return
    print(f"\nDeleting {len(ids)} items...")
    col.delete(ids=ids)
    print("✅ 삭제 완료")


def main():
    parser = argparse.ArgumentParser(description="ChromaDB cleanup helper")
    parser.add_argument("--collection", required=True, help="컬렉션 이름 (e.g., segue_qna_by_question4 / segue_saved_qna13)")
    parser.add_argument("--mode", required=True, choices=["contains_doc", "contains_meta", "similarity", "table"], help="검색/대상 모드")
    parser.add_argument("--query", help="검색어(문서/메타/유사도 모드에서 사용)")
    parser.add_argument("--meta_key", default="answer", help="contains_meta 모드에서 사용할 메타데이터 키 (기본: answer)")
    parser.add_argument("--table", help="table 모드에서 사용할 테이블명 (예: Q7)")
    parser.add_argument("--limit", type=int, default=20, help="가져올 최대 개수(미리보기)")
    parser.add_argument("--n_results", type=int, default=10, help="유사 검색 결과 개수")
    parser.add_argument("--delete", action="store_true", help="미리보기 대신 실제 삭제 실행")
    args = parser.parse_args()

    col = get_chroma_collection(args.collection)

    # 연결 확인용
    print(f"Connected to collection: {args.collection}")
    try:
        print("count:", col.count())
    except Exception:
        pass

    ids = []
    if args.mode == "contains_doc":
        if not args.query:
            raise SystemExit("--query 가 필요합니다.")
        ids = find_by_contains_document(col, args.query, limit=args.limit)

    elif args.mode == "contains_meta":
        if not args.query:
            raise SystemExit("--query 가 필요합니다.")
        ids = find_by_contains_metadata(col, args.meta_key, args.query, limit=args.limit)

    elif args.mode == "similarity":
        if not args.query:
            raise SystemExit("--query 가 필요합니다.")
        ids = find_by_similarity(col, args.query, n_results=args.n_results)

    elif args.mode == "table":
        if not args.table:
            raise SystemExit("--table 이 필요합니다. 예: --table Q7")
        ids = find_by_table(col, args.table, limit=args.limit)

    # 삭제 여부
    if args.delete and ids:
        # 안전장치: 한번 더 출력
        print("\n위 ID 들을 삭제합니다. 계속하려면 'yes' 입력:")
        confirm = input("> ").strip().lower()
        if confirm == "yes":
            delete_by_ids(col, ids)
        else:
            print("취소되었습니다. (미삭제)")
    else:
        print("\n(미리보기만 수행했습니다. 실제 삭제하려면 --delete 를 추가하세요.)")


if __name__ == "__main__":
    main()