import json
import os
from functools import partial, wraps
from importlib import import_module
from typing import Coroutine, TypeVar, Callable, Any, ParamSpec

from .sdk.client import AuthenticatedClient
import julep.sdk.api.default as ops  # This is needed because of the way relative imports work
from .sdk import errors


### GOAL: Convert all autogenerated code into a slightly better interface
# - The generated code is in sdk/api/default and sdk/models
# - each file in sdk/api/default is an operation that contains
#   + a function called `sync_detailed`
#   + a function called `asyncio_detailed`

# - for example, `agents_route_create.py` contains:
#   + def sync_detailed(self, *, client: AuthenticatedClient, body: CreateAgentRequest) -> Response[Agent]:
#   + def asyncio_detailed(self, *, client: AuthenticatedClient, body: CreateAgentRequest) -> Response[Agent]:

# - each file in sdk/models is an `attrs` class that is a model of the API response

# What I want is:
# - a single client that can be used to call all the methods
# - a neat namespaced interface to call the methods

# so the above methods would be called like:
# - julep.agents.create(body=CreateAgentRequest(...))


# Need to get all files from ops module and import them

# Get all the files from the ops module
ops_files = [
    file
    for file in os.listdir(ops.__file__.rstrip("__init__.py"))
    if file.endswith(".py") and file != "__init__.py"
]

# Import all the files
for file in ops_files:
    file_name = file[:-3]
    try:
        import_module(f"julep.sdk.api.default.{file_name}")
    except ImportError:
        breakpoint()
        raise

    setattr(ops, file_name, import_module(f"julep.sdk.api.default.{file_name}"))

# Now we have `ops.agents_route_create` etc.

# These are used to forward the types from the generated code to the decorated functions
T = TypeVar("T", bound=Callable[..., Any])

P = ParamSpec("P")
R = TypeVar("R")


# Decorator to parse the response from the generated code
def parse_response(fn: Callable[P, R]) -> Callable[P, R | dict[str, Any]]:
    @wraps(fn)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | dict[str, Any]:
        #
        # Call the original function
        response = fn(*args, **kwargs)

        # If the response is an error, raise an exception
        if response.status_code.is_client_error or response.status_code.is_server_error:
            raise errors.UnexpectedStatus(response.status_code, response.content)

        else:
            parsed = response.parsed

            # If the response is a model, return it
            if parsed:
                return parsed

            else:
                # Unfortuantely, sometimes the generated code doesn't return a model,
                #   so we need to parse the JSON manually
                return json.loads(response.content)

    return wrapper


# Same decorator but for async functions
def parse_response_async(
    fn: Callable[P, R],
) -> Callable[P, Coroutine[Any, Any, R | dict[str, Any]]]:
    @wraps(fn)
    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | dict[str, Any]:
        response = await fn(*args, **kwargs)

        if response.status_code.is_client_error or response.status_code.is_server_error:
            raise errors.UnexpectedStatus(response.status_code, response.content)
        else:
            parsed = response.parsed

            if parsed:
                return parsed
            else:
                return json.loads(response.content)

    return wrapper


# This is the namespace class that will be used to attach the methods to
# so `julep.agents = JulepNamespace("agents", client)`
class JulepNamespace:
    def __init__(self, name: str, client: AuthenticatedClient):
        self.name = name
        self.client = client


class Julep:
    def __init__(
        self,
        *,
        api_key: str,
        base_url: str = "https://api.julep.ai/api",
        **client_kwargs,
    ):
        self.client = AuthenticatedClient(
            token=api_key, base_url=base_url, **client_kwargs
        )

        # Get a list of all the available operations (attributes of the ops object)
        op_names: list[str] = [
            attr for attr in dir(ops)
            # We only want the operations, not the helpers etc.
            if not attr.startswith("_") and "route" in attr
        ]

        # These look like: agents_route_create / agents_route_list etc.
        # The conventions are:
        # - The first part of the name is the resource name (agents, chat, etc.)
        # - The second part of the name is the method (route_create, route_list, etc.)
        # - The `_route_` prefix can be omitted.

        # We want to create a method on the Julep class for each of these that proxies to the ops object
        # But also ensures that the first argument (self.client) is passed through

        namespaces, operations = list(
            zip(*[name.split("_route_") for name in op_names])
        )

        # So now we have: [("agents", "create"), ("agents", "list"), ...]

        # Some namespaces have aliases
        # This is because the API is organized by resource, but the SDK is organized by namespace
        # And we want to allow the user to use the resource name as the namespace name
        namespace_aliases = {
            "agents_docs_search": "agent_docs",
            "user_docs_search": "user_docs",
            "execution_transitions": "transitions",
            "execution_transitions_stream": "transitions",
            "individual_docs": "docs",
            "job": "jobs",
            "task_executions": "executions",
            "tasks_create_or_update": "tasks",
        }

        # Now let's add the namespaces to the Julep class as attributes
        for namespace in namespaces:
            namespace = namespace_aliases.get(namespace, namespace)
            if not hasattr(self, namespace):
                setattr(self, namespace, JulepNamespace(namespace, self.client))

        # so now we have:
        # - `julep.agents`
        # - `julep.chat`
        # - `julep.agent_docs`
        # - `julep.user_docs`
        # - `julep.transitions`
        # - `julep.jobs`
        # - `julep.executions`
        # - `julep.tasks`

        # Next let's add the operations to the Julep class as attributes
        # We need to add both sync and async versions
        # The sync version is called `sync_detailed` and the async version is called `asyncio_detailed`
        # The `_route_` prefix is omitted
        async_prefix = {"a": "asyncio_detailed", "": "sync_detailed"}

        # However in the final client namespace we will add `a..` as prefix to denote the async version
        # for example: `julep.agents.create` and `julep.agents.acreate`
        for prefix, op_name in async_prefix.items():
            for namespace, operation in zip(namespaces, operations):
                # Get ops.agents_route_create
                op = getattr(ops, f"{namespace}_route_{operation}")
                # Get ops.agents_route_create.sync_detailed
                op = getattr(op, op_name)
                # Add the client argument as a partial
                op = partial(op, client=self.client)
                # Decorate with the parse_response decorator
                op = (parse_response_async if (prefix == "a") else parse_response)(op)
                # Prepare the name for the final client namespace
                final_op_name = prefix + operation

                # Get the namespace object
                namespace = namespace_aliases.get(namespace, namespace)
                ns = getattr(self, namespace)

                # Add the operation to the namespace
                if not hasattr(ns, final_op_name):
                    setattr(ns, final_op_name, op)
