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


env = Env()
env.read_env()


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


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: 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: dict[str, str]

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


class ReusableStep4GPT(BaseModel):
    operation_id: str
    request_type: str
    request: str
    expect: str | None
    description: str
    max_retries: int
    delay_after: int
    name: str
    save: str | None
    continue_on_unexpected: bool
    stage: str
    meta: str


class ReusableStep(BaseModel):
    # case_id: str
    operation_id: str
    # case_type: str
    request_type: str
    request: str
    expect: str | None
    save: str | None
    description: str
    max_retries: int
    delay_after: int
    name: str
    continue_on_unexpected: bool
    stage: str
    meta: str
    generation_id: str
    order: int


class LogMeta(BaseModel):
    task_id: str
    success: bool
    error: str | None
    start_time: str
    end_time: str
    steps_taken: int
    episodes: int
    final_state: dict


class TransitionInput(BaseModel):
    """
    The structure of a single log.json transition, for consumption by various agents.
    The fields are consistent with the log2yaml_design.md document.
    """
    operation_id: str               # Operation ID, original field from log.json
    state_before: dict              # State before transition
    state_after: dict               # State after transition
    request_type: str               # "http" or "cli"
    request: dict | str             # Original request content (dict or str)
    response: Optional[dict | str]  # Original response content (dict or str)
    error: str | None               # Error message
    expect: str | None              # Expected result
    max_retries: int                # Maximum retries
    delay_after: int                # Delay after
    status_code: int                # Response HTTP status code
    operation_selection: dict       # Operation selection
    request_generation: dict        # Request generation
    execution_analysis: dict        # Execution analysis
    target_matching: dict           # Target matching
    resource_manager: dict          # Resource manager

    model_config = ConfigDict(extra='ignore')



class StepJudgeResult(BaseModel):
    """
    The output structure of agent_should_keep.
    keep: Whether to keep this step
    reason: Reason for the decision (for debugging and subsequent agent collaboration)
    """
    reason: str
    keep: bool
    


class StepSaveResult(BaseModel):
    key: str | None
    eval: str | None
    value: str | None


class SaveList(BaseModel):
    placeholder: list[StepSaveResult]