from __future__ import annotations

from abc import ABC, abstractmethod
from datetime import datetime
from os import PathLike
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Literal,
    Sequence,
    TypedDict,
    cast,
    get_args,
    overload,
)

from httpx._types import FileTypes  # type: ignore
from tqdm.auto import tqdm

from ._utils.common import CreateMode, DropMode
from ._utils.tqdm_file_reader import TqdmFileReader
from .client import (
    EmbeddingEvaluationRequest,
    EmbeddingFinetuneConfig,
    EmbedRequest,
    FinetunedEmbeddingModelMetadata,
)
from .client import FinetunedEmbeddingModelTrial as FinetunedEmbeddingModelTrialResponse
from .client import (
    FinetuneEmbeddingModelRequest,
    OrcaClient,
    PostEmbeddingModelUploadRequest,
    PretrainedEmbeddingModelMetadata,
    PretrainedEmbeddingModelName,
    TrialStatus,
)
from .datasource import Datasource
from .job import Job, Status

if TYPE_CHECKING:
    from .classification_model import ClassificationMetrics
    from .memoryset import LabeledMemoryset, ScoredMemoryset
    from .regression_model import RegressionMetrics

LearningRateScheduler = Literal["linear", "cosine", "constant"]
"""
Learning rate scheduler for embedding model finetuning.

* **`"linear"`**: Linearly decays to zero after warmup (default).
* **`"cosine"`**: Cosine annealing to zero.
* **`"constant"`**: Fixed learning rate (warmup is applied when configured).
"""

FinetuningLoss = Literal["prediction", "triplet", "contrastive", "proxy"]
"""
Loss function for embedding model finetuning.

* **`"prediction"`**: Linear prediction head. Works for both categorical labels
  and continuous scores.
* **`"contrastive"`**: In-batch contrastive loss. Often produces better
  embeddings than prediction; trains embeddings directly for similarity
  and scales well to large batches.
* **`"triplet"`**: Batch-hard triplet loss — pulls same-class embeddings
  together. Simpler than contrastive but requires all samples to fit on one
  GPU.
* **`"proxy"`**: Proxy-anchor loss — learns class proxies in embedding space.
  Particularly useful for class-imbalanced datasets.
"""


class EmbeddingFinetuneHyperparams(TypedDict, total=False):
    """
    Training hyperparameters for embedding model finetuning.

    All fields are optional — sensible loss-specific defaults are applied for
    anything you don't set. You only need to override what you care about.

    Sweep mode (trial_count > 1) runs an Optuna hyperparameter search. The
    easiest way to sweep is to just set trial_count — the system auto-injects
    loss-specific search ranges for the most impactful parameters
    (learning_rate, batch_size, epochs, warmup, and loss_scale where
    applicable) up to what the trial budget can support.

    For finer control, pass explicit ranges in this dict:

    - Tuples (min, max) search a continuous range (log-uniform for
      learning_rate, uniform otherwise).
    - Lists [a, b, c] search discrete categorical choices.
    - Scalars fix a parameter to a single value (excluded from the search).

    The sweepable parameters are: learning_rate, epochs, batch_size, warmup,
    weight_decay, loss_scale, normalize_embeddings, and
    learning_rate_scheduler. As a rule of thumb, Optuna needs roughly 2n + 1
    trials for n search parameters.
    """

    epochs: int | tuple[int, int] | list[int]
    """Number of full passes over the training data. Defaults to 1 for single
    runs, 2 for sweeps, or 3 when early stopping is enabled."""

    max_steps: int
    """Maximum training steps. Overrides `epochs` when set."""

    learning_rate: float | tuple[float, float] | list[float]
    """Peak learning rate after warmup."""

    batch_size: int | list[int]
    """Total samples per training step."""

    warmup: int | float | tuple[int, int] | tuple[float, float] | list[int] | list[float]
    """Learning rate warmup. int = steps, float = fraction of total steps (0-1)."""

    weight_decay: float | tuple[float, float] | list[float]
    """L2 regularization strength (typical range: 0.0 to 0.1)."""

    learning_rate_scheduler: LearningRateScheduler | list[LearningRateScheduler]
    """How the learning rate changes after warmup."""

    loss_scale: float | tuple[float, float] | list[float]
    """Inverse temperature for contrastive and proxy losses."""

    normalize_embeddings: bool | list[bool]
    """L2-normalize embeddings before the classification/regression head."""

    max_seq_length: int | Literal["p90", "p95", "p99", "max"]
    """Maximum token length for input text, or a percentile string."""

    truncation_side: Literal["left", "right"]
    """Which end to cut when text exceeds max_seq_length."""

    early_stopping: bool | int
    """Stop when eval metric plateaus. True = patience 2, int = custom patience.
    Auto-enabled for "head" and "loss" eval methods; must be set explicitly
    for "neighbor" eval."""

    early_stopping_threshold: float
    """Minimum improvement to count as progress for early stopping."""

    eval_method: Literal["neighbor", "head", "loss"]
    """How to measure model quality during training. Defaults to "head" for
    single-run classification/regression (fast, auto-enables early stopping),
    "neighbor" for sweeps and metric losses (runs once at trial end by
    default)."""

    eval_steps: int | Literal["epoch", "end", "off"]
    """How often to evaluate (int = every N steps). Defaults to every 50
    steps for "head"/"loss" eval, end-of-training for "neighbor" eval."""

    neighbor_eval_count: int
    """Number of nearest neighbors for neighbor evaluation."""

    neighbor_eval_pool_subsample: int | float
    """Reduce the neighbor search pool. int = sample count, float = fraction."""


class EmbeddingModelBase(ABC):
    embedding_dim: int
    max_seq_length: int
    num_params: int | None
    uses_context: bool
    supports_instructions: bool

    def __init__(
        self,
        *,
        name: str,
        embedding_dim: int,
        max_seq_length: int,
        num_params: int | None,
        uses_context: bool,
        supports_instructions: bool,
    ):
        self.embedding_dim = embedding_dim
        self.max_seq_length = max_seq_length
        self.num_params = num_params
        self.uses_context = uses_context
        self.supports_instructions = supports_instructions

    @classmethod
    @abstractmethod
    def all(cls) -> Sequence[EmbeddingModelBase]:
        pass

    def _get_instruction_error_message(self) -> str:
        """Get error message for instruction not supported"""
        if isinstance(self, FinetunedEmbeddingModel):
            return f"Model {self.name} does not support instructions. Instruction-following is only supported by models based on instruction-supporting models."
        elif isinstance(self, PretrainedEmbeddingModel):
            return f"Model {self.name} does not support instructions. Instruction-following is only supported by instruction-supporting models."
        else:
            raise ValueError("Invalid embedding model")

    @overload
    def embed(self, value: str, max_seq_length: int | None = None, prompt: str | None = None) -> list[float]:
        pass

    @overload
    def embed(
        self, value: list[str], max_seq_length: int | None = None, prompt: str | None = None
    ) -> list[list[float]]:
        pass

    def embed(
        self, value: str | list[str], max_seq_length: int | None = None, prompt: str | None = None
    ) -> list[float] | list[list[float]]:
        """
        Generate embeddings for a value or list of values

        Params:
            value: The value or list of values to embed
            max_seq_length: The maximum sequence length to truncate the input to
            prompt: Optional prompt for prompt-following embedding models.

        Returns:
            A matrix of floats representing the embedding for each value if the input is a list of
                values, or a list of floats representing the embedding for the single value if the
                input is a single value
        """
        payload: EmbedRequest = {
            "values": value if isinstance(value, list) else [value],
            "max_seq_length": max_seq_length,
            "prompt": prompt,
        }
        client = OrcaClient._resolve_client()
        if isinstance(self, PretrainedEmbeddingModel):
            embeddings = client.POST(
                "/gpu/pretrained_embedding_model/{model_name}/embedding",
                params={"model_name": cast(PretrainedEmbeddingModelName, self.name)},
                json=payload,
                timeout=30,  # may be slow in case of cold start
            )
        elif isinstance(self, FinetunedEmbeddingModel):
            embeddings = client.POST(
                "/gpu/finetuned_embedding_model/{name_or_id}/embedding",
                params={"name_or_id": self.id},
                json=payload,
                timeout=30,  # may be slow in case of cold start
            )
        else:
            raise ValueError("Invalid embedding model")
        return embeddings if isinstance(value, list) else embeddings[0]

    @overload
    def evaluate(
        self,
        datasource: Datasource,
        *,
        value_column: str = "value",
        label_column: str,
        score_column: None = None,
        eval_datasource: Datasource | None = None,
        subsample: int | float | None = None,
        neighbor_count: int = 5,
        batch_size: int = 32,
        weigh_memories: bool = True,
        background: Literal[True],
    ) -> Job[ClassificationMetrics]:
        pass

    @overload
    def evaluate(
        self,
        datasource: Datasource,
        *,
        value_column: str = "value",
        label_column: str,
        score_column: None = None,
        eval_datasource: Datasource | None = None,
        subsample: int | float | None = None,
        neighbor_count: int = 5,
        batch_size: int = 32,
        weigh_memories: bool = True,
        background: Literal[False] = False,
    ) -> ClassificationMetrics:
        pass

    @overload
    def evaluate(
        self,
        datasource: Datasource,
        *,
        value_column: str = "value",
        label_column: None = None,
        score_column: str,
        eval_datasource: Datasource | None = None,
        subsample: int | float | None = None,
        neighbor_count: int = 5,
        batch_size: int = 32,
        weigh_memories: bool = True,
        background: Literal[True],
    ) -> Job[RegressionMetrics]:
        pass

    @overload
    def evaluate(
        self,
        datasource: Datasource,
        *,
        value_column: str = "value",
        label_column: None = None,
        score_column: str,
        eval_datasource: Datasource | None = None,
        subsample: int | float | None = None,
        neighbor_count: int = 5,
        batch_size: int = 32,
        weigh_memories: bool = True,
        background: Literal[False] = False,
    ) -> RegressionMetrics:
        pass

    def evaluate(
        self,
        datasource: Datasource,
        *,
        value_column: str = "value",
        label_column: str | None = None,
        score_column: str | None = None,
        eval_datasource: Datasource | None = None,
        subsample: int | float | None = None,
        neighbor_count: int = 5,
        batch_size: int = 32,
        weigh_memories: bool = True,
        background: bool = False,
    ) -> (
        ClassificationMetrics
        | RegressionMetrics
        | Job[ClassificationMetrics]
        | Job[RegressionMetrics]
        | Job[ClassificationMetrics | RegressionMetrics]
    ):
        """
        Evaluate the finetuned embedding model
        """

        payload: EmbeddingEvaluationRequest = {
            "datasource_name_or_id": datasource.id,
            "datasource_label_column": label_column,
            "datasource_value_column": value_column,
            "datasource_score_column": score_column,
            "eval_datasource_name_or_id": eval_datasource.id if eval_datasource is not None else None,
            "subsample": subsample,
            "neighbor_count": neighbor_count,
            "batch_size": batch_size,
            "weigh_memories": weigh_memories,
        }
        client = OrcaClient._resolve_client()
        if isinstance(self, PretrainedEmbeddingModel):
            response = client.POST(
                "/pretrained_embedding_model/{model_name}/evaluation",
                params={"model_name": self.name},
                json=payload,
            )
        elif isinstance(self, FinetunedEmbeddingModel):
            response = client.POST(
                "/finetuned_embedding_model/{name_or_id}/evaluation",
                params={"name_or_id": self.id},
                json=payload,
            )
        else:
            raise ValueError("Invalid embedding model")

        def get_result(job_id: str) -> ClassificationMetrics | RegressionMetrics:
            from .classification_model import ClassificationMetrics
            from .regression_model import RegressionMetrics

            client = OrcaClient._resolve_client()
            if isinstance(self, PretrainedEmbeddingModel):
                res = client.GET(
                    "/pretrained_embedding_model/{model_name}/evaluation/{job_id}",
                    params={"model_name": self.name, "job_id": job_id},
                )["result"]
            elif isinstance(self, FinetunedEmbeddingModel):
                res = client.GET(
                    "/finetuned_embedding_model/{name_or_id}/evaluation/{job_id}",
                    params={"name_or_id": self.id, "job_id": job_id},
                )["result"]
            else:
                raise ValueError("Invalid embedding model")
            assert res is not None
            return RegressionMetrics(res) if "mse" in res else ClassificationMetrics(res)

        job = Job(response["job_id"], lambda: get_result(response["job_id"]))
        return job if background else job.result()


class _ModelDescriptor:
    """
    Descriptor for lazily loading embedding models with IDE autocomplete support.

    This class implements the descriptor protocol to provide lazy loading of embedding models
    while maintaining IDE autocomplete functionality. It delays the actual loading of models
    until they are accessed, which improves startup performance.

    The descriptor pattern works by defining how attribute access is handled. When a class
    attribute using this descriptor is accessed, the __get__ method is called, which then
    retrieves or initializes the actual model on first access.
    """

    def __init__(self, name: str):
        """
        Initialize a model descriptor.

        Args:
            name: The name of the embedding model in PretrainedEmbeddingModelName
        """
        self.name = name
        self.model = None  # Model is loaded lazily on first access

    def __get__(self, instance, owner_class):
        """
        Descriptor protocol method called when the attribute is accessed.

        This method implements lazy loading - the actual model is only initialized
        the first time it's accessed. Subsequent accesses will use the cached model.

        Args:
            instance: The instance the attribute was accessed from, or None if accessed from the class
            owner_class: The class that owns the descriptor

        Returns:
            The initialized embedding model

        Raises:
            AttributeError: If no model with the given name exists
        """
        # When accessed from an instance, redirect to class access
        if instance is not None:
            return self.__get__(None, owner_class)

        # Load the model on first access
        if self.model is None:
            try:
                self.model = PretrainedEmbeddingModel._get(cast(PretrainedEmbeddingModelName, self.name))
            except (KeyError, AttributeError):
                raise AttributeError(f"No embedding model named {self.name}")

        return self.model


class PretrainedEmbeddingModel(EmbeddingModelBase):
    """
    A pretrained embedding model

    **Models:**

    OrcaCloud supports a select number of small to medium sized embedding models that perform well on the
        [Hugging Face MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard).
        These can be accessed as class attributes. We currently support:

    - **`CDE_SMALL`**: Context-aware CDE small model from Hugging Face ([jxm/cde-small-v1](https://huggingface.co/jxm/cde-small-v1))
    - **`CLIP_BASE`**: Multi-modal CLIP model from Hugging Face ([sentence-transformers/clip-ViT-L-14](https://huggingface.co/sentence-transformers/clip-ViT-L-14))
    - **`GTE_BASE`**: Alibaba's GTE model from Hugging Face ([Alibaba-NLP/gte-base-en-v1.5](https://huggingface.co/Alibaba-NLP/gte-base-en-v1.5))
    - **`DISTILBERT`**: DistilBERT embedding model from Hugging Face ([distilbert-base-uncased](https://huggingface.co/distilbert-base-uncased))
    - **`GTE_SMALL`**: GTE-Small embedding model from Hugging Face ([Supabase/gte-small](https://huggingface.co/Supabase/gte-small))
    - **`E5_LARGE`**: E5-Large instruction-tuned embedding model from Hugging Face ([intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct))
    - **`GIST_LARGE`**: GIST-Large embedding model from Hugging Face ([avsolatorio/GIST-large-Embedding-v0](https://huggingface.co/avsolatorio/GIST-large-Embedding-v0))
    - **`MXBAI_LARGE`**: Mixbreas's Large embedding model from Hugging Face ([mixedbread-ai/mxbai-embed-large-v1](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1))
    - **`BGE_BASE`**: BAAI's BGE-Base instruction-tuned embedding model from Hugging Face ([BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5))

    **Instruction Support:**

    Some models support instruction-following for better task-specific embeddings. You can check if a model supports instructions
    using the `supports_instructions` attribute.

    Examples:
        >>> PretrainedEmbeddingModel.CDE_SMALL
        PretrainedEmbeddingModel({name: CDE_SMALL, embedding_dim: 768, max_seq_length: 512})

        >>> # Using instruction with an instruction-supporting model
        >>> model = PretrainedEmbeddingModel.E5_LARGE
        >>> embeddings = model.embed("Hello world", prompt="Represent this sentence for retrieval:")

    Attributes:
        name: Name of the pretrained embedding model
        embedding_dim: Dimension of the embeddings that are generated by the model
        max_seq_length: Maximum input length (in tokens not characters) that this model can process. Inputs that are longer will be truncated during the embedding process
        num_params: Number of parameters in the model
        uses_context: Whether the pretrained embedding model uses context
        supports_instructions: Whether this model supports instruction-following
    """

    # Define descriptors for model access with IDE autocomplete
    CDE_SMALL = _ModelDescriptor("CDE_SMALL")
    CLIP_BASE = _ModelDescriptor("CLIP_BASE")
    GTE_BASE = _ModelDescriptor("GTE_BASE")
    DISTILBERT = _ModelDescriptor("DISTILBERT")
    GTE_SMALL = _ModelDescriptor("GTE_SMALL")
    E5_LARGE = _ModelDescriptor("E5_LARGE")
    GIST_LARGE = _ModelDescriptor("GIST_LARGE")
    MXBAI_LARGE = _ModelDescriptor("MXBAI_LARGE")
    BGE_BASE = _ModelDescriptor("BGE_BASE")

    name: PretrainedEmbeddingModelName

    def __init__(self, metadata: PretrainedEmbeddingModelMetadata):
        # for internal use only, do not document
        self.name = metadata["name"]
        super().__init__(
            name=metadata["name"],
            embedding_dim=metadata["embedding_dim"],
            max_seq_length=metadata["max_seq_length"],
            num_params=metadata["num_params"],
            uses_context=metadata["uses_context"],
            supports_instructions=(
                bool(metadata["supports_instructions"]) if "supports_instructions" in metadata else False
            ),
        )

    def __eq__(self, other) -> bool:
        return isinstance(other, PretrainedEmbeddingModel) and self.name == other.name

    def __repr__(self) -> str:
        num_params_str = f"{self.num_params / 1000000:.0f}M" if self.num_params else "unknown"
        return f"PretrainedEmbeddingModel({{ name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {num_params_str} }})"

    @classmethod
    def all(cls) -> list[PretrainedEmbeddingModel]:
        """
        List all pretrained embedding models in the OrcaCloud

        Returns:
            A list of all pretrained embedding models available in the OrcaCloud
        """
        client = OrcaClient._resolve_client()
        return [cls(metadata) for metadata in client.GET("/pretrained_embedding_model")]

    _instances: dict[str, PretrainedEmbeddingModel] = {}

    @classmethod
    def _get(cls, name: PretrainedEmbeddingModelName) -> PretrainedEmbeddingModel:
        # for internal use only, do not document - we want people to use dot notation to get the model
        cache_key = str(name)
        if cache_key not in cls._instances:
            client = OrcaClient._resolve_client()
            metadata = client.GET(
                "/pretrained_embedding_model/{model_name}",
                params={"model_name": name},
            )
            cls._instances[cache_key] = cls(metadata)
        return cls._instances[cache_key]

    @classmethod
    def open(cls, name: PretrainedEmbeddingModelName) -> PretrainedEmbeddingModel:
        """
        Open an embedding model by name.

        This is an alternative method to access models for environments
        where IDE autocomplete for model names is not available.

        Params:
            name: Name of the model to open (e.g., "GTE_BASE", "CLIP_BASE")

        Returns:
            The embedding model instance

        Examples:
            >>> model = PretrainedEmbeddingModel.open("GTE_BASE")
        """
        try:
            # Always use the _get method which handles caching properly
            return cls._get(name)
        except (KeyError, AttributeError):
            raise ValueError(f"Unknown model name: {name}")

    @classmethod
    def exists(cls, name: str) -> bool:
        """
        Check if a pretrained embedding model exists by name

        Params:
            name: The name of the pretrained embedding model

        Returns:
            True if the pretrained embedding model exists, False otherwise
        """
        return name in get_args(PretrainedEmbeddingModelName)

    @overload
    def finetune(
        self,
        name: str,
        train_datasource: Datasource | LabeledMemoryset | ScoredMemoryset,
        *,
        eval_datasource: Datasource | None = None,
        label_column: str | None = None,
        score_column: str | None = None,
        value_column: str = "value",
        loss: FinetuningLoss = "prediction",
        trial_count: int | None = None,
        hyperparams: EmbeddingFinetuneHyperparams | None = None,
        if_exists: CreateMode = "error",
        background: Literal[True],
        seed: int | None = None,
    ) -> Job[FinetunedEmbeddingModel]:
        pass

    @overload
    def finetune(
        self,
        name: str,
        train_datasource: Datasource | LabeledMemoryset | ScoredMemoryset,
        *,
        eval_datasource: Datasource | None = None,
        label_column: str | None = None,
        score_column: str | None = None,
        value_column: str = "value",
        loss: FinetuningLoss = "prediction",
        trial_count: int | None = None,
        hyperparams: EmbeddingFinetuneHyperparams | None = None,
        if_exists: CreateMode = "error",
        background: Literal[False] = False,
        seed: int | None = None,
    ) -> FinetunedEmbeddingModel:
        pass

    def finetune(  # noqa: C901
        self,
        name: str,
        train_datasource: Datasource | LabeledMemoryset | ScoredMemoryset,
        *,
        eval_datasource: Datasource | None = None,
        label_column: str | None = None,
        score_column: str | None = None,
        value_column: str = "value",
        loss: FinetuningLoss = "prediction",
        trial_count: int | None = None,
        hyperparams: EmbeddingFinetuneHyperparams | None = None,
        if_exists: CreateMode = "error",
        background: bool = False,
        seed: int | None = None,
    ) -> FinetunedEmbeddingModel | Job[FinetunedEmbeddingModel]:
        """
        Finetune an embedding model

        Trains a new embedding model starting from this pretrained base. All
        hyperparameters have sensible loss-specific defaults — in the simplest
        case you only need a name and training data.

        Params:
            name: Name of the finetuned embedding model
            train_datasource: Data to train on
            eval_datasource: Data to evaluate on. When omitted a split is held out from the
                training data automatically.
            label_column: Column name for categorical labels in the datasource
            score_column: Column name for continuous scores in the datasource
            value_column: Column name for the text to embed
            loss: Loss function for training
            trial_count: Number of hyperparameter configurations to try. 1 (default) runs a single
                training job. Values > 1 activate a hyperparameter sweep. The easiest way to sweep
                is to just set trial_count and let the system auto-select which parameters to search
                with sensible loss-specific ranges. For finer control, pass explicit ranges via hyperparams.
            hyperparams: Training hyperparameters to override. All parameters have loss-specific
                defaults so you only need to specify what you want to change. In sweep mode, use
                tuples for continuous ranges and lists for categorical choices.
            if_exists: What to do if a finetuned embedding model with the same name already exists,
                defaults to "error". Other option is "open" to open the existing finetuned embedding model.
            background: Whether to run the operation in the background and return a job handle
            seed: Random seed for reproducibility

        Returns:
            The finetuned embedding model

        Raises:
            ValueError: If the finetuned embedding model already exists and if_exists is "error" or
                if it is "open" but the base model param does not match the existing model
            ValueError: If `train_datasource` is a plain `Datasource` and neither `label_column`
                nor `score_column` is provided

        Examples:
            Minimal single run with default hyperparameters:

            >>> model = PretrainedEmbeddingModel.GTE_BASE
            >>> memoryset = LabeledMemoryset.open("my_memoryset")
            >>> model.finetune("my_model", memoryset)

            Single run with custom hyperparameters:

            >>> datasource = Datasource.open("my_datasource")
            >>> model.finetune("my_model", datasource, label_column="label", loss="contrastive", hyperparams={
            ...     "epochs": 5, "learning_rate": 1e-4, "batch_size": 64,
            ... })

            Default sweep, just set trial_count and the system picks what to search:

            >>> model.finetune("my_model", memoryset, trial_count=9)

            Custom sweep with explicit ranges and choices:

            >>> model.finetune("my_model", memoryset, trial_count=15, hyperparams={
            ...     "learning_rate": (1e-5, 1e-3),
            ...     "batch_size": [32, 64, 128],
            ...     "epochs": 4,
            ... })
        """
        exists = FinetunedEmbeddingModel.exists(name)

        if exists and if_exists == "error":
            raise ValueError(f"Finetuned embedding model '{name}' already exists")
        elif exists and if_exists == "open":
            existing = FinetunedEmbeddingModel.open(name)

            if existing.base_model is not None and existing.base_model.name != self.name:
                raise ValueError(f"Finetuned embedding model '{name}' already exists, but with different base model")

            return existing

        from .memoryset import LabeledMemoryset, ScoredMemoryset

        if score_column is not None and label_column is not None:
            raise ValueError("Only one of label_column or score_column can be provided")
        if isinstance(train_datasource, ScoredMemoryset) or score_column is not None:
            task_type = "regression"
        elif isinstance(train_datasource, LabeledMemoryset) or label_column is not None:
            task_type = "classification"
        else:
            raise ValueError("Must specify either label_column or score_column")

        payload: FinetuneEmbeddingModelRequest = {
            "name": name,
            "base_model": self.name,
            "label_column": label_column,
            "score_column": score_column,
            "value_column": value_column,
        }
        config: EmbeddingFinetuneConfig = {"loss": loss, "task_type": task_type}
        if trial_count is not None:
            config["trial_count"] = trial_count
        if seed is not None:
            config["seed"] = seed
        if hyperparams:
            if "learning_rate" in hyperparams:
                val = hyperparams["learning_rate"]
                config["learning_rate"] = {"min": val[0], "max": val[1]} if isinstance(val, tuple) else val
            if "warmup" in hyperparams:
                val = hyperparams["warmup"]
                config["warmup"] = {"min": val[0], "max": val[1]} if isinstance(val, tuple) else val
            if "epochs" in hyperparams:
                val = hyperparams["epochs"]
                config["epochs"] = {"min": val[0], "max": val[1]} if isinstance(val, tuple) else val
            if "weight_decay" in hyperparams:
                val = hyperparams["weight_decay"]
                config["weight_decay"] = {"min": val[0], "max": val[1]} if isinstance(val, tuple) else val
            if "loss_scale" in hyperparams:
                val = hyperparams["loss_scale"]
                config["loss_scale"] = {"min": val[0], "max": val[1]} if isinstance(val, tuple) else val
            if "max_steps" in hyperparams:
                config["max_steps"] = hyperparams["max_steps"]
            if "batch_size" in hyperparams:
                config["batch_size"] = hyperparams["batch_size"]
            if "learning_rate_scheduler" in hyperparams:
                config["learning_rate_scheduler"] = hyperparams["learning_rate_scheduler"]
            if "normalize_embeddings" in hyperparams:
                config["normalize_embeddings"] = hyperparams["normalize_embeddings"]
            if "max_seq_length" in hyperparams:
                config["max_seq_length"] = hyperparams["max_seq_length"]
            if "truncation_side" in hyperparams:
                config["truncation_side"] = hyperparams["truncation_side"]
            if "early_stopping" in hyperparams:
                config["early_stopping"] = hyperparams["early_stopping"]
            if "early_stopping_threshold" in hyperparams:
                config["early_stopping_threshold"] = hyperparams["early_stopping_threshold"]
            if "eval_method" in hyperparams:
                config["eval_method"] = hyperparams["eval_method"]
            if "eval_steps" in hyperparams:
                config["eval_steps"] = hyperparams["eval_steps"]
            if "neighbor_eval_count" in hyperparams:
                config["neighbor_eval_count"] = hyperparams["neighbor_eval_count"]
            if "neighbor_eval_pool_subsample" in hyperparams:
                config["neighbor_eval_pool_subsample"] = hyperparams["neighbor_eval_pool_subsample"]
        if config:
            payload["config"] = config

        if isinstance(train_datasource, Datasource):
            payload["train_datasource_name_or_id"] = train_datasource.id
        elif isinstance(train_datasource, (LabeledMemoryset, ScoredMemoryset)):
            payload["train_memoryset_name_or_id"] = train_datasource.id
        if eval_datasource is not None:
            payload["eval_datasource_name_or_id"] = eval_datasource.id

        client = OrcaClient._resolve_client()
        res = client.POST(
            "/finetuned_embedding_model",
            json=payload,
        )
        # finetuning_job_id is always present when creating a finetuned model
        finetuning_job_id = res.get("finetuning_job_id")
        assert finetuning_job_id is not None, "finetuning_job_id should be present for finetuned models"
        job = Job(
            finetuning_job_id,
            lambda: FinetunedEmbeddingModel.open(res["id"]),
        )
        return job if background else job.result()


class FinetunedEmbeddingModelTrial:
    """
    A trial for a finetuned embedding model

    Attributes:
        status: The status of the trial
        hyperparameters: The hyperparameters used for the trial
        metrics: The metrics for the trial
        started_at: The start time of the trial
        completed_at: When the trial finished, if known
    """

    status: TrialStatus
    hyperparameters: dict[str, Any]
    metrics: dict[str, float]
    started_at: datetime
    completed_at: datetime | None

    def __init__(self, trial: FinetunedEmbeddingModelTrialResponse):
        self.status = trial["status"]
        self.hyperparameters = trial["hyperparameters"]
        self.metrics = trial["metrics"]
        self.started_at = datetime.fromisoformat(trial["started_at"])
        raw_completed_at = trial.get("completed_at")
        self.completed_at = datetime.fromisoformat(raw_completed_at) if raw_completed_at is not None else None

    def __repr__(self) -> str:
        return (
            "FinetunedEmbeddingModelTrial({\n"
            f"    hyperparameters: {self.hyperparameters},\n"
            f"    metrics: {self.metrics}\n"
            "})"
        )


class FinetunedEmbeddingModel(EmbeddingModelBase):
    """
    A finetuned embedding model in the OrcaCloud

    Attributes:
        name: Name of the finetuned embedding model
        embedding_dim: Dimension of the embeddings that are generated by the model
        max_seq_length: Maximum input length (in tokens not characters) that this model can process.
            Inputs that are longer will be truncated during the embedding process
        uses_context: Whether the model uses the memoryset to contextualize embeddings (acts akin
            to inverse document frequency in TFIDF features)
        id: Unique identifier of the finetuned embedding model
        base_model: Base model the finetuned embedding model was trained on (None for uploaded models)
        created_at: When the model was finetuned
        description: Optional description of the embedding model

    Note:
        For uploaded models (created via `_upload`), `base_model` is None,
        `num_params` may be extracted from the model if possible (otherwise None/0),
        and `supports_instructions` is False since this property cannot be determined
        from the model config alone.
    """

    id: str
    name: str
    created_at: datetime
    updated_at: datetime
    base_model: PretrainedEmbeddingModel | None
    _status: Status | None
    description: str | None

    def __init__(self, metadata: FinetunedEmbeddingModelMetadata):
        # for internal use only, do not document
        self.id = metadata["id"]
        self.name = metadata["name"]
        self.created_at = datetime.fromisoformat(metadata["created_at"])
        self.updated_at = datetime.fromisoformat(metadata["updated_at"])

        # base_model and finetuning_status may be None for uploaded models
        base_model_name = metadata.get("base_model")
        self.base_model = PretrainedEmbeddingModel._get(base_model_name) if base_model_name else None

        finetuning_status = metadata.get("finetuning_status")
        self._status = Status(finetuning_status) if finetuning_status else None

        self.description = metadata.get("description")

        # Get num_params from metadata, fall back to base_model if not present (for backward compatibility)
        num_params = metadata.get("num_params")
        if num_params is None:
            num_params = self.base_model.num_params if self.base_model else None

        super().__init__(
            name=metadata["name"],
            embedding_dim=metadata["embedding_dim"],
            max_seq_length=metadata["max_seq_length"],
            num_params=num_params,
            uses_context=metadata["uses_context"],
            supports_instructions=self.base_model.supports_instructions if self.base_model else False,
        )

    def __eq__(self, other) -> bool:
        return isinstance(other, FinetunedEmbeddingModel) and self.id == other.id

    def __repr__(self) -> str:
        base_model_str = f"PretrainedEmbeddingModel.{self.base_model.name}" if self.base_model else "None"
        return (
            "FinetunedEmbeddingModel({\n"
            f"    name: {self.name},\n"
            f"    embedding_dim: {self.embedding_dim},\n"
            f"    max_seq_length: {self.max_seq_length},\n"
            f"    base_model: {base_model_str}\n"
            "})"
        )

    @classmethod
    def all(cls) -> list[FinetunedEmbeddingModel]:
        """
        List all finetuned embedding model handles in the OrcaCloud

        Returns:
            A list of all finetuned embedding model handles in the OrcaCloud
        """
        client = OrcaClient._resolve_client()
        return [cls(metadata) for metadata in client.GET("/finetuned_embedding_model")]

    @classmethod
    def open(cls, name: str) -> FinetunedEmbeddingModel:
        """
        Get a handle to a finetuned embedding model in the OrcaCloud

        Params:
            name: The name or unique identifier of a finetuned embedding model

        Returns:
            A handle to the finetuned embedding model in the OrcaCloud

        Raises:
            LookupError: If the finetuned embedding model does not exist
        """
        client = OrcaClient._resolve_client()
        metadata = client.GET(
            "/finetuned_embedding_model/{name_or_id}",
            params={"name_or_id": name},
        )
        return cls(metadata)

    @classmethod
    def exists(cls, name_or_id: str) -> bool:
        """
        Check if a finetuned embedding model with the given name or id exists.

        Params:
            name_or_id: The name or id of the finetuned embedding model

        Returns:
            True if the finetuned embedding model exists, False otherwise
        """
        try:
            cls.open(name_or_id)
            return True
        except LookupError:
            return False

    @classmethod
    def drop(cls, name_or_id: str, *, if_not_exists: DropMode = "error", cascade: bool = False):
        """
        Delete the finetuned embedding model from the OrcaCloud

        Params:
            name_or_id: The name or id of the finetuned embedding model
            if_not_exists: What to do if the finetuned embedding model does not exist, defaults to `"error"`.
                Other option is `"ignore"` to do nothing if the model does not exist.
            cascade: If True, also delete all associated memorysets and their predictive models.
                Defaults to False.

        Raises:
            LookupError: If the finetuned embedding model does not exist and `if_not_exists` is `"error"`
            RuntimeError: If the model has associated memorysets and cascade is False
        """
        try:
            client = OrcaClient._resolve_client()
            client.DELETE(
                "/finetuned_embedding_model/{name_or_id}",
                params={"name_or_id": name_or_id, "cascade": cascade},
            )
        except LookupError:
            if if_not_exists == "error":
                raise

    @property
    def trials(self) -> list[FinetunedEmbeddingModelTrial]:
        """
        List the trials for the finetuned embedding model

        Returns:
            A list of finetuned embedding model trials
        """
        client = OrcaClient._resolve_client()
        return [
            FinetunedEmbeddingModelTrial(trial)
            for trial in client.GET(
                "/finetuned_embedding_model/{name_or_id}/trials",
                params={"name_or_id": self.id},
            )
        ]

    @classmethod
    def _upload(
        cls,
        name: str,
        model_path: str | PathLike,
        *,
        if_exists: CreateMode = "error",
        description: str | None = None,
    ) -> FinetunedEmbeddingModel:
        """
        Upload a pretrained embedding model from a local directory.

        This method uploads a Sentence Transformers compatible model directory to OrcaCloud.
        The directory should contain at minimum a config.json file that allows
        extraction of embedding dimension and max sequence length.

        Note: This endpoint requires root API key access.

        Params:
            name: Name for the uploaded embedding model (must be unique)
            model_path: Path to a local directory containing the model files
                (config.json, model weights, tokenizer files, etc.)
            if_exists: What to do if an embedding model with the same name already exists,
                defaults to `"error"`. Other option is `"open"` to open the existing model.
            description: Optional description for the embedding model

        Returns:
            A handle to the uploaded embedding model in OrcaCloud

        Raises:
            ValueError: If the model already exists and if_exists is `"error"`
            ValueError: If the model_path is not a directory
            LookupError: If the model already exists and if_exists is `"open"` but opening fails

        Examples:
            >>> model = FinetunedEmbeddingModel._upload("my_model", "/path/to/model")
            >>> model.embed("Hello world")
        """
        # Check if model already exists and handle accordingly
        if cls.exists(name):
            if if_exists == "error":
                raise ValueError(f"Embedding model with name '{name}' already exists")
            elif if_exists == "open":
                return cls.open(name)

        model_path = Path(model_path)
        if not model_path.is_dir():
            raise ValueError(f"model_path must be a directory, got: {model_path}")

        # Collect all files in the model directory
        file_paths = [f for f in model_path.rglob("*") if f.is_file()]
        if not file_paths:
            raise ValueError(f"No files found in model directory: {model_path}")

        files: list[tuple[Literal["files"], FileTypes]] = []
        opened_files: list = []

        # Calculate total size for progress bar
        total_size = sum(f.stat().st_size for f in file_paths)

        try:
            with tqdm(total=total_size, unit="B", unit_scale=True, desc="upload") as pbar:
                for file_path in file_paths:
                    # Preserve relative path structure within the model directory
                    relative_path = file_path.relative_to(model_path)
                    buffered_reader = open(file_path, "rb")
                    opened_files.append(buffered_reader)
                    tqdm_reader = TqdmFileReader(buffered_reader, pbar)
                    files.append(("files", (relative_path.as_posix(), cast(bytes, tqdm_reader))))

                client = OrcaClient._resolve_client()
                data: PostEmbeddingModelUploadRequest = {"name": name}
                if description is not None:
                    data["description"] = description
                metadata = client.POST(
                    "/embedding_model/upload",
                    files=files,
                    data=data,
                    timeout=None,
                )
        finally:
            for f in opened_files:
                f.close()

        return cls(metadata)
