# lollms_client/lollms_tools_binding.py
from abc import abstractmethod
import importlib
from pathlib import Path
from typing import Optional, List, Dict, Any, Union, Callable
from ascii_colors import trace_exception, ASCIIColors
import yaml
from lollms_client.lollms_base_binding import LollmsBaseBinding

class LollmsToolBinding(LollmsBaseBinding):
    """
    Abstract Base Class for LOLLMS Model Context Protocol (MCP) Bindings.
    """

    def __init__(self,
                 binding_name: str,
                 debug:Optional[bool] = False,
                 **kwargs
                 ):
        """
        Initialize the LollmsToolBinding.
        """
        super().__init__(binding_name=binding_name, debug=debug, **kwargs)
        self.settings = kwargs


    @abstractmethod
    def discover_tools(self, **kwargs) -> List[Dict[str, Any]]:
        """
        Discover available tools compliant with the MCP specification.
        Returns a list of raw tool definitions.
        """
        pass

    @abstractmethod
    def list_tools(self, **kwargs) -> List[Dict[str, Any]]:
        """
        Return a list of tools formatted for consumption by the discussion chat module.

        Each dictionary should contain at least:
            - "name": Unique tool identifier (string)
            - "description": Human‑readable description (string)
            - "parameters": List of parameter specifications, where each parameter is a dict with
                "name", "type", "description", "optional" (bool) and optionally "default".
            - "output": List of output field specifications (optional)

        The returned list can be a subset or a transformed view of the data produced by
        ``discover_tools``; its purpose is to provide a ready‑to‑use schema for the chat
        system.
        """
        pass

    @abstractmethod
    def execute_tool(self,
                     tool_name: str,
                     params: Dict[str, Any],
                     **kwargs) -> Dict[str, Any]:
        """
        Execute a specified tool with the given parameters.
        """
        pass

    # ------------------------------------------------------------------
    # Conversion helper
    # ------------------------------------------------------------------

    def to_chat_tool_specs(self, **discover_kwargs) -> Dict[str, Dict[str, Any]]:
        """
        Convert this binding's tools into the ``tools`` dict format expected by
        ``LollmsDiscussion.chat()``.

        Returns a dict keyed by tool name where every value is a spec dict
        compatible with the ``tools`` parameter of ``chat()``:

        .. code-block:: python

            {
                "alias::tool_name": {
                    "name":        "alias::tool_name",
                    "description": "...",
                    "parameters":  [{"name": "p", "type": "str", ...}, ...],
                    "output":      [],
                    "callable":    <bound execute function>,
                }
            }

        The ``callable`` is a closure over ``self.execute_tool`` so that the
        ``chat()`` agentic loop can call it transparently without knowing about
        the underlying binding.

        Raw MCP ``input_schema`` (JSON Schema) is translated to the flat
        parameter-list format used by ``chat()``.  If a tool already exposes
        the flat ``parameters`` list (e.g. ``local_mcp`` style), it is used
        as-is.
        """
        raw_tools = self.discover_tools(**discover_kwargs)
        specs: Dict[str, Dict[str, Any]] = {}

        for tool in raw_tools:
            name = tool.get("name", "")
            if not name:
                continue

            description = tool.get("description", f"Execute {name}")

            # ── Parameter list ────────────────────────────────────────────────
            # Two formats arrive here:
            #   A) flat list  – already in chat() format: list of {"name","type",...}
            #   B) input_schema – JSON Schema object (from MCP/OpenAPI style bindings)
            if "parameters" in tool and isinstance(tool["parameters"], list):
                # Format A – use as-is
                parameters = tool["parameters"]
            elif "input_schema" in tool and isinstance(tool["input_schema"], dict):
                # Format B – convert JSON Schema properties → flat list
                parameters = _json_schema_to_param_list(tool["input_schema"])
            else:
                parameters = []

            output = tool.get("output", [])

            # ── Callable closure ──────────────────────────────────────────────
            # Capture tool_name by value so the lambda is safe inside the loop.
            def _make_callable(tool_name: str) -> Callable[..., Dict[str, Any]]:
                def _call(**kwargs: Any) -> Dict[str, Any]:
                    try:
                        result = self.execute_tool(tool_name, kwargs)
                        # Normalise: always return a dict
                        if not isinstance(result, dict):
                            result = {"output": result}
                        if "success" not in result:
                            result["success"] = "error" not in result or not result.get("error")
                        return result
                    except Exception as exc:
                        return {"error": str(exc), "success": False}
                return _call

            specs[name] = {
                "name":        name,
                "description": description,
                "parameters":  parameters,
                "output":      output,
                "callable":    _make_callable(name),
                # Keep raw metadata for introspection / show_tools
                "_source":     "personality_tools",
                "_binding":    self.binding_name,
            }

        return specs

    # ------------------------------------------------------------------

    def get_binding_config(self) -> Dict[str, Any]:
        """
        Returns the configuration of the binding.
        """
        return self.settings

    def list_models(self) -> List[Any]:
        """
        For MCP, list_models returns an empty list as it primarily deals with tools.
        Or could be implemented to return available tools as models if desired.
        """
        return []
    
    def get_zoo(self) -> List[Dict[str, Any]]:
        """
        Returns a list of models available for download.
        each entry is a dict with:
        name, description, size, type, link
        """
        return []

    def download_from_zoo(self, index: int, progress_callback: Callable[[dict], None] = None) -> dict:
        """
        Downloads a model from the zoo using its index.
        """
        return {"status": False, "message": "Not implemented"}


# ---------------------------------------------------------------------------
# Internal helper
# ---------------------------------------------------------------------------

def _json_schema_to_param_list(schema: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Convert a JSON Schema ``object`` definition into the flat parameter list
    used by ``chat()``.

    Handles ``properties``, ``required``, ``description``, ``default``, and
    ``enum`` fields.  Nested objects are represented as ``type: object`` with
    no further recursion (the LLM receives the full schema in the description).
    """
    properties: Dict[str, Any] = schema.get("properties", {})
    required_set = set(schema.get("required", []))
    params: List[Dict[str, Any]] = []

    _JSON_TO_PYTHON: Dict[str, str] = {
        "string":  "str",
        "integer": "int",
        "number":  "float",
        "boolean": "bool",
        "array":   "list",
        "object":  "dict",
        "null":    "None",
    }

    for prop_name, prop_schema in properties.items():
        raw_type = prop_schema.get("type", "any")
        # Handle {"type": ["string", "null"]} union style
        if isinstance(raw_type, list):
            non_null = [t for t in raw_type if t != "null"]
            raw_type = non_null[0] if non_null else "any"

        py_type  = _JSON_TO_PYTHON.get(raw_type, raw_type)
        is_optional = prop_name not in required_set

        desc_parts = []
        if prop_schema.get("description"):
            desc_parts.append(prop_schema["description"])
        if prop_schema.get("enum"):
            desc_parts.append(f"Allowed values: {prop_schema['enum']}")

        entry: Dict[str, Any] = {
            "name":        prop_name,
            "type":        py_type,
            "description": " ".join(desc_parts),
            "optional":    is_optional,
        }
        if "default" in prop_schema:
            entry["default"] = prop_schema["default"]

        params.append(entry)

    return params


class LollmsTOOLBindingManager:
    """
    Manages discovery and instantiation of MCP bindings.
    """

    def __init__(self, tools_bindings_dir: Union[str, Path] = Path(__file__).parent / "tools_bindings"):
        """
        Initialize the LollmsTOOLBindingManager.
        """
        self.tools_bindings_dir = Path(tools_bindings_dir)
        if not self.tools_bindings_dir.is_absolute():
            self.tools_bindings_dir = (Path(__file__).parent.parent / tools_bindings_dir).resolve()

        self.available_bindings: Dict[str, type[LollmsToolBinding]] = {}


    def _load_binding_class(self, binding_name: str) -> Optional[type[LollmsToolBinding]]:
        """
        Dynamically load a specific MCP binding class.
        """
        binding_dir = self.tools_bindings_dir / binding_name
        if binding_dir.is_dir():
            init_file_path = binding_dir / "__init__.py"
            if init_file_path.exists():
                try:
                    module_spec = importlib.util.spec_from_file_location(
                        f"lollms_client.tools_bindings.{binding_name}",
                        str(init_file_path)
                    )
                    if module_spec and module_spec.loader:
                        module = importlib.util.module_from_spec(module_spec)
                        module_spec.loader.exec_module(module)
                        
                        if not hasattr(module, 'BindingName'):
                            ASCIIColors.warning(f"Binding '{binding_name}' __init__.py does not define BindingName variable.")
                            return None
                            
                        binding_class_name = module.BindingName
                        if not hasattr(module, binding_class_name):
                            ASCIIColors.warning(f"Binding '{binding_name}' __init__.py defines BindingName='{binding_class_name}', but class not found.")
                            return None

                        binding_class = getattr(module, binding_class_name)
                        if not issubclass(binding_class, LollmsToolBinding):
                             ASCIIColors.warning(f"Class {binding_class_name} in {binding_name} is not a subclass of LollmsToolBinding.")
                             return None
                        return binding_class
                    else:
                        ASCIIColors.warning(f"Could not create module spec for MCP binding '{binding_name}'.")
                except Exception as e:
                    ASCIIColors.error(f"Failed to load MCP binding '{binding_name}': {e}")
                    trace_exception(e)
        return None

    def create_binding(self,
                       binding_name: str,
                       **kwargs
                       ) -> Optional[LollmsToolBinding]:
        """
        Create an instance of a specific MCP binding.
        """
        if binding_name not in self.available_bindings:
            binding_class = self._load_binding_class(binding_name)
            if binding_class:
                self.available_bindings[binding_name] = binding_class
            else:
                ASCIIColors.error(f"TOOLS binding '{binding_name}' class not found or failed to load.")
                return None
        
        binding_class_to_instantiate = self.available_bindings.get(binding_name)
        if binding_class_to_instantiate:
            try:
                return binding_class_to_instantiate(
                    **kwargs
                )
            except Exception as e:
                ASCIIColors.error(f"Failed to instantiate MCP binding '{binding_name}': {e}")
                trace_exception(e)
                return None
        return None


    def get_available_bindings(self) -> List[str]:
        """
        Return list of available MCP binding names.
        """
        available = []
        if self.tools_bindings_dir.is_dir():
            for item in self.tools_bindings_dir.iterdir():
                if item.is_dir() and (item / "__init__.py").exists():
                    available.append(item.name)
        return available

    def get_binding_description(self, binding_name: str) -> Optional[Dict[str, Any]]:
        """
        Loads and returns the content of the description.yaml file for a given binding.
        """
        binding_dir = self.tools_bindings_dir / binding_name
        description_file = binding_dir / "description.yaml"

        if not description_file.exists():
            ASCIIColors.warning(f"No description.yaml found for MCP binding '{binding_name}'.")
            return None

        try:
            with open(description_file, 'r', encoding='utf-8') as f:
                description = yaml.safe_load(f)
            return description
        except yaml.YAMLError as e:
            ASCIIColors.error(f"Error parsing description.yaml for MCP binding '{binding_name}': {e}")
            trace_exception(e)
            return None
        except Exception as e:
            ASCIIColors.error(f"Error reading description.yaml for MCP binding '{binding_name}': {e}")
            trace_exception(e)
            return None

def list_binding_models(tools_binding_name: str, tools_binding_config: Optional[Dict[str, any]]|None = None, tools_bindings_dir: str|Path = Path(__file__).parent / "tools_bindings") -> List[Dict]:
    """
    Lists all available models/tools for a specific binding.
    """
    binding = LollmsTOOLBindingManager(tools_bindings_dir).create_binding(
        binding_name=tools_binding_name,
        **{
            k: v
            for k, v in (tools_binding_config or {}).items()
            if k != "binding_name"
        }
    )

    return binding.discover_tools() if binding else []