import toml
import yaml
import asyncio
import concurrent.futures
import json
import os
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Any, Dict, List, Optional, Callable

from swarms.utils.workspace_utils import get_workspace_dir


import psutil

try:
    import gzip
except ImportError as error:
    print(f"Error importing gzip: {error}")
# from pydantic import BaseModel


class BaseStructure:
    """Base structure.


    Attributes:
        name (Optional[str]): _description_
        description (Optional[str]): _description_
        save_metadata (bool): _description_
        save_artifact_path (Optional[str]): _description_
        save_metadata_path (Optional[str]): _description_
        save_error_path (Optional[str]): _description_

    Methods:
        run: _description_
        save_to_file: _description_
        load_from_file: _description_
        save_metadata: _description_
        load_metadata: _description_
        log_error: _description_
        save_artifact: _description_
        load_artifact: _description_
        log_event: _description_
        run_async: _description_
        save_metadata_async: _description_
        load_metadata_async: _description_
        log_error_async: _description_
        save_artifact_async: _description_
        load_artifact_async: _description_
        log_event_async: _description_
        asave_to_file: _description_
        aload_from_file: _description_
        run_in_thread: _description_
        save_metadata_in_thread: _description_
        run_concurrent: _description_
        compress_data: _description_
        decompres_data: _description_
        run_batched: _description_
        load_config: _description_
        backup_data: _description_
        monitor_resources: _description_
        run_with_resources: _description_
        run_with_resources_batched: _description_

    Examples:
    >>> base_structure = BaseStructure()
    >>> base_structure
    BaseStructure(name=None, description=None, save_metadata=True, save_artifact_path='./artifacts', save_metadata_path='./metadata', save_error_path='./errors')
    """

    def __init__(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
        save_metadata_on: bool = True,
        save_artifact_path: Optional[str] = "./artifacts",
        save_metadata_path: Optional[str] = "./metadata",
        save_error_path: Optional[str] = "./errors",
    ):
        super().__init__()
        self.name = name
        self.description = description
        self.save_metadata_on = save_metadata_on
        self.save_artifact_path = save_artifact_path
        self.save_metadata_path = save_metadata_path
        self.save_error_path = save_error_path
        self.workspace_dir = get_workspace_dir()

    def run(self, *args, **kwargs):
        """Run the structure."""

    def save_to_file(self, data: Any, file_path: str):
        """Save data to file.

        Args:
            data (Any): _description_
            file_path (str): _description_
        """
        with open(file_path, "w") as file:
            json.dump(data, file)

    def load_from_file(self, file_path: str) -> Any:
        """Load data from file.

        Args:
            file_path (str): _description_

        Returns:
            Any: _description_
        """
        with open(file_path) as file:
            return json.load(file)

    def save_metadata(self, metadata: Dict[str, Any]):
        """Save metadata to file.

        Args:
            metadata (Dict[str, Any]): _description_
        """
        if self.save_metadata:
            file_path = os.path.join(
                self.save_metadata_path, f"{self.name}_metadata.json"
            )
            self.save_to_file(metadata, file_path)

    def load_metadata(self) -> Dict[str, Any]:
        """Load metadata from file.

        Returns:
            Dict[str, Any]: _description_
        """
        file_path = os.path.join(
            self.save_metadata_path, f"{self.name}_metadata.json"
        )
        return self.load_from_file(file_path)

    def log_error(self, error_message: str):
        """Log error to file.

        Args:
            error_message (str): _description_
        """
        file_path = os.path.join(
            self.save_error_path, f"{self.name}_errors.log"
        )
        with open(file_path, "a") as file:
            file.write(f"{error_message}\n")

    def save_artifact(self, artifact: Any, artifact_name: str):
        """Save artifact to file.

        Args:
            artifact (Any): _description_
            artifact_name (str): _description_
        """
        file_path = os.path.join(
            self.save_artifact_path, f"{artifact_name}.json"
        )
        self.save_to_file(artifact, file_path)

    def load_artifact(self, artifact_name: str) -> Any:
        """Load artifact from file.

        Args:
            artifact_name (str): _description_

        Returns:
            Any: _description_
        """
        file_path = os.path.join(
            self.save_artifact_path, f"{artifact_name}.json"
        )
        return self.load_from_file(file_path)

    def _current_timestamp(self):
        """Current timestamp.

        Returns:
            _type_: _description_
        """
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    def log_event(
        self,
        event: str,
        event_type: str = "INFO",
    ):
        """Log event to file.

        Args:
            event (str): _description_
            event_type (str, optional): _description_. Defaults to "INFO".
        """
        timestamp = self._current_timestamp()
        log_message = f"[{timestamp}] [{event_type}] {event}\n"
        file = os.path.join(
            self.save_metadata_path, f"{self.name}_events.log"
        )
        with open(file, "a") as file:
            file.write(log_message)

    async def run_async(self, *args, **kwargs):
        """Run the structure asynchronously."""
        return await asyncio.to_thread(self.run, *args, **kwargs)

    async def save_metadata_async(self, metadata: Dict[str, Any]):
        """Save metadata to file asynchronously.

        Args:
            metadata (Dict[str, Any]): _description_
        """
        return await asyncio.to_thread(self.save_metadata, metadata)

    async def load_metadata_async(self) -> Dict[str, Any]:
        """Load metadata from file asynchronously.

        Returns:
            Dict[str, Any]: _description_
        """
        return await asyncio.to_thread(self.load_metadata)

    async def log_error_async(self, error_message: str):
        """Log error to file asynchronously.

        Args:
            error_message (str): _description_
        """
        return await asyncio.to_thread(self.log_error, error_message)

    async def save_artifact_async(
        self, artifact: Any, artifact_name: str
    ):
        """Save artifact to file asynchronously.

        Args:
            artifact (Any): _description_
            artifact_name (str): _description_
        """
        return await asyncio.to_thread(
            self.save_artifact, artifact, artifact_name
        )

    async def load_artifact_async(self, artifact_name: str) -> Any:
        """Load artifact from file asynchronously.

        Args:
            artifact_name (str): _description_

        Returns:
            Any: _description_
        """
        return await asyncio.to_thread(
            self.load_artifact, artifact_name
        )

    async def log_event_async(
        self,
        event: str,
        event_type: str = "INFO",
    ):
        """Log event to file asynchronously.

        Args:
            event (str): _description_
            event_type (str, optional): _description_. Defaults to "INFO".
        """
        return await asyncio.to_thread(
            self.log_event, event, event_type
        )

    async def asave_to_file(
        self, data: Any, file: str, *args, **kwargs
    ):
        """Save data to file asynchronously.

        Args:
            data (Any): _description_
            file (str): _description_
        """
        await asyncio.to_thread(
            self.save_to_file,
            data,
            file,
            *args,
        )

    async def aload_from_file(
        self,
        file: str,
    ) -> Any:
        """Async load data from file.

        Args:
            file (str): _description_

        Returns:
            Any: _description_
        """
        return await asyncio.to_thread(self.load_from_file, file)

    def run_in_thread(self, *args, **kwargs):
        """Run the structure in a thread."""
        with concurrent.futures.ThreadPoolExecutor() as executor:
            return executor.submit(self.run, *args, **kwargs)

    def save_metadata_in_thread(self, metadata: Dict[str, Any]):
        """Save metadata to file in a thread.

        Args:
            metadata (Dict[str, Any]): _description_
        """
        with concurrent.futures.ThreadPoolExecutor() as executor:
            return executor.submit(self.save_metadata, metadata)

    def run_concurrent(self, *args, **kwargs):
        """Run the structure concurrently."""
        return asyncio.run(self.run_async(*args, **kwargs))

    def compress_data(
        self,
        data: Any,
    ) -> bytes:
        """Compress data.

        Args:
            data (Any): _description_

        Returns:
            bytes: _description_
        """
        return gzip.compress(json.dumps(data).encode())

    def decompres_data(self, data: bytes) -> Any:
        """Decompress data.

        Args:
            data (bytes): _description_

        Returns:
            Any: _description_
        """
        return json.loads(gzip.decompress(data).decode())

    def run_batched(
        self,
        batched_data: List[Any],
        batch_size: int = 10,
        *args,
        **kwargs,
    ):
        """Run batched data.

        Args:
            batched_data (List[Any]): _description_
            batch_size (int, optional): _description_. Defaults to 10.

        Returns:
            _type_: _description_
        """
        with ThreadPoolExecutor(max_workers=batch_size) as executor:
            futures = [
                executor.submit(self.run, data)
                for data in batched_data
            ]
            return [future.result() for future in futures]

    def load_config(
        self, config: str = None, *args, **kwargs
    ) -> Dict[str, Any]:
        """Load config from file.

        Args:
            config (str, optional): _description_. Defaults to None.

        Returns:
            Dict[str, Any]: _description_
        """
        return self.load_from_file(config)

    def backup_data(
        self, data: Any, backup_path: str = None, *args, **kwargs
    ):
        """Backup data to file.

        Args:
            data (Any): _description_
            backup_path (str, optional): _description_. Defaults to None.
        """
        timestamp = self._current_timestamp()
        backup_file_path = f"{backup_path}/{timestamp}.json"
        self.save_to_file(data, backup_file_path)

    def monitor_resources(self):
        """Monitor resource usage."""
        memory = psutil.virtual_memory().percent
        cpu_usage = psutil.cpu_percent(interval=1)
        self.log_event(
            f"Resource usage - Memory: {memory}%, CPU: {cpu_usage}%"
        )

    def run_with_resources(self, *args, **kwargs):
        """Run the structure with resource monitoring."""
        self.monitor_resources()
        return self.run(*args, **kwargs)

    def run_with_resources_batched(
        self,
        batched_data: List[Any],
        batch_size: int = 10,
        *args,
        **kwargs,
    ):
        """Run batched data with resource monitoring.

        Args:
            batched_data (List[Any]): _description_
            batch_size (int, optional): _description_. Defaults to 10.

        Returns:
            _type_: _description_
        """
        self.monitor_resources()
        return self.run_batched(
            batched_data, batch_size, *args, **kwargs
        )

    def _serialize_callable(
        self, attr_value: Callable
    ) -> Dict[str, Any]:
        """
        Serializes callable attributes by extracting their name and docstring.

        Args:
            attr_value (Callable): The callable to serialize.

        Returns:
            Dict[str, Any]: Dictionary with name and docstring of the callable.
        """
        return {
            "name": getattr(
                attr_value, "__name__", type(attr_value).__name__
            ),
            "doc": getattr(attr_value, "__doc__", None),
        }

    def _serialize_attr(self, attr_name: str, attr_value: Any) -> Any:
        """
        Serializes an individual attribute, handling non-serializable objects.

        Args:
            attr_name (str): The name of the attribute.
            attr_value (Any): The value of the attribute.

        Returns:
            Any: The serialized value of the attribute.
        """
        try:
            if callable(attr_value):
                return self._serialize_callable(attr_value)
            elif hasattr(attr_value, "to_dict"):
                return (
                    attr_value.to_dict()
                )  # Recursive serialization for nested objects
            else:
                json.dumps(
                    attr_value
                )  # Attempt to serialize to catch non-serializable objects
                return attr_value
        except (TypeError, ValueError):
            return f"<Non-serializable: {type(attr_value).__name__}>"

    def to_dict(self) -> Dict[str, Any]:
        """
        Converts all attributes of the class, including callables, into a dictionary.
        Handles non-serializable attributes by converting them or skipping them.

        Returns:
            Dict[str, Any]: A dictionary representation of the class attributes.
        """
        return {
            attr_name: self._serialize_attr(attr_name, attr_value)
            for attr_name, attr_value in self.__dict__.items()
        }

    def to_json(self, indent: int = 4, *args, **kwargs):
        return json.dumps(
            self.to_dict(), indent=indent, *args, **kwargs
        )

    def to_yaml(self, indent: int = 4, *args, **kwargs):
        return yaml.dump(
            self.to_dict(), indent=indent, *args, **kwargs
        )

    def to_toml(self, *args, **kwargs):
        return toml.dumps(self.to_dict(), *args, **kwargs)

    # def model_dump_json(self):
    #     logger.info(
    #         f"Saving {self.agent_name} model to JSON in the {self.workspace_dir} directory"
    #     )

    #     create_file_in_folder(
    #         self.workspace_dir,
    #         f"{self.agent_name}.json",
    #         str(self.to_json()),
    #     )

    #     return (
    #         f"Model saved to {self.workspace_dir}/{self.agent_name}.json"
    #     )

    # def model_dump_yaml(self):
    #     logger.info(
    #         f"Saving {self.agent_name} model to YAML in the {self.workspace_dir} directory"
    #     )

    #     create_file_in_folder(
    #         self.workspace_dir,
    #         f"{self.agent_name}.yaml",
    #         self.to_yaml(),
    #     )

    #     return (
    #         f"Model saved to {self.workspace_dir}/{self.agent_name}.yaml"
    #     )
