from typing import List, Optional, Type, Union, cast

from arthur_common.models.llm_model_providers import (
    LLMConfigSettings,
    LLMResponseFormat,
    MessageRole,
    ModelProvider,
    OpenAIMessage,
)
from arthur_common.models.task_eval_schemas import LLMEval
from litellm import supports_response_schema
from pydantic import BaseModel
from sqlalchemy.orm import Session

from db_models.llm_eval_models import DatabaseLLMEval, DatabaseLLMEvalVersionTag
from repositories.base_llm_repository import BaseLLMRepository
from repositories.model_provider_repository import ModelProviderRepository
from schemas.agentic_prompt_schemas import AgenticPrompt
from schemas.llm_eval_schemas import ReasonedScore
from schemas.request_schemas import (
    BaseCompletionRequest,
    CreateEvalRequest,
    PromptCompletionRequest,
)
from schemas.response_schemas import (
    LLMEvalRunResponse,
    LLMEvalsVersionListResponse,
    LLMVersionResponse,
)
from services.prompt.chat_completion_service import ChatCompletionService

llm_eval_system_instructions = """
You are a helpful assistant that adheres exactly to the following instructions.

You will be given a set of instructions and respond with a boolean score and a reason for your score.
The detail for how and what you should score will be provided in the instructions.
"""


class LLMEvalsRepository(
    BaseLLMRepository[DatabaseLLMEval, DatabaseLLMEvalVersionTag, CreateEvalRequest],
):
    db_model: Type[DatabaseLLMEval] = DatabaseLLMEval
    tag_db_model: Type[DatabaseLLMEvalVersionTag] = DatabaseLLMEvalVersionTag
    version_list_response_model: Type[BaseModel] = LLMEvalsVersionListResponse

    def __init__(self, db_session: Session):
        super().__init__(db_session)
        self.model_provider_repo = ModelProviderRepository(db_session)
        self.chat_completion_service = ChatCompletionService()

    def from_db_model(self, db_eval: DatabaseLLMEval) -> LLMEval:
        tags = self._get_all_tags_for_item_version(db_eval)

        return LLMEval(
            name=db_eval.name,
            model_name=db_eval.model_name,
            model_provider=ModelProvider(db_eval.model_provider),
            instructions=db_eval.instructions,
            variables=db_eval.variables,
            tags=tags,
            config=db_eval.config if db_eval.config else None,
            created_at=db_eval.created_at,
            deleted_at=db_eval.deleted_at,
            version=db_eval.version,
        )

    def _to_versions_reponse_item(
        self,
        db_item: DatabaseLLMEval,
        tags: Optional[List[str]] = None,
    ) -> LLMVersionResponse:
        tags = self._get_all_tags_for_item_version(db_item)

        return LLMVersionResponse(
            version=db_item.version,
            created_at=db_item.created_at,
            deleted_at=db_item.deleted_at,
            model_provider=ModelProvider(db_item.model_provider),
            model_name=db_item.model_name,
            tags=tags,
        )

    def _clear_db_item_data(self, db_item: DatabaseLLMEval) -> None:
        db_item.model_name = ""
        db_item.instructions = ""
        db_item.config = None

    def _extract_variables_from_item(self, item: CreateEvalRequest) -> List[str]:
        return list(
            self.chat_completion_service.find_undeclared_variables_in_text(
                item.instructions,
            ),
        )

    def from_llm_eval_to_agentic_prompt(
        self,
        llm_eval: LLMEval,
        response_format: Optional[Union[LLMResponseFormat, Type[BaseModel]]] = None,
    ) -> AgenticPrompt:
        messages = [
            OpenAIMessage(
                role=MessageRole.SYSTEM,
                content=llm_eval_system_instructions,
            ),
            OpenAIMessage(role=MessageRole.USER, content=llm_eval.instructions),
        ]

        config_dict = {}
        if llm_eval.config:
            config_dict = llm_eval.config.model_dump(exclude_none=True)

        if response_format is not None:
            config_dict["response_format"] = response_format

        return AgenticPrompt(
            name=llm_eval.name,
            model_name=llm_eval.model_name,
            model_provider=llm_eval.model_provider,
            messages=messages,
            version=llm_eval.version,
            created_at=llm_eval.created_at,
            deleted_at=llm_eval.deleted_at,
            config=LLMConfigSettings(**config_dict),
            tools=None,
        )

    def save_llm_item(
        self,
        task_id: str,
        item_name: str,
        item: CreateEvalRequest,
        commit: bool = True,
    ) -> LLMEval:
        return cast(
            LLMEval,
            super().save_llm_item(task_id, item_name, item, commit=commit),
        )

    def run_llm_eval(
        self,
        task_id: str,
        eval_name: str,
        version: str = "latest",
        completion_request: Optional[BaseCompletionRequest] = None,
    ) -> LLMEvalRunResponse:
        if not self.model_provider_repo:
            raise ValueError("Model provider repository not initialized")

        # get the llm eval
        llm_eval = self.get_llm_item(task_id, eval_name, version)

        if llm_eval.deleted_at is not None:
            raise ValueError(
                f"Cannot run this llm eval because it was deleted on: {llm_eval.deleted_at}",
            )

        # get the llm client
        llm_client = self.model_provider_repo.get_model_provider_client(
            provider=llm_eval.model_provider,
        )

        # NOTE: We currently don't set litellm.enable_json_schema_validation=True, which has litellm validate schemas for models that support structured outputs
        # Some vertex ai models return true for the function below, but don't do any schema validations: https://docs.litellm.ai/docs/completion/json_mode?#validate-json-schema
        # If we choose to support vertex ai in the future, we should set the above flag to True to have litellm validate the schema
        if not supports_response_schema(
            model=llm_eval.model_name,
            custom_llm_provider=llm_eval.model_provider,
        ):
            raise ValueError(
                f"Model {llm_eval.model_name} with provider {llm_eval.model_provider} does not support structured outputs",
            )

        variables = []
        if completion_request and completion_request.variables:
            variables = completion_request.variables

        # create the full completion request
        prompt_completion_request = PromptCompletionRequest(
            variables=variables,
            stream=False,
            strict=True,
        )

        # run the chat completion
        agentic_prompt = self.from_llm_eval_to_agentic_prompt(
            llm_eval=llm_eval,
            response_format=ReasonedScore,
        )

        llm_model_response = (
            self.chat_completion_service.run_chat_completion_raw_response(
                agentic_prompt,
                llm_client,
                prompt_completion_request,
            )
        )

        if llm_model_response.structured_output_response is None:
            raise ValueError(
                f"No structured output response from model {llm_eval.model_name} with provider {llm_eval.model_provider}",
            )

        if not isinstance(
            llm_model_response.structured_output_response,
            ReasonedScore,
        ):
            raise TypeError("Structured output is not a ReasonedScore instance")

        return LLMEvalRunResponse(
            reason=llm_model_response.structured_output_response.reason,
            score=llm_model_response.structured_output_response.score,
            cost=llm_model_response.cost or "",
        )

    def get_llm_item(
        self,
        task_id: str,
        item_name: str,
        item_version: str = "latest",
    ) -> LLMEval:
        return cast(
            LLMEval,
            super().get_llm_item(task_id, item_name, item_version),
        )
