import json
import logging
import time
from dataclasses import dataclass
from typing import Optional, Generator, Any, cast, Tuple

from . import api_support
from .completions import PromptTemplates, CompletionResponse, CompletionChunk, PromptTemplateWithMetadata, \
    ChatCompletionResponse, ChatMessage
from .errors import TemplateNotFoundError, APIKeyMissingError, AuthorizationError, FreeplayError
from .flavors import Flavor, ChatFlavor
from .llm_parameters import LLMParameters

JsonDom = dict[str, Any]
Variables = dict[str, str]

logger = logging.getLogger(__name__)
default_tag = 'latest'


class CallSupport:
    def __init__(
            self,
            freeplay_api_key: str,
            api_base: str,
            **kwargs: Any
    ) -> None:
        self.api_base = api_base
        self.freeplay_api_key = freeplay_api_key
        self.client_params = LLMParameters(kwargs)

    @staticmethod
    def find_template_by_name(prompts: PromptTemplates, template_name: str) -> PromptTemplateWithMetadata:
        templates = [t for t in prompts.templates if t.name == template_name]
        if len(templates) == 0:
            raise TemplateNotFoundError(f'Could not find template with name "{template_name}"')
        return templates[0]

    def create_session(self, project_id: str, tag: str, test_run_id: Optional[str] = None) -> JsonDom:
        request_body = {'test_run_id': test_run_id} if test_run_id is not None else None
        response = api_support.post_raw(api_key=self.freeplay_api_key,
                                        url=f'{self.api_base}/projects/{project_id}/sessions/tag/{tag}',
                                        payload=request_body)

        if response.status_code == 201:
            return cast(dict[str, Any], json.loads(response.content))
        elif response.status_code == 401:
            raise AuthorizationError()
        else:
            raise FreeplayError(f'Unknown response while creating a session. Response code: {response.status_code}')

    def get_prompts(self, project_id: str, tag: str) -> PromptTemplates:
        prompts = api_support.get(
            target_type=PromptTemplates,
            api_key=self.freeplay_api_key,
            url=f'{self.api_base}/projects/{project_id}/templates/all/{tag}'
        )
        return prompts

    def continue_chat(
            self,
            session_id: str,
            flavor: ChatFlavor,
            tag: str,
            target_template: PromptTemplateWithMetadata,
            variables: Variables,
            message_history: list[ChatMessage],
            new_messages: Optional[list[ChatMessage]],
            test_run_id: Optional[str] = None,
            completion_parameters: Optional[LLMParameters] = None) -> ChatCompletionResponse:
        # make call
        start = int(time.time())
        try:
            params = self.client_params.merge_and_override(completion_parameters)
            completion_response = flavor.continue_chat(message_history, new_messages=new_messages,
                                                       llm_parameters=params)
        except Exception as e:
            raise FreeplayError("Error calling service") from e
        end = int(time.time())

        formatted_prompt = json.dumps(completion_response.message_history)
        # record data
        self.__record_call(
            completion_response.content,
            completion_response.is_complete,
            end,
            formatted_prompt,
            session_id,
            start,
            target_template,
            variables,
            flavor.record_format_type,
            tag,
            test_run_id
        )

        return completion_response

    def prepare_and_make_call(
            self,
            session_id: str,
            prompts: PromptTemplates,
            template_name: str,
            variables: dict[str, str],
            flavor: Flavor,
            tag: str,
            test_run_id: Optional[str] = None,
            completion_parameters: Optional[LLMParameters] = None
    ) -> CompletionResponse:
        target_template = self.find_template_by_name(prompts, template_name)
        formatted_prompt = flavor.format(target_template, variables)

        # make call
        start = int(time.time())
        try:
            params = self.client_params.merge_and_override(completion_parameters)
            completion_response = flavor.call_service(
                formatted_prompt=formatted_prompt, llm_parameters=params)
        except Exception as e:
            raise FreeplayError("Error calling service") from e
        end = int(time.time())

        # record data
        self.__record_call(
            completion_response.content,
            completion_response.is_complete,
            end,
            formatted_prompt,
            session_id,
            start,
            target_template,
            variables,
            flavor.record_format_type,
            tag,
            test_run_id
        )

        return completion_response

    def prepare_and_make_call_stream(
            self,
            session_id: str,
            prompts: PromptTemplates,
            template_name: str,
            variables: dict[str, str],
            flavor: Flavor,
            tag: str,
            test_run_id: Optional[str] = None,
            completion_parameters: Optional[LLMParameters] = None
    ) -> Generator[CompletionChunk, None, None]:
        target_template = self.find_template_by_name(prompts, template_name)
        formatted_prompt = flavor.format(target_template, variables)

        # make call
        start = int(time.time())
        params = self.client_params.merge_and_override(completion_parameters)
        completion_response = flavor.call_service_stream(formatted_prompt=formatted_prompt, llm_parameters=params)
        text_chunks = []
        last_is_complete = False
        for chunk in completion_response:
            text_chunks.append(chunk.text)
            last_is_complete = chunk.is_complete
            yield chunk
        end = int(time.time())

        self.__record_call(''.join(text_chunks),
                           last_is_complete,
                           end,
                           formatted_prompt,
                           session_id,
                           start,
                           target_template,
                           variables,
                           flavor.record_format_type,
                           tag,
                           test_run_id)

    def __record_call(
            self,
            completion_content: str,
            completion_is_complete: bool,
            end: int,
            formatted_prompt: str,
            session_id: str,
            start: int,
            target_template: PromptTemplateWithMetadata,
            variables: dict[str, str],
            record_format_type: str,
            tag: str,
            test_run_id: Optional[str]
    ) -> None:

        record_payload = {
            "session_id": session_id,
            "project_version_id": target_template.project_version_id,
            "prompt_template_id": target_template.prompt_template_id,
            "start_time": start,
            "end_time": end,
            "tag": tag,
            "inputs": variables,
            "prompt_content": formatted_prompt,
            "return_content": completion_content,
            "format_type": record_format_type,
            "is_complete": completion_is_complete
        }

        if test_run_id is not None:
            record_payload['test_run_id'] = test_run_id

        try:
            recorded_response = api_support.post_raw(
                api_key=self.freeplay_api_key,
                url=f'{self.api_base}/v1/record',
                payload=record_payload
            )
            recorded_response.raise_for_status()
        except Exception as e:
            status_code = -1
            if hasattr(e, 'response') and hasattr(e.response, 'status_code'):
                status_code = e.response.status_code
            logger.warning(f'There was an error recording to Freeplay. Call will not be logged. '
                           f'Status: {status_code}. {e.__class__}')


class Session:
    def __init__(
            self,
            call_support: CallSupport,
            session_id: str,
            prompts: PromptTemplates,
            flavor: Flavor,
            tag: str = default_tag,
            test_run_id: Optional[str] = None
    ) -> None:
        self.tag = tag
        self.call_support = call_support
        self.session_flavor = flavor
        self.session_id = session_id
        self.prompts = prompts
        self.test_run_id = test_run_id

    def get_completion(
            self,
            template_name: str,
            variables: dict[str, str],
            flavor: Optional[Flavor] = None,
            **kwargs: Any
    ) -> CompletionResponse:
        completion_flavor = flavor or self.session_flavor
        return self.call_support.prepare_and_make_call(self.session_id,
                                                       self.prompts,
                                                       template_name,
                                                       variables,
                                                       completion_flavor,
                                                       self.tag,
                                                       self.test_run_id,
                                                       completion_parameters=LLMParameters(kwargs))

    def get_completion_stream(
            self,
            template_name: str,
            variables: dict[str, str],
            flavor: Optional[Flavor] = None,
            **kwargs: Any
    ) -> Generator[CompletionChunk, None, None]:
        completion_flavor = flavor or self.session_flavor
        return self.call_support.prepare_and_make_call_stream(self.session_id,
                                                              self.prompts,
                                                              template_name,
                                                              variables,
                                                              completion_flavor,
                                                              self.tag,
                                                              self.test_run_id,
                                                              completion_parameters=LLMParameters(kwargs))


class ChatSession(Session):
    def __init__(
            self,
            call_support: CallSupport,
            session_id: str,
            prompts: PromptTemplates,
            flavor: ChatFlavor,
            template_name: str,
            variables: Variables,
            tag: str = default_tag,
            test_run_id: Optional[str] = None
    ) -> None:
        super().__init__(call_support, session_id, prompts, flavor, tag, test_run_id)
        self.flavor = flavor  # Ensures a ChatFlavor is set for this session.
        # A Chat Session tracks the template_name and variables for a set of chat completions.
        # Assumes these will be the same for subsequent chat messages.
        self.template_name = template_name
        self.variables = variables

    def start_chat(self, **kwargs: Any) -> ChatCompletionResponse:
        target_template = self.call_support.find_template_by_name(self.prompts, self.template_name)
        formatted_prompt = self.session_flavor.format(target_template, self.variables)
        messages = json.loads(formatted_prompt)
        return self.continue_chat(message_history=messages, new_messages=None, **kwargs)

    def continue_chat(
            self,
            message_history: list[ChatMessage],
            new_messages: Optional[list[ChatMessage]] = None,
            **kwargs: Any
    ) -> ChatCompletionResponse:
        target_template = self.call_support.find_template_by_name(self.prompts, self.template_name)

        return self.call_support.continue_chat(
            self.session_id,
            flavor=self.flavor,
            tag=self.tag,
            test_run_id=self.test_run_id,
            target_template=target_template,
            variables=self.variables,
            message_history=message_history,
            new_messages=new_messages,
            completion_parameters=LLMParameters(kwargs)
        )


@dataclass()
class FreeplayTestRun:
    def __init__(
            self,
            call_support: CallSupport,
            flavor: Flavor,
            test_run_id: str,
            inputs: list[dict[str, str]]
    ):
        self.call_support = call_support
        self.flavor = flavor
        self.test_run_id = test_run_id
        self.inputs = inputs

    def get_inputs(self) -> list[dict[str, str]]:
        return self.inputs

    def create_session(self, project_id: str, tag: str = default_tag) -> Session:
        project_session = self.call_support.create_session(project_id, tag, self.test_run_id)
        prompts = self.call_support.get_prompts(project_id, tag)
        return Session(self.call_support, project_session['session_id'], prompts, self.flavor, tag, self.test_run_id)


# This SDK prototype does not support full functionality of either OpenAI's API or Freeplay's
# The simplifications are:
#  - Always assumes there is a single choice returned, does not support multiple
#  - Does not support an "escape hatch" to allow use of features we don't explicitly expose
class Freeplay:
    def __init__(
            self,
            flavor: Flavor,
            freeplay_api_key: str,
            api_base: str,
            **kwargs: Any
    ) -> None:
        if not freeplay_api_key or not freeplay_api_key.strip():
            raise APIKeyMissingError("Freeplay API key not set. It must be set to the Freeplay API.")

        self.call_support = CallSupport(freeplay_api_key, api_base, **kwargs)
        self.client_flavor = flavor
        self.freeplay_api_key = freeplay_api_key
        self.api_base = api_base

    def create_session(self, project_id: str, tag: str = default_tag) -> Session:
        project_session = self.call_support.create_session(project_id, tag)
        prompts = self.call_support.get_prompts(project_id, tag)
        return Session(self.call_support, project_session['session_id'], prompts, self.client_flavor, tag)

    def get_completion(
            self,
            project_id: str,
            template_name: str,
            variables: dict[str, str],
            tag: str = default_tag,
            flavor: Optional[Flavor] = None,
            **kwargs: Any
    ) -> CompletionResponse:
        project_session = self.call_support.create_session(project_id, tag)
        prompts = self.call_support.get_prompts(project_id, tag)
        completion_flavor = flavor or self.client_flavor

        return self.call_support.prepare_and_make_call(project_session['session_id'],
                                                       prompts,
                                                       template_name,
                                                       variables,
                                                       completion_flavor,
                                                       tag,
                                                       completion_parameters=LLMParameters(kwargs))

    def get_completion_stream(
            self,
            project_id: str,
            template_name: str,
            variables: dict[str, str],
            tag: str = default_tag,
            flavor: Optional[Flavor] = None,
            **kwargs: Any
    ) -> Generator[CompletionChunk, None, None]:
        project_session = self.call_support.create_session(project_id, tag)
        prompts = self.call_support.get_prompts(project_id, tag)
        completion_flavor = flavor or self.client_flavor

        return self.call_support.prepare_and_make_call_stream(project_session['session_id'],
                                                              prompts,
                                                              template_name,
                                                              variables,
                                                              completion_flavor,
                                                              tag,
                                                              completion_parameters=LLMParameters(kwargs))

    def create_test_run(self, project_id: str, playlist: str) -> FreeplayTestRun:
        response = api_support.post_raw(
            api_key=self.freeplay_api_key,
            url=f'{self.api_base}/projects/{project_id}/test-runs',
            payload={'playlist_name': playlist},
        )

        json_dom = response.json()

        return FreeplayTestRun(self.call_support, self.client_flavor, json_dom['test_run_id'], json_dom['inputs'])

    def start_chat(
            self,
            project_id: str,
            template_name: str,
            variables: Variables,
            tag: str = default_tag,
            **kwargs: Any
    ) -> Tuple[ChatSession, ChatCompletionResponse]:
        if not isinstance(self.client_flavor, ChatFlavor):
            raise FreeplayError('A Chat flavor is required to start a chat session.')

        project_session = self.call_support.create_session(project_id, tag)
        prompts = self.call_support.get_prompts(project_id, tag)
        session = ChatSession(
            self.call_support,
            project_session['session_id'],
            prompts,
            self.client_flavor,
            template_name,
            variables,
            tag)
        completion_response = session.start_chat(**kwargs)
        return session, completion_response
