from pydantic import BaseModel
from typing import List,Optional,Any,Dict
from .request import TextData

class Response(BaseModel):
    status: str

################################################

class VectorData(BaseModel):
    text: str
    vector: List[float]
    
class EmbeddingResponse(Response):
    vectors: List[VectorData]
    
################################################

class SearchResult(BaseModel):
    id: str
    title: str
    author: str
    site_name: str
    site_domain: str
    pub_date: str
    article_text: str
    score: float
    url: str

class SearchResultMeta(BaseModel):
    id: str
    title: str
    pub_date: str
    url: str

class SearchResultContent(BaseModel):
    id: str
    title: str
    article_text: str

class SearchResponse(Response):
    results: List[SearchResult]
    
class SearchResponseV2(BaseModel):
    meta: List[SearchResultMeta]
    content: List[SearchResultContent]
    
################################################

    
class GetDocumentResponse(Response):
    title: str
    text: str
    url: Optional[str]
    
class GetMDocumentResponse(Response):
    results: List[GetDocumentResponse]
    
################################################

################################################
SUCCESS = 200
FAILED = 500
PROCESSING = 202
class StandardResponse(BaseModel):
    code: int = 200
    message: str = "success"
    meta: Dict = {}
    records: List[Any] = []
    total: Optional[int] = None

    def __init__(self, code=SUCCESS, success=True, err_code=-1, err_message="", meta={}, records=None, message="success", total=None):
        super().__init__()
        self.code = code
        self.message = message
        self.meta = meta
        if records:
            recordsItem = []
            for item in records:
                if isinstance(item, dict):
                    recordsItem.append(item)
                    continue
                item_dict = {}
                for key in item.__dict__:
                    if not key.startswith("_"):
                        value = item.__dict__[key]
                        if isinstance(value, str):
                            item_dict[key] = value
                        elif isinstance(value, list):
                            values = []
                            for v in value:
                                values.append(v)
                            item_dict[key] = values
                        else:
                            item_dict[key] = value
                recordsItem.append(item_dict)
            self.records = recordsItem

        if success == False:
            self.code = FAILED
        if err_code >= 0:
            self.code = err_code
            self.message = err_message
        self.total = total
        
        

            
    
class TaskResponse(StandardResponse):
    task_id: str
    status: str
    result: Optional[Any] = None