from typing import Dict, Any, List
import json

from llmops.base import (
    Step,
    Step4GPT,
    HttpRequest,
    CliRequest,
    StepType,
    RequestType,
    Auth,
)
from llmops.utils.llm_client import LlmClient, LlmConfig
from llmops.config import prompt_path
from llmops.business.business_expert import BusinessExpert


class RequestGenerator:
    def __init__(self, business_expert: BusinessExpert, model: str):
        self.system_prompt_template = prompt_path / "request_generator.md"
        self.client = LlmClient(self.system_prompt_template, llm_config=LlmConfig(model=model))
        self.business_expert = business_expert
        self.trajectory: List[Dict[str, Any]] = []
        if business_expert.biz_system_prompt_request_generator != "":
            self.client.update_system_prompt(
                self.client.system_prompt
                + "\nBusiness Logic:\n"
                + business_expert.biz_system_prompt_request_generator
            )

    def generate_step(
        self,
        operation_id: str,
        current_state: Dict[str, Any],
        target_operation: Dict[str, Any],
        feedback: dict,
        reason: str,
    ) -> Step:
        operation = self.business_expert.extract_operation(operation_id)
        if operation.get("paths"):
            request_type = RequestType.HTTP
        else:
            request_type = RequestType(next(iter(operation.values()))["cli_type"])

        # Prepare prompt
        user_prompt = {
            "operation": operation,
            "selected_operation_reason": reason,
            "request_type": request_type,
            "current_state": current_state,
            "target_operation": target_operation,
            "feedback": json.dumps(feedback) if feedback else "None",
            "history": [
                {"output": attempt["output"], "feedback": attempt["input"]["feedback"]}
                for attempt in self.trajectory[-3:]
            ],
        }

        try:
            llm_response = self.client.generate(user_prompt, Step4GPT)
            expect = json.dumps(llm_response.expect.model_dump()) if llm_response.expect is not None else None
            if llm_response.request_type == RequestType.HTTP:
                request = HttpRequest(
                    auth=Auth(auth_type=self.business_expert.executor.auth_type, content=llm_response.request.role),
                    method=llm_response.request.method,
                    path=llm_response.request.path,
                    path_params=llm_response.request.path_params,
                    query_params=llm_response.request.query_params,
                    headers=llm_response.request.headers,
                    body=llm_response.request.body,
                )
                result = Step(
                    operation_id=operation_id,
                    step_type=StepType.GENERATED,
                    request_type=llm_response.request_type,
                    request=request,
                    description=llm_response.description,
                    max_retries=0,
                    delay_after=0,
                    expect=expect,
                )
                result.step_type = StepType.GENERATED
                if result.request.method.value == "GET":
                    result.max_retries = 120
                    result.delay_after = 30
                else:
                    result.max_retries = 0
                    result.delay_after = 0
            elif llm_response.request_type in [RequestType.AWS_CLI, RequestType.AZURE_CLI, RequestType.GCP_CLI, RequestType.MYSQLSH]:
                request = CliRequest(
                    service=llm_response.request.args[0],
                    command=llm_response.request.command,
                    args=llm_response.request.args,
                    args_params="{}",
                )
                result = Step(
                    operation_id=operation_id,
                    step_type=StepType.GENERATED,
                    request_type=llm_response.request_type,
                    request=request,
                    description=llm_response.description,
                    max_retries=0,
                    delay_after=0,
                    expect=expect,
                )
            self.trajectory.append(
                {
                    "input": {
                        "operation_id": operation_id,
                        "current_state": current_state,
                        "target_operation": target_operation,
                        "feedback": json.dumps(feedback) if feedback else "None",
                    },
                    "output": llm_response.model_dump(),
                }
            )
            return result
        except Exception:
            raise

    def get_trajectory(self) -> Dict[str, Any]:
        """Get the trajectory of request generation process"""
        return self.trajectory