# hiveos/utilities/wrapper_serialization.py

import inspect
import ast
import uuid
import textwrap
import re
import json
import copy
import hmac
import hashlib

REQUIRED_METHODS = [
    "execute_chunk",
    "get_status",
    "report_capabilities",
    "health_check",
    "enter_cooldown",
    "is_available",
    "get_summary",
]

def serialize_wrapper(cls):
    wrapper_spec = {
        "name": cls.__name__,
        "base_class": cls.__bases__[0].__name__,
        "capabilities": getattr(cls, "capabilities", []),
        "intents": getattr(cls, "intents", []),
        "methods": {},
        "metadata": {
            "description": (getattr(cls, "__doc__", "") or "").strip(),
            "version": "0.1.0",
            "origin": "manual"
        }
    }

    for method_name in REQUIRED_METHODS:
        method = cls.__dict__.get(method_name, None)
        if method is None:
            continue  # Skip inherited methods

        try:
            method_src = inspect.getsource(method)
            method_src = textwrap.dedent(method_src)

            # 🔁 Replace super() call with explicit base call
            base_class_name = wrapper_spec["base_class"]

            # Replace super().<method>() with BaseClass.<method>(self, ...)
            if f"super().{method_name}" in method_src:
                method_src = method_src.replace(
                    f"super().{method_name}(",
                    f"{base_class_name}.{method_name}(self, "
                )

            # If base class method is explicitly called but missing self
            if f"{base_class_name}.{method_name}(" in method_src:
                # Regex patch: Add 'self' if not already present
                pattern = rf"{base_class_name}\.{method_name}\((?!self)"
                method_src = re.sub(pattern, f"{base_class_name}.{method_name}(self, ", method_src)


            wrapper_spec["methods"][method_name] = method_src

        except (TypeError, OSError):
            print(f"[Serialize] Skipping method: {method_name} (unavailable source)")

    return wrapper_spec


def hydrate_wrapper(wrapper_spec: dict):
    """
    Reconstructs a live agent wrapper from a serialized wrapper_spec dictionary.
    Returns an instance of the hydrated class.
    """
    name = wrapper_spec.get("name", "HydratedWrapper")
    base_class_name = wrapper_spec.get("base_class", "AgentInterface")
    methods = wrapper_spec.get("methods", {})

    from hiveos.agent_interface import AgentInterface

    # Inject AgentInterface into globals for exec()
    exec_globals = globals().copy()
    exec_globals["AgentInterface"] = AgentInterface

    method_namespace = {}


    # Compile each method into the namespace
    for method_name, method_src in methods.items():
        # Fix super() usage if it exists
        if f"super().{method_name}" in method_src:
            method_src = method_src.replace(
                f"super().{method_name}(",
                f"{base_class_name}.{method_name}(self, "
            )

        # Add 'self' if missing in direct base class calls
        if f"{base_class_name}.{method_name}(" in method_src:
            pattern = rf"{base_class_name}\.{method_name}\((?!self)"
            method_src = re.sub(pattern, f"{base_class_name}.{method_name}(self, ", method_src)

        try:
            exec(method_src, exec_globals, method_namespace)
        except Exception as e:
            print(f"[hydrate_wrapper] Failed to compile method '{method_name}': {e}")

    # Build final method map: pull only actual callables
    callable_methods = {
        k: v for k, v in method_namespace.items()
        if callable(v) and k in methods
    }

    # Dynamically import the base class (defaults to AgentInterface)
    from hiveos.agent_interface import AgentInterface  # This may change if dynamic import needed
    base_class = AgentInterface  # Support for more base classes can be added later

    # Create the class dynamically
    HydratedWrapperClass = type(name, (base_class,), callable_methods)

    agent = HydratedWrapperClass()

    # Ensure dynamic summary attributes exist
    for attr in ("latest_output", "latest_input", "current_intent"):
        if not hasattr(agent, attr):
            setattr(agent, attr, None)

    # Pull type from wrapper_spec['name'] if provided
    agent_type = wrapper_spec.get("name", "HydratedWrapper")
    uid = uuid.uuid4().hex[:8]

    agent.agent_id = f"{agent_type}-{uid}"
    agent.name = f"{agent_type}-wrap"

    # Finalize + normalize
    # agent.agent_id = getattr(agent, "agent_id", None) or f"wrapper-{uuid.uuid4().hex[:8]}"
    # agent.name = getattr(agent, "name", None) or f"{type(agent).__name__}-{agent.agent_id[:4]}"
    agent.status = getattr(agent, "status", None) or "idle"
    agent.zone = getattr(agent, "zone", None) or None
    agent._available = getattr(agent, "_available", None)
    if agent._available is None:
        agent._available = True
    if not hasattr(agent, "capabilities") or not agent.capabilities:
        agent.capabilities = wrapper_spec.get("capabilities", ["generic"])

    return agent

def generate_wrapper_py(wrapper_spec: dict) -> str:
    """
    Converts a WrapperSpec into a valid Python class definition.
    """
    class_name = wrapper_spec.get("name", "GeneratedWrapper")
    base_class = wrapper_spec.get("base_class", "AgentInterface")
    methods = wrapper_spec.get("methods", {})
    meta = wrapper_spec.get("metadata", {})

    # Header comments
    header = [
        "# Generated by HiveOS Wrapper Wizard",
        f"# Description: {meta.get('description', '')}",
        f"# Version: {meta.get('version', '0.1.0')}",
        f"# Origin: {meta.get('origin', 'manual')}",
        "",
        f"from hiveos.agent_interface import {base_class}",
        "",
        f"class {class_name}({base_class}):"
    ]

    # Handle method source code
    method_blocks = []
    for name, src in methods.items():
        if not src.strip().startswith("def "):
            continue  # Skip invalid methods

        # Indent each line (to sit inside class)
        indented = ["    " + line if line.strip() else "" for line in src.splitlines()]
        method_blocks.append("\n".join(indented))

    # Handle empty case
    if not method_blocks:
        method_blocks.append("    pass")

    return "\n".join(header + method_blocks) + "\n"

def validate_py_wrapper(py_source: str) -> dict:
    """
    Parses and validates a .py source string as a potential wrapper class.
    Returns a tuple (is_valid: bool, result: WrapperSpec | error_str).
    """
    try:
        tree = ast.parse(py_source)
    except SyntaxError as e:
        return (False, f"Syntax error: {e}")

    class_defs = [node for node in tree.body if isinstance(node, ast.ClassDef)]
    if not class_defs:
        return (False, "No class definition found.")

    if len(class_defs) > 1:
        return (False, "Multiple classes found. Only one wrapper per file allowed.")

    cls = class_defs[0]
    class_name = cls.name
    base_classes = [b.id for b in cls.bases if isinstance(b, ast.Name)]

    if "AgentInterface" not in base_classes:
        return (False, f"Class {class_name} does not inherit from AgentInterface.")

    method_names = [n.name for n in cls.body if isinstance(n, ast.FunctionDef)]
    missing_methods = [m for m in REQUIRED_METHODS if m not in method_names]
    if missing_methods:
        return (False, f"Missing required methods: {missing_methods}")

    # Extract methods as raw strings
    methods = {}
    for node in cls.body:
        if isinstance(node, ast.FunctionDef):
            start_line = node.lineno - 1
            end_line = node.end_lineno
            method_lines = py_source.splitlines()[start_line:end_line]
            method_src = "\n".join(method_lines)
            methods[node.name] = textwrap.dedent(method_src)


    wrapper_spec = {
        "name": class_name,
        "base_class": "AgentInterface",
        "capabilities": [],
        "intents": [],
        "methods": methods,
        "metadata": {
            "description": f"Validated wrapper from {class_name}",
            "version": "0.1.0",
            "origin": "uploaded"
        }
    }

    return (True, wrapper_spec)


def validate_wrapper_spec_policy(wrapper_spec, allowlist=None, require_signature=False, signing_key=""):
    if not isinstance(wrapper_spec, dict):
        return (False, "wrapper_spec must be an object")

    name = str(wrapper_spec.get("name") or "").strip()
    if not name:
        return (False, "wrapper_spec.name is required")

    allowlist = list(allowlist or [])
    if allowlist and name not in allowlist:
        return (False, f"wrapper_spec '{name}' is not allowlisted")

    if require_signature:
        if not signing_key:
            return (False, "wrapper spec signature is required but signing key is not configured")
        metadata = wrapper_spec.get("metadata")
        if not isinstance(metadata, dict):
            return (False, "wrapper_spec.metadata must be an object when signature is required")
        signature = str(metadata.get("signature") or "").strip()
        if not signature:
            return (False, "wrapper_spec.metadata.signature is required")
        expected = compute_wrapper_spec_signature(wrapper_spec, signing_key)
        if not hmac.compare_digest(signature, expected):
            return (False, "wrapper spec signature mismatch")

    return (True, "")


def compute_wrapper_spec_signature(wrapper_spec, signing_key):
    if not isinstance(wrapper_spec, dict):
        raise ValueError("wrapper_spec must be an object")
    if not signing_key:
        raise ValueError("signing_key is required")

    canonical = _canonical_wrapper_spec_json(wrapper_spec)
    digest = hmac.new(
        signing_key.encode("utf-8"),
        canonical.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    return digest


def _canonical_wrapper_spec_json(wrapper_spec):
    cloned = copy.deepcopy(wrapper_spec)
    metadata = cloned.get("metadata")
    if isinstance(metadata, dict):
        metadata.pop("signature", None)
    return json.dumps(cloned, sort_keys=True, separators=(",", ":"))
