from __future__ import annotations

import json
import logging
import os
from abc import abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Generic, TypeVar

from typing_extensions import get_args, get_origin, get_original_bases

from aibs_informatics_core.collections import PostInitMixin
from aibs_informatics_core.env import EnvBaseMixins
from aibs_informatics_core.models.aws.s3 import S3Path
from aibs_informatics_core.models.base import ModelProtocol
from aibs_informatics_core.utils.json import JSON, JSONObject, is_json_str, load_json_object

logger = logging.getLogger(__name__)


REQUEST = TypeVar("REQUEST", bound=ModelProtocol)
RESPONSE = TypeVar("RESPONSE", bound=ModelProtocol)

E = TypeVar("E", bound="BaseExecutor")


@dataclass  # type: ignore[misc] # mypy #5374
class BaseExecutor(EnvBaseMixins, PostInitMixin, Generic[REQUEST, RESPONSE]):
    """Abstract base executor for handling typed request/response workflows.

    Subclasses must parameterize with request and response model types and
    implement the ``handle`` method. Provides serialization, deserialization,
    and CLI/remote I/O utilities.
    """

    @abstractmethod
    def handle(self, request: REQUEST) -> RESPONSE | None:  # pragma: no cover
        """Core logic for handling request

        NOT IMPLEMENTED

        Args:
            request (API_REQUEST): Request object expected

        Returns:
            response object returned, Optional
        """
        raise NotImplementedError("Must implement handler logic here")

    # --------------------------------------------------------------------
    # Request & Response De-/Serialization
    # --------------------------------------------------------------------

    @classmethod
    def _get_generic_args(cls) -> tuple[type, ...]:
        """Walk the MRO looking for the parameterised ``BaseExecutor[REQ, RESP]`` base."""
        for klass in cls.__mro__:
            for base in get_original_bases(klass):
                origin = get_origin(base)
                if isinstance(origin, type) and issubclass(origin, BaseExecutor):
                    args = get_args(base)
                    if args and not any(isinstance(a, TypeVar) for a in args):
                        return args
        raise TypeError(
            f"Could not determine generic type arguments for {cls}. "
            "Ensure the class directly parameterises BaseExecutor[REQUEST, RESPONSE]."
        )

    @classmethod
    def get_request_cls(cls) -> type[REQUEST]:
        """Return the request model class from the generic type arguments.

        Returns:
            The request model class.
        """
        return cls._get_generic_args()[0]  # type: ignore

    @classmethod
    def get_response_cls(cls) -> type[RESPONSE]:
        """Return the response model class from the generic type arguments.

        Returns:
            The response model class.
        """
        return cls._get_generic_args()[1]  # type: ignore

    @classmethod
    def deserialize_request(cls, request: JSON) -> REQUEST:
        """Deserialize a raw request

        Supported request types:
        Union[JSONObject, S3Path, str, Path, REQUEST]
            1. dict object
            2. S3 Path object (must implement s3 deserialization method)
            3. stringified object
            4. path to file of json object

        Args:
            request (Union[JSONObject, S3Path, str, Path]): raw request

        Returns:
            Deserialized request
        """
        if isinstance(request, cls.get_request_cls()):
            return request

        request = cls.load_input(request)
        if isinstance(request, dict):
            try:
                return cls.get_request_cls().from_dict(request)
            except Exception as e:
                logger.error(f"Could not deserialize object. error: {e}")
                raise e
        else:
            raise ValueError(f"Cannot deserialize request: {request}, type: {type(request)}")

    @classmethod
    def serialize_response(cls, response: RESPONSE) -> JSONObject:
        """Serializes response object into JSON

        Args:
            response (API_RESPONSE): The returned response object

        Returns:
            serialized form of response object
        """
        return response.to_dict()

    @classmethod
    def load_input(cls, input: Any) -> JSON:
        """Loads input from various sources

        Supported input types:
            1. dict object
            2. S3 Path object (must implement s3 deserialization method)
            3. stringified object
            4. path to file of json object

        Args:
            input (Any): raw input. Can be any of the supported types.

        Returns:
            content loaded from input
        """
        if isinstance(input, dict):
            return input
        elif isinstance(input, str) and is_json_str(input):
            return load_json_object(input)
        elif isinstance(input, S3Path) or (isinstance(input, str) and S3Path.is_valid(input)):
            return cls.load_input__remote(S3Path(input))
        elif isinstance(input, (str, Path)):
            return cls.load_input__file(Path(input))
        else:
            raise ValueError(f"Cannot read input: {input}, type: {type(input)}")

    @classmethod
    def load_input__remote(cls, remote_path: S3Path) -> JSON:
        """Reads input from S3 location

        Args:
            s3_path (S3Path): S3 location to read input from

        Returns:
            JSON data read from S3
        """
        raise NotImplementedError(
            f"{cls.__name__} has not provided an implementation for this form of input"
        )

    @classmethod
    def load_input__file(cls, local_path: Path) -> JSON:
        """Reads input from file

        Args:
            local_path (Path): Local path to read input from

        Returns:
            JSON data read from file
        """
        if not local_path.exists():
            raise ValueError(f"local path specified {local_path} does not exist")
        if local_path.is_dir():
            raise ValueError(f"local path specified {local_path} cannot be read. Must be a file.")
        return load_json_object(local_path)

    @classmethod
    def write_output(cls, output: JSON, path: str | Path) -> None:
        """Writes output to location

        Args:
            output (JSON): serialized output object
            path (Path | str): location to write output to. Can be S3 or local file
        """

        if isinstance(path, S3Path) or (isinstance(path, str) and S3Path.is_valid(path)):
            return cls.write_output__remote(output, S3Path(path))
        else:
            # Here we assume that it could be string or path to file
            return cls.write_output__file(output, Path(path))

    @classmethod
    def write_output__remote(cls, output: JSON, remote_path: S3Path) -> None:
        """Writes output to Remote location

        Args:
            output (RESPONSE): output object
            remote_path (S3Path): S3 location to write response to
        """
        raise NotImplementedError(
            f"{cls.__name__} has not provided an implementation for this form of serialization"
        )

    @classmethod
    def write_output__file(cls, output: JSON, local_path: Path) -> None:
        """Writes output to file

        Args:
            output (JSON): serialized output json
            local_path (S3Path): Local path to write response to
        """

        if local_path.is_dir() or (
            local_path.parent.exists() and not os.access(local_path.parent, os.W_OK)
        ):
            raise ValueError(
                f"local path specified {local_path} cannot be written to. Must be a file "
            )
        local_path.parent.mkdir(parents=True, exist_ok=True)
        local_path.write_text(json.dumps(output, sort_keys=True))

    @classmethod
    def get_executor_name(cls) -> str:
        """Returns a distinguishable name of the executor class

        By default, it returns class name.

        """
        return cls.__name__

    @classmethod
    def build_from_env(cls: type[E], **kwargs) -> E:
        """Creates an executor from environment

        Must be able to create an executor from environment variables

        Returns:
            Executor object
        """
        return cls(**kwargs)

    @classmethod
    def run_executor(
        cls, input: JSON, output_location: str | Path | None = None, **kwargs
    ) -> JSON | None:
        """Runs an executor

        Args:
            input (JSON): input to executor
            output_location (Optional[str], optional): Optional output location to write
                response to. Defaults to None.
        """
        executor = cls.build_from_env(**kwargs)

        request = cls.deserialize_request(input)
        response = executor.handle(request)
        output = executor.serialize_response(response) if response else None

        if output_location:
            if output is None:
                logger.warning(
                    "Response is None but output location is specified. Writing empty dict"
                )
            executor.write_output(output or {}, output_location)

        return output
