# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Université du Québec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.embeddings import EmbeddingManager, EmbeddingModelRegistry
from src.storage.database import DatabaseManager
from src.storage.parquet_manager import ParquetManager
from src.utils.logging import setup_logging

logger = setup_logging("build_embeddings")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build embedding vectors for filing chunks or sections"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="all-MiniLM-L6-v2",
        help="Embedding model name (default: all-MiniLM-L6-v2)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=256,
        help="Batch size for encoding (default: 256)",
    )
    parser.add_argument(
        "--source",
        type=str,
        choices=["chunk", "section"],
        default="chunk",
        help="Source type to embed: chunk or section (default: chunk)",
    )
    parser.add_argument(
        "--accession",
        type=str,
        default=None,
        help="Comma-separated accession numbers to embed (default: all)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    registry = EmbeddingModelRegistry()
    logger.info("Available models:")
    for model in registry.list_models():
        default_marker = " [default]" if model.get("is_default") else ""
        logger.info(
            "  %s (dim=%d, provider=%s)%s",
            model["name"], model["dimension"], model["provider"], default_marker,
        )

    db = DatabaseManager()
    db.initialize()
    parquet = ParquetManager()

    manager = EmbeddingManager(db, parquet, model_name=args.model)

    accession_numbers = None
    if args.accession:
        accession_numbers = [a.strip() for a in args.accession.split(",")]
        logger.info("Embedding specific accessions: %s", accession_numbers)

    try:
        if args.source == "chunk":
            count = manager.embed_chunks(
                accession_numbers=accession_numbers,
                model_name=args.model,
                batch_size=args.batch_size,
            )
        else:
            count = manager.embed_sections(
                accession_numbers=accession_numbers,
                model_name=args.model,
            )

        logger.info("Embedded %d %ss with model '%s'", count, args.source, args.model)

    except NotImplementedError as e:
        logger.error("Model not supported: %s", e)
        sys.exit(1)
    except Exception:
        logger.error("Embedding failed", exc_info=True)
        sys.exit(1)

    stats = manager.get_embedding_stats()
    logger.info("Embedding statistics:")
    logger.info("  Total embeddings: %d", stats["total_embeddings"])
    for model_name, sources in stats.get("by_model", {}).items():
        for source_type, info in sources.items():
            logger.info(
                "  %s / %s: %d vectors",
                model_name, source_type, info["count"],
            )

    db.close()
    logger.info("Embedding build complete")


if __name__ == "__main__":
    main()