from __future__ import annotations


import enum
import json
import time
import asyncio
import contextlib
from lzl.api import aiohttpx
from typing import Optional, Type, Any, Union, List, Dict, Iterator, AsyncIterator, Generator, AsyncGenerator, TYPE_CHECKING

from lzo.types import (
    Literal, 
    eproperty, 
    BaseModel,
    validator, 
    root_validator, 
    field_validator,
    Field,
    PYDANTIC_VERSION, 
)
from lzl.api.openai.version import DEFAULT_COMPLETION_MODEL
from lzl.api.openai.types.handlers import ModelContextHandler
from lzl.api.openai.types.base import BaseResource, Usage
from lzl.api.openai.types.responses import BaseResponse
from lzl.api.openai.types.routes import BaseRoute
from lzl.api.openai.utils import logger, parse_stream, aparse_stream


__all__ = [
    'CompletionChoice',
    'CompletionObject',
    'CompletionResponse',
    'CompletionRoute',
]


class CompletionKind(str, enum.Enum):
    # Only have one choice
    CONTENT = 'content'

class StreamedCompletionChoice(BaseResource):
    kind: CompletionKind = CompletionKind.CONTENT
    value: str


class CompletionChoice(BaseResource):
    text: str
    index: int
    logprobs: Optional[Any] = None
    finish_reason: Optional[str] = None



class CompletionObject(BaseResource):
    prompt: Union[List[str], str] = '<|endoftext|>'
    model: Optional[str] = DEFAULT_COMPLETION_MODEL
    # prompt: Optional[Union[List[str], str]] = '<|endoftext|>'
    suffix: Optional[str] = None
    max_tokens: Optional[int] = 16
    temperature: Optional[float] = 1.0
    top_p: Optional[float] = 1.0
    n: Optional[int] = 1
    stream: Optional[bool] = False
    logprobs: Optional[int] = None
    echo: Optional[bool] = False
    stop: Optional[Union[List[str], str]] = None
    presence_penalty: Optional[float] = 0.0
    frequency_penalty: Optional[float] = 0.0
    best_of: Optional[int] = 1
    logit_bias: Optional[Dict[str, float]] = None
    user: Optional[str] = None

    validate_max_tokens: Optional[bool] = Field(default = False, exclude = True)
    # validate_model_aliases: Optional[bool] = Field(default = False, exclude = True)

    # @validator('model', pre=True, always=True)
    @field_validator('model')
    def validate_model(cls, v, values: Dict[str, Any]) -> str:
        """
        Validate the model
        """
        if not v:
            if values.get('engine'):
                v = values.get('engine')
            elif values.get('deployment'):
                v = values.get('deployment')
        with contextlib.suppress(KeyError):
            v = ModelContextHandler.resolve_model_name(v)
        # if values.get('validate_model_aliases', False):
        #     v = ModelContextHandler[v].name
        return v

    
    @validator('temperature')
    def validate_temperature(cls, v: float) -> float:
        """
        Min Temperature is 0.0
        https://beta.openai.com/docs/api-reference/completions/create#completions/create-temperature
        """
        return None if v is None else max(0.0, v)
    
    @validator('top_p')
    def validate_top_p(cls, v: float) -> float:
        """
        Min Top Probability is 0.0
        https://beta.openai.com/docs/api-reference/completions/create#completions/create-top-probability
        """
        return None if v is None else max(0.0, v)

    @validator('logprobs')
    def validate_logprobs(cls, v: int) -> int:
        """
        Max logprobs is 5
        https://beta.openai.com/docs/api-reference/completions/create#completions/create-logprobs
        """
        return None if v is None else max(0, min(v, 5))
    
    @validator('presence_penalty')
    def validate_presence_penalty(cls, v: float) -> float:
        """
        Min Presence Penalty is -2.0, Max is 2.0
        https://beta.openai.com/docs/api-reference/completions/create#completions/create-presence-penalty
        """
        return None if v is None else max(0.0, min(v, 2.0))
    
    @validator('frequency_penalty')
    def validate_frequency_penalty(cls, v: float) -> float:
        """
        Min Frequency Penalty is -2.0, Max is 2.0
        https://beta.openai.com/docs/api-reference/completions/create#completions/create-frequency-penalty
        """
        return None if v is None else max(0.0, min(v, 2.0))
    

    def dict(self, *args, exclude: Any = None, **kwargs):
        """
        Returns the dict representation of the response
        """
        return super().dict(*args, exclude = exclude, **kwargs)


    

class CompletionResponse(BaseResponse):
    choices: Optional[List[CompletionChoice]] = None
    choice_model: Optional[Type[BaseResource]] = CompletionChoice
    input_object: Optional[CompletionObject] = None

    @eproperty
    def text(self) -> str:
        """
        Returns the text for the completions
        """
        if self.choices_results:
            return ''.join([choice.text for choice in self.choices])
        return self.response.text
    
    @eproperty
    def openai_model(self):
        """
        Returns the model for the completions
        """
        return self.headers.get('openai-model', self.completion_model)

    @eproperty
    def completion_model(self):
        """
        Returns the model for the completions
        """
        return self.input_object.model or None

    def _validate_usage(self):
        """
        Validate usage
        """
        if self.usage and self.usage.total_tokens: return
        if self.response.status_code == 200:
            self.usage = Usage(
                prompt_tokens = ModelContextHandler.count_tokens(self.input_object.prompt),
                completion_tokens = ModelContextHandler.count_tokens(self.text),
            )
            self.usage.total_tokens = self.usage.prompt_tokens + self.usage.completion_tokens

    @eproperty
    def consumption(self) -> int:
        """
        Returns the consumption for the completions
        """ 
        self._validate_usage()
        return ModelContextHandler.get_consumption_cost(
            model_name = self.openai_model,
            usage = self.usage,
        )

    def dict(self, *args, exclude: Any = None, **kwargs):
        """
        Returns the dict representation of the response
        """
        return super().dict(*args, exclude = exclude, **kwargs)

    def parse_stream_item(self, item: Union[Dict, Any], **kwargs) -> Optional[StreamedCompletionChoice]:
        """
        Parses a single stream item
        """
        choice = item['choices'][0]
        if choice['finish_reason'] == 'stop':
            return None
        return StreamedCompletionChoice(
            value = choice['text']
        )

    def handle_stream(
        self,
        response: aiohttpx.Response,
        streaming: Optional[bool] = False,
    ) -> Iterator[Dict]:
        """
        Handles the streaming response
        """
        texts = {}
        for line in parse_stream(response):
            try:
                item = json.loads(line)
                if streaming:
                    yield item
                self.handle_stream_metadata(item)
                for n, choice in enumerate(item['choices']):
                    if not texts.get(n):
                        texts[n] = {
                            'index': choice['index'],
                            'text': choice['text'],
                        }
                        self.usage.completion_tokens += 1
                    elif choice['finish_reason'] != 'stop':
                        texts[n]['text'] += choice['text']
                        self.usage.completion_tokens += 1

                    else:
                        texts[n]['finish_reason'] = choice['finish_reason']
                        compl = texts.pop(n)
                        if streaming:
                            self.handle_resource_item(item = compl)
                        else:
                            yield compl

            except Exception as e:
                logger.error(f'Error: {line}: {e}')

        self._stream_consumed = True
        for remaining_text in texts.values():
            if streaming:
                self.handle_resource_item(item = remaining_text)
            else:
                yield remaining_text
        self.usage.total_tokens = self.usage.completion_tokens
    

    async def ahandle_stream(
        self,
        response: aiohttpx.Response,
        streaming: Optional[bool] = False,
    ) -> AsyncIterator[Dict]:
        """
        Handles the streaming response
        """
        texts = {}
        async for line in aparse_stream(response):
            try:
                item = json.loads(line)
                if streaming:
                    yield item
                self.handle_stream_metadata(item)
                for n, choice in enumerate(item['choices']):
                    if not texts.get(n):
                        texts[n] = {
                            'index': choice['index'],
                            'text': choice['text'],
                        }
                        self.usage.completion_tokens += 1
                    elif choice['finish_reason'] != 'stop':
                        texts[n]['text'] += choice['text']
                        self.usage.completion_tokens += 1

                    else:
                        texts[n]['finish_reason'] = choice['finish_reason']
                        compl = texts.pop(n)
                        if streaming:
                            self.handle_resource_item(item = compl)
                        else:
                            yield compl

            except Exception as e:
                logger.error(f'Error: {line}: {e}')

        self._stream_consumed = True
        for remaining_text in texts.values():
            if streaming:
                self.handle_resource_item(item = remaining_text)
            else:
                yield remaining_text
        self.usage.total_tokens = self.usage.completion_tokens

    if TYPE_CHECKING:
        def stream(self, **kwargs) -> Generator[StreamedCompletionChoice, None, None]:
            """
            Streams the response
            """
            ...

        async def astream(self, **kwargs) -> AsyncGenerator[StreamedCompletionChoice, None]:
            """
            Streams the response
            """
            ...
    

class CompletionRoute(BaseRoute):
    input_model: Optional[Type[BaseResource]] = CompletionObject
    response_model: Optional[Type[BaseResource]] = CompletionResponse

    api_resource: Optional[str] = Field(default = 'completions')
    root_name: Optional[str] = Field(default = 'completion')

    # @eproperty
    # def api_resource(self):
    #     return 'completions'

    # @eproperty
    # def root_name(self):
    #     return 'completion'
    
    def create(
        self, 
        input_object: Optional[CompletionObject] = None,
        parse_stream: Optional[bool] = True,
        timeout: Optional[int] = None,
        auto_retry: Optional[bool] = False,
        auto_retry_limit: Optional[int] = None,
        header_cache_keys: Optional[List[str]] = None,
        **kwargs
    ) -> CompletionResponse:
        """
        Creates a completion for the provided prompt and parameters

        Usage:

        ```python
        >>> result = OpenAI.completions.create(
        >>>    prompt = 'say this is a test',
        >>>    max_tokens = 4,
        >>>    stream = True
        >>> )
        ```

        **Parameters:**

        :model (required): ID of the model to use. You can use the List models API 
        to see all of your available models,  or see our Model overview for descriptions of them.
        Default: `text-curie-001`
        
        :prompt (optional): The prompt(s) to generate completions for, encoded as 
        a string, array of strings, array of tokens, or array of token arrays.
        Note that `<|endoftext|>` is the document separator that the model sees during 
        training, so if a prompt is not specified the model will generate as if from 
        the beginning of a new document.
        Default: `<|endoftext|>`

        :suffix (optional): The suffix that comes after a completion of inserted text.
        Default: `None`

        :max_tokens (optional): The maximum number of tokens to generate in the completion.
        The token count of your prompt plus `max_tokens` cannot exceed the model's context length. 
        Most models have a context length of 2048 tokens (except for the newest models, which 
        support 4096).
        Default: `16`

        :temperature (optional): What sampling temperature to use. Higher values means 
        the model will take more risks. Try 0.9 for more creative applications, and 0 (argmax sampling) 
        for ones with a well-defined answer. We generally recommend altering this or `top_p` but not both.
        Default: `1.0`

        :top_p (optional): An alternative to sampling with `temperature`, called nucleus 
        sampling, where the model considers the results of the tokens with `top_p` probability mass. 
        So `0.1` means only  the tokens comprising the top 10% probability mass are considered.
        We generally recommend altering this or `temperature` but not both
        Default: `1.0`

        :n (optional): How many completions to generate for each prompt.
        Note: Because this parameter generates many completions, it can quickly 
        consume your token quota. Use carefully and ensure that you have reasonable 
        settings for `max_tokens` and stop.
        Default: `1`

        :stream (optional): Whether to stream back partial progress. 
        If set, tokens will be sent as data-only server-sent events as they become 
        available, with the stream terminated by a `data: [DONE]` message. This is 
        handled automatically by the Client and enables faster response processing.
        Default: `False`

        :logprobs (optional): Include the log probabilities on the `logprobs` 
        most likely tokens, as well the chosen tokens. For example, if `logprobs` is 5, 
        the API will return a list of the 5 most likely tokens. The API will always 
        return the logprob of the sampled token, so there may be up to `logprobs+1` 
        elements in the response. The maximum value for `logprobs` is 5.
        Default: `None`

        :echo (optional): Echo back the prompt in addition to the completion
        Default: `False`

        :stop (optional): Up to 4 sequences where the API will stop generating 
        further tokens. The returned text will not contain the stop sequence.
        Default: `None`

        :presence_penalty (optional): Number between `-2.0` and `2.0`. Positive values 
        penalize new tokens based on whether they appear in the text so far, increasing the 
        model's likelihood to talk about new topics
        Default: `0.0`

        :frequency_penalty (optional): Number between `-2.0` and `2.0`. Positive 
        values penalize new tokens based on their existing frequency in the text so 
        far, decreasing the model's likelihood to repeat the same line verbatim.
        Default: `0.0`

        :best_of (optional): Generates `best_of` completions server-side and returns 
        the "best" (the one with the highest log probability per token). Results cannot be streamed.
        When used with `n`, `best_of` controls the number of candidate completions and n specifies how 
        many to return – `best_of` must be greater than `n`.
        Note: Because this parameter generates many completions, it can quickly consume your token quota. 
        Use carefully and ensure that you have reasonable settings for `max_tokens` and `stop`.
        Default: `1`

        :logit_bias (optional): Modify the likelihood of specified tokens appearing in the completion.
        Accepts a json object that maps tokens (specified by their token ID in the GPT tokenizer) to an associated 
        bias value from -100 to 100. You can use this tokenizer tool (which works for both GPT-2 and GPT-3) to 
        convert text to token IDs. Mathematically, the bias is added to the logits generated by the model prior 
        to sampling. The exact effect will vary per model, but values between -1 and 1 should decrease or increase 
        likelihood of selection; values like -100 or 100 should result in a ban or exclusive selection of the 
        relevant token.
        As an example, you can pass `{"50256": -100}` to prevent the `<|endoftext|>` token from being generated.
        Default: `None`

        :user (optional): A unique identifier representing your end-user, which can help OpenAI to 
        monitor and detect abuse.
        Default: `None`

        :auto_retry (optional): Whether to automatically retry the request if it fails due to a rate limit error.

        :auto_retry_limit (optional): The maximum number of times to retry the request if it fails due to a rate limit error.

        Returns: `CompletionResult`
        """
        return super().create(
            input_object = input_object,
            parse_stream = parse_stream,
            timeout = timeout,
            auto_retry = auto_retry,
            auto_retry_limit = auto_retry_limit,
            header_cache_keys = header_cache_keys,
            **kwargs
        )


    async def async_create(
        self, 
        input_object: Optional[CompletionObject] = None,
        parse_stream: Optional[bool] = True,
        timeout: Optional[int] = None,
        auto_retry: Optional[bool] = False,
        auto_retry_limit: Optional[int] = None,
        header_cache_keys: Optional[List[str]] = None,
        **kwargs
    ) -> CompletionResponse:
        """
        Creates a completion for the provided prompt and parameters

        Usage:

        ```python
        >>> result = await OpenAI.completions.async_create(
        >>>    prompt = 'say this is a test',
        >>>    max_tokens = 4,
        >>>    stream = True
        >>> )
        ```

        **Parameters:**

        :model (required): ID of the model to use. You can use the List models API 
        to see all of your available models,  or see our Model overview for descriptions of them.
        Default: `text-curie-001`
        
        :prompt (optional): The prompt(s) to generate completions for, encoded as 
        a string, array of strings, array of tokens, or array of token arrays.
        Note that `<|endoftext|>` is the document separator that the model sees during 
        training, so if a prompt is not specified the model will generate as if from 
        the beginning of a new document.
        Default: `<|endoftext|>`

        :suffix (optional): The suffix that comes after a completion of inserted text.
        Default: `None`

        :max_tokens (optional): The maximum number of tokens to generate in the completion.
        The token count of your prompt plus `max_tokens` cannot exceed the model's context length. 
        Most models have a context length of 2048 tokens (except for the newest models, which 
        support 4096).
        Default: `16`

        :temperature (optional): What sampling temperature to use. Higher values means 
        the model will take more risks. Try 0.9 for more creative applications, and 0 (argmax sampling) 
        for ones with a well-defined answer. We generally recommend altering this or `top_p` but not both.
        Default: `1.0`

        :top_p (optional): An alternative to sampling with `temperature`, called nucleus 
        sampling, where the model considers the results of the tokens with `top_p` probability mass. 
        So `0.1` means only  the tokens comprising the top 10% probability mass are considered.
        We generally recommend altering this or `temperature` but not both
        Default: `1.0`

        :n (optional): How many completions to generate for each prompt.
        Note: Because this parameter generates many completions, it can quickly 
        consume your token quota. Use carefully and ensure that you have reasonable 
        settings for `max_tokens` and stop.
        Default: `1`

        :stream (optional): Whether to stream back partial progress. 
        If set, tokens will be sent as data-only server-sent events as they become 
        available, with the stream terminated by a `data: [DONE]` message. This is 
        handled automatically by the Client and enables faster response processing.
        Default: `False`

        :logprobs (optional): Include the log probabilities on the `logprobs` 
        most likely tokens, as well the chosen tokens. For example, if `logprobs` is 5, 
        the API will return a list of the 5 most likely tokens. The API will always 
        return the logprob of the sampled token, so there may be up to `logprobs+1` 
        elements in the response. The maximum value for `logprobs` is 5.
        Default: `None`

        :echo (optional): Echo back the prompt in addition to the completion
        Default: `False`

        :stop (optional): Up to 4 sequences where the API will stop generating 
        further tokens. The returned text will not contain the stop sequence.
        Default: `None`

        :presence_penalty (optional): Number between `-2.0` and `2.0`. Positive values 
        penalize new tokens based on whether they appear in the text so far, increasing the 
        model's likelihood to talk about new topics
        Default: `0.0`

        :frequency_penalty (optional): Number between `-2.0` and `2.0`. Positive 
        values penalize new tokens based on their existing frequency in the text so 
        far, decreasing the model's likelihood to repeat the same line verbatim.
        Default: `0.0`

        :best_of (optional): Generates `best_of` completions server-side and returns 
        the "best" (the one with the highest log probability per token). Results cannot be streamed.
        When used with `n`, `best_of` controls the number of candidate completions and n specifies how 
        many to return – `best_of` must be greater than `n`.
        Note: Because this parameter generates many completions, it can quickly consume your token quota. 
        Use carefully and ensure that you have reasonable settings for `max_tokens` and `stop`.
        Default: `1`

        :logit_bias (optional): Modify the likelihood of specified tokens appearing in the completion.
        Accepts a json object that maps tokens (specified by their token ID in the GPT tokenizer) to an associated 
        bias value from -100 to 100. You can use this tokenizer tool (which works for both GPT-2 and GPT-3) to 
        convert text to token IDs. Mathematically, the bias is added to the logits generated by the model prior 
        to sampling. The exact effect will vary per model, but values between -1 and 1 should decrease or increase 
        likelihood of selection; values like -100 or 100 should result in a ban or exclusive selection of the 
        relevant token.
        As an example, you can pass `{"50256": -100}` to prevent the `<|endoftext|>` token from being generated.
        Default: `None`

        :user (optional): A unique identifier representing your end-user, which can help OpenAI to 
        monitor and detect abuse.
        Default: `None`

        :auto_retry (optional): Whether to automatically retry the request if it fails due to a rate limit error.

        :auto_retry_limit (optional): The maximum number of times to retry the request if it fails due to a rate limit error.

        Returns: `CompletionResult`
        """
        return await super().async_create(
            input_object = input_object,
            parse_stream = parse_stream,
            timeout = timeout,
            auto_retry = auto_retry,
            auto_retry_limit = auto_retry_limit,
            header_cache_keys = header_cache_keys,
            **kwargs
        )
    
