from dataclasses import dataclass
from datetime import datetime, timezone
 
# 埋め込みモデルの定義を一箇所に集約します。
# モデルを増やすときはここに追記するだけにします。
EMBEDDING_MODELS = {
    "v1": {"name": "text-embedding-004", "dim": 768},
    "v2": {"name": "gemini-embedding-001", "dim": 1536},
}
 
# 現在「正」とするバージョン。切り替えはこの1行の変更で完結させます。
ACTIVE_EMBEDDING_VERSION = "v1"
 
@dataclass
class VectorRecord:
    doc_id: str
    embedding: list[float]
    model_version: str          # "v1" / "v2" など。空間の身分証明書です。
    indexed_at: str
 
    def is_compatible_with(self, query_version: str) -> bool:
        # 同じモデルバージョンでなければ比較してはいけません。
        return self.model_version == query_version

from google import genai
 
client = genai.Client(api_key="YOUR_GEMINI_API_KEY")
 
def embed_query(text: str, version: str) -> list[float]:
    """指定バージョンのモデルでクエリを埋め込みます。"""
    model_name = EMBEDDING_MODELS[version]["name"]
    result = client.models.embed_content(
        model=model_name,
        contents=text,
        config={"task_type": "RETRIEVAL_QUERY"},
    )
    return result.embeddings[0].values
 
def search(query: str, store, version: str = ACTIVE_EMBEDDING_VERSION, top_k: int = 5):
    """指定バージョンの空間内だけで検索します。"""
    q_vec = embed_query(query, version)
    # store 側も同じ version のベクトルだけを対象に絞り込みます。
    candidates = store.fetch_by_version(version)
    scored = [
        (rec.doc_id, _cosine(q_vec, rec.embedding))
        for rec in candidates
        if rec.is_compatible_with(version)   # 念のための二重チェック
    ]
    scored.sort(key=lambda x: x[1], reverse=True)
    return scored[:top_k]
 
def _cosine(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(y * y for y in b) ** 0.5
    return dot / (na * nb) if na and nb else 0.0

import time
import logging
 
logger = logging.getLogger("reindex")
 
def reembed_corpus(
    store,
    checkpoint_store,
    target_version: str = "v2",
    batch_size: int = 100,
    max_retries: int = 5,
):
    """全ドキュメントを target_version のモデルで埋め込み直します。
    チェックポイントから再開できるため、途中で落ちても安全です。
    """
    model_name = EMBEDDING_MODELS[target_version]["name"]
    last_done = checkpoint_store.get(f"reindex:{target_version}") or ""
 
    # まだ新バージョンで埋め込んでいないドキュメントだけを取り出します。
    pending = store.iter_documents(after_id=last_done)
 
    batch = []
    for doc in pending:
        batch.append(doc)
        if len(batch) >= batch_size:
            _flush_batch(store, checkpoint_store, batch, model_name,
                         target_version, max_retries)
            batch = []
    if batch:
        _flush_batch(store, checkpoint_store, batch, model_name,
                     target_version, max_retries)
 
def _flush_batch(store, checkpoint_store, batch, model_name,
                 target_version, max_retries):
    texts = [d.text for d in batch]
    for attempt in range(max_retries):
        try:
            res = client.models.embed_content(
                model=model_name,
                contents=texts,
                config={"task_type": "RETRIEVAL_DOCUMENT"},
            )
            for doc, emb in zip(batch, res.embeddings):
                store.upsert(VectorRecord(
                    doc_id=doc.id,
                    embedding=emb.values,
                    model_version=target_version,
                    indexed_at=datetime.now(timezone.utc).isoformat(),
                ))
            # バッチが書き終わってからチェックポイントを進めます。
            # 順序が逆だと、落ちたときに取りこぼしが出ます。
            checkpoint_store.set(f"reindex:{target_version}", batch[-1].id)
            logger.info("reembedded up to %s", batch[-1].id)
            return
        except Exception as e:
            wait = min(2 ** attempt, 60)   # 指数バックオフ（上限60秒）
            logger.warning("batch failed (%s), retry in %ss", e, wait)
            time.sleep(wait)
    raise RuntimeError(f"batch permanently failed at {batch[-1].id}")

def upsert_document(store, doc, migration_active: bool):
    """移行期間中は両方の空間へ書き込みます。"""
    versions = ["v1"]
    if migration_active:
        versions.append("v2")   # 新空間にも同時に反映します
 
    for ver in versions:
        model_name = EMBEDDING_MODELS[ver]["name"]
        res = client.models.embed_content(
            model=model_name,
            contents=doc.text,
            config={"task_type": "RETRIEVAL_DOCUMENT"},
        )
        store.upsert(VectorRecord(
            doc_id=doc.id,
            embedding=res.embeddings[0].values,
            model_version=ver,
            indexed_at=datetime.now(timezone.utc).isoformat(),
        ))