from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional

import dacite


class Provider(Enum):
    OPENAI = "OPENAI"
    COHERE = "COHERE"
    UNKNOWN = "UNKNOWN"


class FinishReason(Enum):
    LENGTH = "LENGTH"
    STOP = "STOP"
    UNKNOWN = "UNKNOWN"


@dataclass
class TokenLogProbs:
    token: str
    logprob: Optional[float]
    top_logprobs: Optional[Dict[str, float]] = field(repr=False)
    text_offset: int = field(repr=False)


@dataclass
class LogProbs:
    tokens: List[TokenLogProbs] = field(repr=False)
    likelihood: float


@dataclass
class Completion:
    id: str = field(repr=False)
    external_id: Optional[str] = field(repr=False)
    model_version_id: str = field(repr=False)
    text: str
    finish_reason: FinishReason = field(repr=False)
    logprobs: Optional[LogProbs] = field(repr=False)


@dataclass
class CompletionResultData:
    completions: List[Completion]


@dataclass
class CompletionResultError:
    message: str


@dataclass
class CompletionResult:
    data: Optional[CompletionResultData]
    error: Optional[CompletionResultError]


@dataclass
class GenerateResult:
    # The completions generated by the LLM provider, normalized to Vellum's schema.
    results: List[CompletionResult]

    @classmethod
    def from_raw(cls, raw_result: dict) -> GenerateResult:
        return dacite.from_dict(data_class=cls, data=raw_result, config=dacite.Config(cast=[Provider, FinishReason]))

    @property
    def texts(self) -> List[str]:
        return [
            completion.text
            for result in self.results
            for completion in (result.data.completions if result.data else [])
        ]

    @property
    def text(self) -> str:
        if len(self.texts) != 1:
            raise ValueError(f"Expected exactly one completion, but got {len(self.texts)}")

        return self.texts[0]


@dataclass
class GenerateRequest:
    # A mapping between prompt template variables and the values they should be substituted with
    input_values: Dict[str, Any]
    # Optionally provide an external ID for each prompt. This ID will be returned in the result
    # and can later be used to uniquely identify the generation request that produced the result.
    external_ids: Optional[List[str]] = None


@dataclass
class CompletionActual:
    # The id of the completion request, as originally returned by Vellum when making the Generate call.
    # You must provide either this or external_id.
    id: Optional[str] = None
    # The external id of the completion request, as originally provided by you when making the Generate call.
    # You must provide either this or id.
    external_id: Optional[str] = None
    # The quality of the completion, as a number between 0 and 1. 1 is good, 0 is bad.
    quality: Optional[float] = None
    # The desired output text of the completion. You would typically include the original completion text if also
    # passing quality = 1.0. If you deem quality to be < 1.0, you might include what you wish the
    # generation had produced.
    text: Optional[str] = None
    # Optionally include a timestamp representing when this actual was determined.
    timestamp: Optional[datetime] = None

    def __post_init__(self) -> None:
        if self.id is None and self.external_id is None:
            raise ValueError("Either id or external_id must be provided")

        if self.text is None and self.quality is None:
            raise ValueError("Either text and/or quality must be provided")

        if self.quality is not None and (self.quality < 0 or self.quality > 1):
            raise ValueError("quality must be between 0 and 1")


@dataclass
class SubmitCompletionActualsResult:
    success: bool


@dataclass
class SearchWeights:
    semantic_similarity: float = 0.8
    keywords: float = 0.2

    def __post_init__(self) -> None:
        if sum([self.semantic_similarity, self.keywords]) != 1.0:
            raise ValueError("Weights must sum to 1.0")


@dataclass
class ResultMerging:
    enabled: bool


@dataclass
class SearchFilters:
    external_ids: Optional[List[str]] = None


@dataclass
class SearchOptions:
    limit: Optional[int] = 3
    weights: Optional[SearchWeights] = field(default_factory=SearchWeights)
    result_merging: Optional[ResultMerging] = None
    filters: Optional[SearchFilters] = None


@dataclass
class SearchResultDocument:
    id: str
    label: str
    external_id: Optional[str] = None


@dataclass
class SearchResult:
    document: SearchResultDocument = field(repr=False)
    text: str
    keywords: List[str] = field(repr=False)
    score: float = field(repr=False)


@dataclass
class SearchResults:
    results: List[SearchResult]

    @classmethod
    def from_raw(cls, raw_result: dict) -> SearchResults:
        return dacite.from_dict(data_class=cls, data=raw_result)


@dataclass
class DocumentUploadResult:
    document_id: str

    @classmethod
    def from_raw(cls, raw_result: dict) -> DocumentUploadResult:
        return dacite.from_dict(data_class=cls, data=raw_result)
