"""
Create the ProductProfile Weaviate collection with Gemini embeddings.

Uses Google AI Studio (Gemini API) with:
  - model: gemini-embedding-001
  - task type: RETRIEVAL_DOCUMENT (optimized for indexing documents)

Usage:
    python setup_weaviate.py              # create collection if it does not exist
    python setup_weaviate.py --recreate    # delete and recreate (required after embedding config changes)
"""

from __future__ import annotations

import argparse
import os
import sys
from typing import Dict

try:
    import weaviate
except ModuleNotFoundError:
    print(
        "[FATAL] Missing dependency 'weaviate'. Activate the project virtual environment "
        "and install requirements:\n"
        "  venv\\Scripts\\activate\n"
        "  pip install -r requirements.txt",
        file=sys.stderr,
    )
    sys.exit(1)

from dotenv import load_dotenv
from weaviate.classes.config import Configure, DataType, Property

COLLECTION_NAME = "ProductProfile"
EMBEDDING_MODEL = "gemini-embedding-001"
EMBEDDING_TASK_TYPE = "RETRIEVAL_DOCUMENT"


def load_config() -> Dict[str, str]:
    load_dotenv()
    gemini_api_key = os.getenv("GEMINI_API_KEY", "")
    if not gemini_api_key:
        raise EnvironmentError(
            "Missing required environment variable: GEMINI_API_KEY. "
            "Copy .env.example to .env and fill in the values."
        )
    return {
        "GEMINI_API_KEY": gemini_api_key,
        "WEAVIATE_HOST": os.getenv("WEAVIATE_HOST", "localhost"),
        "WEAVIATE_HTTP_PORT": os.getenv("WEAVIATE_HTTP_PORT", "8080"),
        "WEAVIATE_GRPC_PORT": os.getenv("WEAVIATE_GRPC_PORT", "50051"),
    }


def connect_weaviate(cfg: Dict[str, str]) -> weaviate.WeaviateClient:
    return weaviate.connect_to_local(
        host=cfg["WEAVIATE_HOST"],
        port=int(cfg["WEAVIATE_HTTP_PORT"]),
        grpc_port=int(cfg["WEAVIATE_GRPC_PORT"]),
        headers={"X-Goog-Studio-Api-Key": cfg["GEMINI_API_KEY"]},
    )


def ensure_product_profile_collection(
    client: weaviate.WeaviateClient,
    *,
    recreate: bool = False,
) -> bool:
    """
    Ensure the ProductProfile collection exists with the expected embedding config.

    Returns True if a new collection was created, False if an existing one was kept.
    """
    if client.collections.exists(COLLECTION_NAME):
        if recreate:
            client.collections.delete(COLLECTION_NAME)
        else:
            return False

    client.collections.create(
        name=COLLECTION_NAME,
        properties=[
            Property(
                name="product_id",
                data_type=DataType.TEXT,
                skip_vectorization=True,
                index_filterable=True,
            ),
            Property(
                name="product_name",
                data_type=DataType.TEXT,
                skip_vectorization=True,
                index_filterable=True,
            ),
            Property(
                name="retrieval_summary",
                data_type=DataType.TEXT,
            ),
        ],
        vector_config=Configure.Vectors.text2vec_google_gemini(
            source_properties=["retrieval_summary"],
            model=EMBEDDING_MODEL,
            task_type=EMBEDDING_TASK_TYPE,
            vectorize_collection_name=False,
        ),
    )
    return True


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create the ProductProfile Weaviate collection with Gemini embeddings."
    )
    parser.add_argument(
        "--recreate",
        action="store_true",
        help="Delete and recreate the collection (use after changing embedding settings).",
    )
    return parser.parse_args(argv)


def main() -> None:
    args = parse_args()
    cfg = load_config()
    client = connect_weaviate(cfg)

    try:
        had_collection = client.collections.exists(COLLECTION_NAME)
        created = ensure_product_profile_collection(client, recreate=args.recreate)
        if created and had_collection and args.recreate:
            print(
                f"Recreated collection '{COLLECTION_NAME}' "
                f"(model={EMBEDDING_MODEL}, task_type={EMBEDDING_TASK_TYPE})."
            )
        elif created:
            print(
                f"Created collection '{COLLECTION_NAME}' "
                f"(model={EMBEDDING_MODEL}, task_type={EMBEDDING_TASK_TYPE})."
            )
        else:
            print(f"Collection '{COLLECTION_NAME}' already exists; no changes made.")
    finally:
        client.close()


if __name__ == "__main__":
    try:
        main()
    except EnvironmentError as exc:
        print(f"[FATAL] {exc}", file=sys.stderr)
        sys.exit(1)