"""
Query side of the Space: the same pre-processing as retrieval/submission.py
(desc-io cleanup, then the model's query prefix), L2-normalised output.

No torch import: the model is passed in (anything with a sentence-transformers style
`encode`), so tests can use a fake.
"""
from __future__ import annotations

from typing import Any

import numpy as np

from query_clean import clean_query

DEFAULT_QUERY_PREFIX = "Represent this query for searching relevant code: "
DEFAULT_QUERY_CLEAN = "desc-io"


class STQueryEncoder:
    def __init__(self, model: Any, query_prefix: str = DEFAULT_QUERY_PREFIX,
                 query_clean: str = DEFAULT_QUERY_CLEAN):
        self.model = model
        self.query_prefix = query_prefix
        self.query_clean = query_clean

    def preprocess(self, text: str) -> str:
        return self.query_prefix + clean_query(text.strip(), self.query_clean)

    def encode_query(self, text: str) -> np.ndarray:
        emb = self.model.encode([self.preprocess(text)], convert_to_numpy=True,
                                normalize_embeddings=True, show_progress_bar=False)
        return np.asarray(emb, dtype=np.float32)[0]