from typing import Any, Type, Optional

import urllib3
from pydantic import BaseModel, Field

from codemie_tools.base.codemie_tool import CodeMieTool
from codemie_tools.open_api.models import OpenApiConfig
from codemie_tools.open_api.tools_vars import OPEN_API_TOOL, OPEN_API_SPEC_TOOL
from codemie_tools.base.utils import parse_to_dict


class OpenApiInput(BaseModel):
    method: str = Field(
        description="""
            String text. It's the verb of a HTTP request generated by LLM which 
            will be used to invoke appropriate REST API to fulfill user request. 
            Important: string formatting and escaping SHOULD NOT be used when passing query to the tool.
            """.strip(),
    )
    url: str = Field(
        description="""
            String text. It's the URL of a HTTP request generated by LLM which 
            will be used to invoke appropriate REST API to fulfill user request. 
            Important: string formatting and escaping SHOULD NOT be used when passing query to the tool.
            """.strip(),
    )
    headers: Optional[str] = Field(
        default="",
        description="""
            MUST be String text. It's the headers of a HTTP request generated by LLM which
            will be used to invoke appropriate REST API to fulfill user request.
            Important: string formatting and escaping SHOULD NOT be used when passing query to the tool.
            """.strip(),
    )
    fields: Optional[str] = Field(
        default="",
        description="""
            MUST be String text. It's the query parameters of a HTTP request generated by LLM which 
            will be used to invoke appropriate REST API to fulfill user request.
            Important: string formatting and escaping SHOULD NOT be used when passing query to the tool.
             """.strip(),
    )
    body: Optional[str] = Field(
        default="",
        description="""
            JSON that should be passed as a string. It's the body of a HTTP request generated by LLM which 
            will be used to invoke appropriate REST API to fulfill user request. 
            Should be passes as an empty string if there is no body. 
            Important: string formatting and escaping SHOULD NOT be used when passing query to the tool.
            """.strip(),
    )
    
class GetOpenApiSpecToolInput(BaseModel):
    query: Optional[str] = Field(
        default="",
        description="User initial request should be passed as a string.",
    )


class InvokeRestApiBySpec(CodeMieTool):
    openapi_config: OpenApiConfig = Field(exclude=True, default=None)
    name: str = OPEN_API_TOOL.name
    description: str = OPEN_API_TOOL.description
    args_schema: Type[BaseModel] = OpenApiInput

    def execute(self, method: str, url: str, headers: Any = "", fields: Any = "", body: str = "", **kwargs: Any) -> Any:
        manager = urllib3.PoolManager()
        encoded_data = body.encode('utf-8') if body else None
        headers_param = parse_to_dict(headers) if headers else {}
        fields_param = parse_to_dict(fields) if fields else None

        if self.openapi_config.api_key:
            headers_param['Authorization'] = "Bearer " + self.openapi_config.api_key

        response = manager.request(method=method, url=url, fields=fields_param, headers=headers_param,
                                   body=encoded_data)
        return response.data


class GetOpenApiSpec(CodeMieTool):
    openapi_spec: str = Field(exclude=True, default=None)
    name: str = OPEN_API_SPEC_TOOL.name
    description: str = OPEN_API_SPEC_TOOL.description
    args_schema: Type[BaseModel] = GetOpenApiSpecToolInput

    def __init__(self, openapi_config: OpenApiConfig):
        super().__init__()
        self.openapi_spec = openapi_config.spec

    def execute(self, *args, **kwargs: Any) -> Any:
        return self.openapi_spec
