#!/usr/bin/env python
"""Stage 02 — extract embeddings.

The one GPU-heavy stage: runs the frozen EmbeddingGemma + SigLIP2 encoders
over every row and caches the results as .npy files, so training (stage 03)
never has to re-run the encoders.

Order of operations (deliberate, not incidental):
  1. Load the text encoder, extract text embeddings for train AND test.
  2. Free the text encoder from RAM/VRAM.
  3. Load the image encoder, extract image embeddings for train AND test.
  4. Free the image encoder.
This keeps peak memory to "one encoder at a time" instead of both loaded
simultaneously, and means a crash during image extraction doesn't force
re-running text extraction too.

Every batch is checkpointed to disk (see src/encoders/batch_extract.py) —
if this script crashes or the session times out, re-running the exact same
command resumes from the last checkpoint instead of starting over.

Usage:
    python scripts/02_extract_embeddings.py --config configs/base.yaml
"""

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.encoders.batch_extract import free_encoder, resumable_batch_encode
from src.encoders.registry import build_image_encoder, build_text_encoder
from src.utils.config import load_config
from src.utils.exceptions import PricePredictorError
from src.utils.logging import get_logger
from src.utils.seed import set_seed

logger = get_logger(__name__)

SPLITS = ("train", "test")


def _load_split_df(processed_dir: Path, split: str) -> pd.DataFrame:
    parquet_path = processed_dir / f"{split}_clean.parquet"
    if not parquet_path.exists():
        raise PricePredictorError(f"{parquet_path} not found — run scripts/01_preprocess.py first")
    return pd.read_parquet(parquet_path)


def run(config_path: str) -> None:
    config = load_config(config_path)
    set_seed(config["seed"])

    data_cfg = config["data"]
    processed_dir = Path(data_cfg["processed_dir"])
    embeddings_dir = Path(data_cfg["embeddings_dir"])
    embeddings_dir.mkdir(parents=True, exist_ok=True)

    dfs = {split: _load_split_df(processed_dir, split) for split in SPLITS}
    for split, df in dfs.items():
        logger.info("Loaded split=%s: %d rows", split, len(df))
        df[["sample_id"]].to_csv(embeddings_dir / f"{split}_ids.csv", index=False)
        if split == "train":
            np.save(embeddings_dir / "train_price.npy", df["price"].to_numpy(dtype=np.float32))

    # ---- Phase 1: text embeddings (train + test), then free the encoder ----
    logger.info("=== Phase 1/2: text embeddings ===")
    text_cfg = config["encoders"]["text"]
    text_encoder = build_text_encoder(text_cfg)
    text_output_dim = text_cfg.get("mrl_truncate") or text_cfg["output_dim"]

    for split, df in dfs.items():
        texts = df["catalog_content"].tolist()
        resumable_batch_encode(
            items=texts,
            encode_batch_fn=lambda batch: text_encoder.encode(batch).numpy(),
            output_dim=text_output_dim,
            output_path=str(embeddings_dir / f"{split}_text.npy"),
            batch_size=text_cfg["batch_size"],
            desc=f"{split} text",
        )

    free_encoder(text_encoder)
    logger.info("=== Phase 1/2 complete, text encoder freed from memory ===")

    # ---- Phase 2: image embeddings (train + test), then free the encoder ----
    logger.info("=== Phase 2/2: image embeddings ===")
    image_cfg = config["encoders"]["image"]
    image_encoder = build_image_encoder(image_cfg)
    image_output_dim = image_cfg["output_dim"]

    def _encode_image_batch(paths):
        images = [Image.open(p).convert("RGB") for p in paths]
        return image_encoder.encode(images).numpy()

    for split, df in dfs.items():
        image_paths = df["image_path"].tolist()
        resumable_batch_encode(
            items=image_paths,
            encode_batch_fn=_encode_image_batch,
            output_dim=image_output_dim,
            output_path=str(embeddings_dir / f"{split}_image.npy"),
            batch_size=image_cfg["batch_size"],
            desc=f"{split} image",
        )

    free_encoder(image_encoder)
    logger.info("=== Phase 2/2 complete, image encoder freed from memory ===")

    # ---- Final sanity check: every split's text/image/row counts line up ----
    for split, df in dfs.items():
        text_matrix = np.load(embeddings_dir / f"{split}_text.npy")
        image_matrix = np.load(embeddings_dir / f"{split}_image.npy")
        if text_matrix.shape[0] != len(df) or image_matrix.shape[0] != len(df):
            raise PricePredictorError(
                f"Embedding count mismatch for split={split}: "
                f"text={text_matrix.shape[0]} image={image_matrix.shape[0]} rows={len(df)}"
            )
        logger.info(
            "Verified split=%s: text %s, image %s", split, text_matrix.shape, image_matrix.shape,
        )

    logger.info("All embeddings extracted and verified.")


def main() -> None:
    parser = argparse.ArgumentParser(description="Stage 02: extract and cache embeddings")
    parser.add_argument("--config", default="configs/base.yaml")
    args = parser.parse_args()

    try:
        run(args.config)
    except PricePredictorError as e:
        logger.error("Embedding extraction failed: %s", e)
        logger.error("Progress up to the last checkpoint was saved — re-run this exact command to resume.")
        sys.exit(1)
    except Exception as e:
        logger.exception("Unexpected error during embedding extraction: %s", e)
        logger.error("Progress up to the last checkpoint was saved — re-run this exact command to resume.")
        sys.exit(1)


if __name__ == "__main__":
    main()