from pydantic import BaseModel, Field, model_validator


class CampaignRequest(BaseModel):
    text: str
    campaign_id: str | None = None


class CampaignVector(BaseModel):
    features: dict[str, float] = Field(default_factory=dict)


class StrategyEvaluateRequest(BaseModel):
    campaign_text: str = Field(min_length=1)
    centroids: list[list[float]] = Field(default_factory=list)
    covariances: list[list[list[float]]] = Field(default_factory=list)
    cluster_top_doctors: dict[str, list[str]] | None = None
    rejection_distance_threshold: float | None = Field(default=None, gt=0.0)


class HeatmapRequest(BaseModel):
    campaign_vector: list[float] = Field(min_length=1)
    centroids: list[list[float]] = Field(min_length=1)
    covariances: list[list[list[float]]] = Field(min_length=1)
    cluster_top_doctors: dict[str, list[str]] | None = None

    @model_validator(mode="after")
    def validate_shapes(self) -> "HeatmapRequest":
        if len(self.centroids) != len(self.covariances):
            raise ValueError("centroids and covariances length mismatch")

        expected_dim = len(self.campaign_vector)
        for idx, centroid in enumerate(self.centroids):
            if len(centroid) != expected_dim:
                raise ValueError(
                    f"centroid at index {idx} must have {expected_dim} dimensions"
                )

            covariance = self.covariances[idx]
            if len(covariance) != expected_dim:
                raise ValueError(
                    f"covariance at index {idx} must have {expected_dim} rows"
                )
            for row_idx, row in enumerate(covariance):
                if len(row) != expected_dim:
                    raise ValueError(
                        f"covariance at index {idx} row {row_idx} must have {expected_dim} columns"
                    )
        return self


class CampaignMemoryStoreRequest(BaseModel):
    campaign_text: str = Field(min_length=1)
    campaign_id: str | None = None
    outcome: str = ""
    success_score: float = Field(default=0.0, ge=0.0, le=1.0)
    is_successful: bool | None = None
    cluster_id: int | None = None
    extra_metadata: dict[str, str | int | float | bool | None] = Field(default_factory=dict)