"""Qdrant 向量存储模块的数据模型 (M6)。"""

from __future__ import annotations

from datetime import datetime
from typing import Any

from pydantic import BaseModel, Field


class SearchFilter(BaseModel):
    """可选检索过滤条件,全部为 AND 关系。"""

    source_id: str | None = None
    source_ids: list[str] | None = None
    stock_codes: list[str] | None = Field(default=None, description="match any")
    company_names: list[str] | None = Field(default=None, description="match any")
    industries: list[str] | None = Field(default=None, description="match any")
    sentiment: str | None = None          # positive / neutral / negative
    importance_min: int | None = None      # >= N
    event_types: list[str] | None = None   # match any
    publish_date_from: str | None = None   # YYYY-MM-DD
    publish_date_to: str | None = None     # YYYY-MM-DD


class SearchResult(BaseModel):
    """单条检索结果。"""

    url_hash: str
    score: float
    title: str
    url: str
    source_id: str
    publish_time: datetime | None = None
    event: dict[str, Any] | None = None   # EventExtraction 展开的 dict
    char_count: int | None = None
    word_count: int | None = None

    def short_summary(self) -> str:
        codes = (
            ",".join((self.event or {}).get("stock_codes", []))
            if self.event else "-"
        )
        return f"[{self.source_id}] score={self.score:.4f} 《{self.title[:40]}》 {codes}"


class CollectionInfo(BaseModel):
    """Collection 概览信息。"""

    name: str
    exists: bool
    vectors_count: int
    indexed_vectors_count: int | None = None
    segments_count: int | None = None