from typing import Dict, Any, List, Optional, Union
from enum import Enum
from pydantic import BaseModel
from environs import Env

# from llmops.business.singleton import get_knowledge_assistant
from llmops.utils.json_utils import format_json_string_by_ai

env = Env()
env.read_env()


# knowledge_assistant = get_knowledge_assistant()


# OperationID = Enum("OperationID", {op: op for op in knowledge_assistant.operation_ids})


class AbstractOperation(BaseModel):
    # operation_id: Enum  # OperationID
    # attributes: str  # json string
    # original_description: str

    # def model_dump(self, **kwargs):
    #     return {
    #         "operation_id": self.operation_id.value,
    #         "attributes": self.attributes,
    #         "original_description": self.original_description,
    #     }
    original_description: str


class Judgement4GPT(BaseModel):
    reason: str
    judgement: bool


class StepType(str, Enum):
    GENERATED = "generated"
    DEFINED = "defined"


class RequestType(str, Enum):
    HTTP = "http"
    GRPC = "grpc"
    AWS_CLI = "aws_cli"
    AZURE_CLI = "azure_cli"
    GCP_CLI = "gcp_cli"
    MYSQLSH = "mysqlsh"


class HttpMethod(str, Enum):
    """not the universal set of HTTP methods"""

    GET = "GET"
    POST = "POST"
    PUT = "PUT"
    PATCH = "PATCH"
    DELETE = "DELETE"


class Auth(BaseModel):
    auth_type: str
    content: Optional[str]

    def model_dump(self, **kwargs):
        data = super().model_dump(**kwargs)
        data["content"] = (
            format_json_string_by_ai(self.content) if self.content else "{}"
        )
        return data


class HttpRequest(BaseModel):
    auth: Auth
    method: HttpMethod
    path: str
    path_params: Optional[str]
    query_params: Optional[str]
    headers: Optional[str]
    body: Optional[str]

    def model_dump(self, **kwargs):
        data = super().model_dump(**kwargs)
        data["method"] = self.method.value
        data["path_params"] = (
            format_json_string_by_ai(self.path_params) if self.path_params else "{}"
        )
        data["query_params"] = (
            format_json_string_by_ai(self.query_params) if self.query_params else "{}"
        )
        data["headers"] = (
            format_json_string_by_ai(self.headers) if self.headers else "{}"
        )
        data["body"] = format_json_string_by_ai(self.body) if self.body else "{}"
        return data


class CliRequest(BaseModel):
    service: str
    command: str
    args: List[str]
    args_params: Optional[str]


class CliRequest4GPT(BaseModel):
    command: str
    args: List[str]


class Step(BaseModel):
    operation_id: str
    step_type: StepType
    request_type: RequestType
    request: Union[HttpRequest, CliRequest]  # HttpRequest
    description: str
    max_retries: int
    delay_after: int
    expect: Optional[str]

    def model_dump(self, **kwargs):
        data = super().model_dump(**kwargs)
        # Convert the placeholder to a string
        data["request_type"] = self.request_type.value
        data["step_type"] = self.step_type.value
        data["request"] = self.request.model_dump()
        return data


class HttpRequest4GPT(BaseModel):
    method: HttpMethod
    path: str
    path_params: Optional[str]
    query_params: Optional[str]
    headers: Optional[str]
    body: Optional[str]
    role: Optional[str]


class Expect(BaseModel):
    status_code: int
    cel: str


class Step4GPT(BaseModel):
    analysis: str
    request: Union[HttpRequest4GPT, CliRequest4GPT]  # HttpRequest4GPT
    description: str
    expect: Optional[Expect]
    request_type: RequestType


class ExecutionAnalysisType(str, Enum):
    OPERATION_ERROR = "operation_error"  # select wrong operation
    REQUEST_ERROR = "request_error"  # generate wrong request
    SUCCESS = "success"  # execute successfully


class ExecutionJudgement(BaseModel):
    """Critic 的评估结果"""

    reasoning: str
    result_type: ExecutionAnalysisType
    confidence: float


class ExecutionResult(BaseModel):
    status_code: int
    json: Dict[str, Any] | List[Dict[str, Any]]
    headers: Dict[str, str]
    error: Optional[str] = None


class Resource(BaseModel):
    id: str
    resource_type: str
    attributes: str
    state: str
    parent_id: str
    reference_ids: List[str]


class ResourceList4GPT(BaseModel):
    resources: List[Resource]