"""
Zubbl SDK - Main Client

The ZubblClient is the primary interface for integrating Zubbl with any AI agent.
It provides a simple 3-line integration that enables automatic learning and improvement.

Example:
    from zubbl import ZubblClient

    zubbl = ZubblClient(api_key="your-api-key")
    agent = zubbl.wrap(your_agent)  # That's it!
"""

import asyncio
import hashlib
import uuid
from datetime import datetime
from functools import wraps
from typing import Any, Callable, Dict, List, Optional

import httpx
from pydantic import BaseModel, Field

from zubbl.interceptor import FailureInterceptor, RecoveryResult
from zubbl.types import (
    Feedback,
    Policy,
    RLMode,
    Step,
    TaskStatus,
    Trajectory,
    UsageStats,
)


class ZubblConfig(BaseModel):
    """Configuration for ZubblClient"""

    api_key: str
    base_url: str = "https://api.zubbl.ai"
    rl_mode: RLMode = RLMode.LOCAL
    auto_ingest: bool = True  # Automatically ingest trajectories
    auto_query: bool = True  # Automatically query for policies before tasks
    auto_recover: bool = True  # Automatically recover from failures
    timeout: float = 30.0
    retry_attempts: int = 3
    local_cache: bool = True  # Cache policies locally for speed
    max_recovery_retries: int = 3  # Max retries for auto-recovery
    fallback_models: List[str] = Field(default_factory=list)  # Fallback models for recovery


class ZubblClient:
    """
    Zubbl Client - Make any AI agent self-evolving.

    Usage:
        # Basic (3 lines)
        from zubbl import ZubblClient
        zubbl = ZubblClient(api_key="your-api-key")
        agent = zubbl.wrap(your_agent)

        # Advanced
        zubbl = ZubblClient(
            api_key="your-api-key",
            rl_mode=RLMode.COREWEAVE,  # Use cloud RL
        )
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.zubbl.ai",
        rl_mode: RLMode = RLMode.LOCAL,
        auto_ingest: bool = True,
        auto_query: bool = True,
        timeout: float = 30.0,
        # SDK metadata headers (all optional, backward compatible)
        sdk_version: str = "0.1.0",
        agent_framework: Optional[str] = None,
        model: Optional[str] = None,
        environment: Optional[str] = None,
        **kwargs,
    ):
        self.config = ZubblConfig(
            api_key=api_key,
            base_url=base_url,
            rl_mode=rl_mode,
            auto_ingest=auto_ingest,
            auto_query=auto_query,
            timeout=timeout,
            **kwargs,
        )

        # Build headers with auth + SDK metadata
        headers = {"Authorization": f"Bearer {api_key}"}
        if sdk_version:
            headers["X-SDK-Version"] = sdk_version
        if agent_framework:
            headers["X-Agent-Framework"] = agent_framework
        if model:
            headers["X-Model"] = model
        if environment:
            headers["X-Environment"] = environment

        self._client = httpx.AsyncClient(
            base_url=base_url,
            headers=headers,
            timeout=timeout,
        )
        self._sync_client = httpx.Client(
            base_url=base_url,
            headers=headers,
            timeout=timeout,
        )
        self._policy_cache: Dict[str, Policy] = {}
        self._current_trajectory: Optional[Trajectory] = None
        self._steps_buffer: List[Step] = []
        self._current_policy: Optional[Policy] = None
        self._interceptor = FailureInterceptor(
            max_retries=kwargs.get("max_recovery_retries", 3),
            fallback_models=kwargs.get("fallback_models", []),
        )
        self._last_recovery: Optional[RecoveryResult] = None

    @property
    def current_policy(self) -> Optional[Policy]:
        """Get the policy that was queried for the current task.

        Use this to access recommendations without them being
        injected into agent method kwargs.
        """
        return self._current_policy

    @property
    def last_recovery(self) -> Optional[RecoveryResult]:
        """Get the result of the last auto-recovery attempt.

        Check this after a wrapped agent call to see if Zubbl
        intercepted a failure and recovered automatically.

        Example:
            result = agent.run("classify this image")
            if zubbl.last_recovery and zubbl.last_recovery.fix_applied:
                print(f"Zubbl auto-fixed: {zubbl.last_recovery.fix_applied}")
        """
        return self._last_recovery

    @property
    def interceptor(self) -> FailureInterceptor:
        """Access the failure interceptor for manual pattern management"""
        return self._interceptor

    # =========================================================================
    # Core API Methods
    # =========================================================================

    def wrap(self, agent: Any, auto_recover: Optional[bool] = None) -> Any:
        """
        Wrap an agent to enable automatic learning AND failure recovery.

        This is the main entry point - wrapping an agent automatically:
        1. Queries for relevant policies before each task
        2. Records the execution trajectory
        3. If agent fails: auto-recovers using known fix patterns
        4. Ingests all trajectories (success + failure) for learning
        5. Over time, pre-applies fixes before failures even happen

        Args:
            agent: Any AI agent (LangGraph, CrewAI, AutoGPT, custom, etc.)
            auto_recover: Override auto_recover setting for this agent.
                         When True, failures trigger automatic retry/fix.
                         Defaults to config.auto_recover.

        Returns:
            Wrapped agent with Zubbl learning + recovery capabilities

        Example:
            zubbl = ZubblClient(api_key="your-key")
            agent = zubbl.wrap(your_agent)  # auto_recover=True by default

            # Agent fails internally → Zubbl retries with a fix
            # → returns correct result. No code change. No redeploy.
        """
        recover = auto_recover if auto_recover is not None else self.config.auto_recover
        return ZubblWrappedAgent(self, agent, auto_recover=recover)

    def wrap_crew(self, crew: Any, agent_names: Optional[List[str]] = None) -> Any:
        """
        Wrap a multi-agent crew for automatic learning.

        Works with CrewAI, AutoGen, or any custom multi-agent system.
        Records the entire crew execution as a trajectory with per-agent steps.

        Args:
            crew: The crew/team object (e.g., CrewAI Crew)
            agent_names: Optional list of agent names for tracking

        Returns:
            Wrapped crew with Zubbl learning capabilities

        Example:
            smart_crew = zubbl.wrap_crew(crew, agent_names=["researcher", "writer"])
            result = smart_crew.kickoff()
        """
        return ZubblWrappedCrew(self, crew, agent_names=agent_names)

    def start_trajectory(self, task_description: str, metadata: Optional[Dict] = None) -> str:
        """Start recording a new trajectory"""
        trajectory_id = f"traj_{uuid.uuid4().hex[:16]}"
        self._current_trajectory = Trajectory(
            trajectory_id=trajectory_id,
            tenant_id=self._extract_tenant_id(),
            task_description=task_description,
            steps=[],
            status=TaskStatus.PARTIAL,
            metadata=metadata or {},
        )
        self._steps_buffer = []
        return trajectory_id

    def record_step(
        self,
        action: str,
        tool_name: Optional[str] = None,
        tool_input: Optional[Dict] = None,
        tool_output: Optional[Any] = None,
        reasoning: Optional[str] = None,
        tokens_used: int = 0,
        latency_ms: int = 0,
    ) -> None:
        """Record a single step in the current trajectory"""
        if self._current_trajectory is None:
            raise RuntimeError("No active trajectory. Call start_trajectory() first.")

        step = Step(
            step_id=f"step_{len(self._steps_buffer) + 1}",
            action=action,
            tool_name=tool_name,
            tool_input=tool_input,
            tool_output=tool_output,
            reasoning=reasoning,
            tokens_used=tokens_used,
            latency_ms=latency_ms,
        )
        self._steps_buffer.append(step)

    def end_trajectory(
        self,
        status: TaskStatus,
        actual_output: Optional[Any] = None,
        expected_output: Optional[Any] = None,
        agent_confidence: Optional[float] = None,
    ) -> Trajectory:
        """End the current trajectory and optionally ingest it.

        Args:
            status: Task execution status
            actual_output: What the agent actually produced (for validation)
            expected_output: What the correct answer should be (for validation)
            agent_confidence: Agent's self-reported confidence (0.0-1.0)
        """
        if self._current_trajectory is None:
            raise RuntimeError("No active trajectory.")

        self._current_trajectory.steps = self._steps_buffer
        self._current_trajectory.status = status
        self._current_trajectory.total_tokens = sum(s.tokens_used for s in self._steps_buffer)
        self._current_trajectory.total_latency_ms = sum(s.latency_ms for s in self._steps_buffer)
        self._current_trajectory.expected_output = expected_output
        self._current_trajectory.actual_output = actual_output
        self._current_trajectory.agent_confidence = agent_confidence

        trajectory = self._current_trajectory

        # Auto-ingest ALL trajectories (not just success)
        # Failures carry learning signal too, especially with expected_output
        if self.config.auto_ingest:
            self.ingest(trajectory)

        self._current_trajectory = None
        self._steps_buffer = []

        return trajectory

    # =========================================================================
    # API Endpoints
    # =========================================================================

    def ingest(self, trajectory: Trajectory) -> Dict[str, Any]:
        """
        Ingest a trajectory into Zubbl for learning.

        POST /ingest
        """
        response = self._sync_client.post(
            "/ingest",
            json=trajectory.model_dump(mode="json"),
        )
        response.raise_for_status()
        return response.json()

    async def ingest_async(self, trajectory: Trajectory) -> Dict[str, Any]:
        """Async version of ingest"""
        response = await self._client.post(
            "/ingest",
            json=trajectory.model_dump(mode="json"),
        )
        response.raise_for_status()
        return response.json()

    def query(self, task_description: str, context: Optional[Dict] = None) -> Optional[Policy]:
        """
        Query for a policy recommendation before starting a task.

        GET /query
        """
        task_hash = self._hash_task(task_description)

        # Check cache first
        if self.config.local_cache and task_hash in self._policy_cache:
            return self._policy_cache[task_hash]

        response = self._sync_client.get(
            "/query",
            params={
                "task": task_description,
                "context": str(context) if context else None,
                "rl_mode": self.config.rl_mode.value,
            },
        )

        if response.status_code == 404:
            return None  # No policy found for this task type

        response.raise_for_status()
        data = response.json()

        if data:
            policy = Policy(**data)
            if self.config.local_cache:
                self._policy_cache[task_hash] = policy
            return policy
        return None

    async def query_async(
        self, task_description: str, context: Optional[Dict] = None
    ) -> Optional[Policy]:
        """Async version of query"""
        task_hash = self._hash_task(task_description)

        if self.config.local_cache and task_hash in self._policy_cache:
            return self._policy_cache[task_hash]

        response = await self._client.get(
            "/query",
            params={
                "task": task_description,
                "context": str(context) if context else None,
                "rl_mode": self.config.rl_mode.value,
            },
        )

        if response.status_code == 404:
            return None

        response.raise_for_status()
        data = response.json()

        if data:
            policy = Policy(**data)
            if self.config.local_cache:
                self._policy_cache[task_hash] = policy
            return policy
        return None

    def feedback(
        self, trajectory_id: str, rating: float, comment: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Submit feedback for a trajectory.

        POST /feedback
        """
        feedback = Feedback(
            feedback_id=f"fb_{uuid.uuid4().hex[:12]}",
            trajectory_id=trajectory_id,
            rating=rating,
            comment=comment,
        )

        response = self._sync_client.post(
            "/feedback",
            json=feedback.model_dump(mode="json"),
        )
        response.raise_for_status()
        return response.json()

    def validate(
        self,
        trajectory_id: str,
        expected_output: Any,
        actual_output: Optional[Any] = None,
        agent_confidence: float = 0.0,
    ) -> Dict[str, Any]:
        """
        Post-hoc validation of a trajectory's output.

        Use this when the ground truth becomes available after the agent ran.
        Detects confident-but-wrong cases automatically.

        POST /validate

        Args:
            trajectory_id: ID of the trajectory to validate
            expected_output: The correct answer
            actual_output: What the agent produced (fetched from DB if not provided)
            agent_confidence: Agent's self-reported confidence (0.0-1.0)

        Returns:
            Dict with correctness_score, is_confidence_mismatch, etc.
        """
        response = self._sync_client.post(
            "/validate",
            json={
                "trajectory_id": trajectory_id,
                "expected_output": expected_output,
                "actual_output": actual_output,
                "agent_confidence": agent_confidence,
            },
        )
        response.raise_for_status()
        return response.json()

    def get_usage(self) -> UsageStats:
        """Get current usage statistics"""
        response = self._sync_client.get("/usage")
        response.raise_for_status()
        return UsageStats(**response.json())

    def sync_failure_patterns(self) -> Dict[str, Any]:
        """
        Sync failure patterns with the Zubbl server.

        Uploads local patterns learned this session, downloads patterns
        from previous sessions (including fixes that worked).

        Call this on startup to pre-load known fixes, and periodically
        to share new learnings with the server.

        Returns:
            Dict with sync stats (total patterns, patterns with fixes, etc.)
        """
        # Export local patterns
        local_patterns = self._interceptor.export_patterns()

        try:
            response = self._sync_client.post(
                "/failure-patterns/sync",
                json={"patterns": local_patterns},
            )
            response.raise_for_status()
            data = response.json()

            # Load server patterns into interceptor
            self._interceptor.load_patterns_from_server(data.get("patterns", []))
            self._interceptor.load_prompt_patches(data.get("prompt_patches", {}))

            return data.get("stats", {})
        except Exception:
            # Sync failure is non-critical — interceptor works locally
            return {"synced": False, "local_patterns": len(local_patterns)}

    def get_failure_stats(self) -> Dict[str, Any]:
        """
        Get failure recovery statistics.

        Shows how many failures Zubbl has auto-recovered,
        what fix types work best, and failure categories.
        """
        try:
            response = self._sync_client.get("/failure-patterns/stats")
            response.raise_for_status()
            return response.json()
        except Exception:
            return {"error": "Could not fetch stats"}

    # =========================================================================
    # Helper Methods
    # =========================================================================

    def _extract_tenant_id(self) -> str:
        """Extract tenant ID from API key (format: zubbl_<tenant_id>_<secret>)"""
        parts = self.config.api_key.split("_")
        if len(parts) >= 3:
            return "_".join(parts[1:-1])
        if len(parts) == 2:
            return parts[1]
        return hashlib.sha256(self.config.api_key.encode()).hexdigest()[:16]

    def _hash_task(self, task_description: str) -> str:
        """Create a hash for task caching"""
        normalized = task_description.lower().strip()
        return hashlib.md5(normalized.encode()).hexdigest()[:12]

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._sync_client.close()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self._client.aclose()


class ZubblWrappedAgent:
    """
    A wrapped agent that automatically:
    1. Records execution trajectories
    2. Queries policies before execution
    3. Intercepts failures and auto-recovers
    4. Learns from both successes and failures

    The auto_recover flag is the key difference from a simple wrapper:
    - auto_recover=False: Record + Learn only (old behavior)
    - auto_recover=True: Record + Learn + Fix failures in real-time (the product)
    """

    def __init__(self, client: ZubblClient, agent: Any, auto_recover: bool = True):
        self._client = client
        self._agent = agent
        self._auto_recover = auto_recover
        self._policy_cache = {}      # category → policy
        self._cache_expiry = {}      # category → timestamp
        self._CACHE_TTL = 300        # 5 minutes

    def _get_cached_policy(self, task_desc: str):
        """Get policy from cache or fetch (saves ~500ms per call)."""
        import time

        from zubbl.task_utils import normalize_task_type
        full_type = normalize_task_type(task_desc)
        category = full_type.split("_")[0] if "_" in full_type else full_type

        now = time.time()
        if category in self._policy_cache and now < self._cache_expiry.get(category, 0):
            return self._policy_cache[category]

        policy = self._client.query(task_desc)
        self._policy_cache[category] = policy
        self._cache_expiry[category] = now + self._CACHE_TTL
        return policy

    def __getattr__(self, name: str) -> Any:
        """Proxy all attribute access to the underlying agent"""
        attr = getattr(self._agent, name)

        # Wrap callable methods to record trajectories + auto-recover
        if callable(attr):
            return self._wrap_method(attr, name)
        return attr

    def __call__(self, *args, **kwargs):
        """Make the wrapped agent callable with auto-recovery"""
        task_desc = str(args[0]) if args else "default_task"

        # Query for policy (cached — avoids ~500ms network call on repeated tasks)
        policy = None
        if self._client.config.auto_query:
            policy = self._get_cached_policy(task_desc)
            self._client._current_policy = policy

        # Inject learned policy into the task (Training-Free GRPO approach)
        # Compact experiential context: <100 tokens, prepended to input
        enhanced_args = self._inject_policy(args, policy)

        # Start trajectory
        self._client.start_trajectory(task_desc)

        if self._auto_recover:
            return self._call_with_recovery(task_desc, enhanced_args, kwargs)
        else:
            return self._call_without_recovery(enhanced_args, kwargs)

    def _inject_policy(self, args: tuple, policy) -> tuple:
        """Inject learned policy into agent input as compact prefix.

        Based on Training-Free GRPO (Cai et al., 2025): distill experiences
        into compact natural-language priors (<100 tokens) prepended to input.
        This improves the agent WITHOUT changing any code.

        Example: "Write a function to sort" becomes:
        "[Strategy: plan first, use code_generator, validate output] Write a function to sort"
        """
        if not policy or not policy.recommended_actions or not args:
            return args

        # Build compact strategy prefix (<100 tokens)
        actions = policy.recommended_actions[:3]  # Max 3 recommendations
        strategy = "[Strategy: " + ", ".join(
            a.replace("Use ", "").replace(" for ", "→")
            for a in actions
        ) + "] "

        # Prepend to first argument (the task/prompt)
        first_arg = str(args[0])
        enhanced = strategy + first_arg
        return (enhanced,) + args[1:]

    def _call_with_recovery(self, task_desc: str, args: tuple, kwargs: dict) -> Any:
        """Execute with auto-recovery — the core product behavior"""
        from zubbl.task_utils import normalize_task_type

        task_type = normalize_task_type(task_desc)

        # Use the interceptor for automatic failure recovery
        recovery = self._client._interceptor.intercept_and_recover(
            func=self._agent,
            args=args,
            kwargs=kwargs,
            task_description=task_desc,
            task_type=task_type,
        )

        # Store recovery result for inspection
        self._client._last_recovery = recovery

        if recovery.success:
            output = str(recovery.recovered_output or "")

            # Auto-extract rich steps from input + output
            self._record_rich_steps(task_desc, output, recovery.latency_ms)

            # If a fix was applied, record it as a separate step
            if recovery.fix_applied:
                self._client.record_step(
                    action="auto_recovery",
                    tool_name="zubbl_interceptor",
                    tool_input={
                        "fix_type": recovery.fix_applied.value,
                        "original_error": recovery.original_error,
                    },
                    tool_output="recovery_successful",
                    reasoning=f"Auto-recovered from: {recovery.original_error}",
                )

            self._client.end_trajectory(TaskStatus.SUCCESS)
            return recovery.recovered_output
        else:
            self._client.record_step(
                action="execute",
                tool_name="agent",
                tool_output=f"ERROR: {recovery.original_error}",
                latency_ms=recovery.latency_ms,
            )
            self._client.end_trajectory(TaskStatus.FAILURE)
            raise RuntimeError(
                f"Agent failed and auto-recovery exhausted "
                f"({recovery.attempts} attempts): {recovery.original_error}"
            )

    def _record_rich_steps(self, task_desc: str, output: str, total_latency_ms: int):
        """Auto-extract meaningful steps from task description and output.

        This is what makes wrap() powerful — users get rich trajectory data
        without writing any recording code. The steps are inferred from:
        1. Task description keywords → what the agent was asked to do
        2. Output analysis → what the agent produced
        """
        task_lower = task_desc.lower()
        output_lower = output.lower()
        step_ms = total_latency_ms // 3 if total_latency_ms else 100

        # Step 1: Analyze task (what was the agent asked to do?)
        task_action = "analyze"
        if any(w in task_lower for w in ["write", "create", "generate", "build", "make"]):
            task_action = "generate"
        elif any(w in task_lower for w in ["fix", "debug", "repair", "resolve", "patch"]):
            task_action = "diagnose"
        elif any(w in task_lower for w in ["review", "check", "audit", "scan", "inspect"]):
            task_action = "review"
        elif any(w in task_lower for w in ["optimize", "improve", "speed", "reduce", "refactor"]):
            task_action = "optimize"
        elif any(w in task_lower for w in ["explain", "describe", "what is", "how does"]):
            task_action = "explain"
        elif any(w in task_lower for w in ["test", "validate", "verify"]):
            task_action = "test"
        elif any(w in task_lower for w in ["deploy", "release", "publish", "ship"]):
            task_action = "deploy"

        self._client.record_step(
            action="understand_task",
            tool_name="planner",
            tool_input={"task": task_desc[:200]},
            tool_output=f"Task type: {task_action}",
            latency_ms=step_ms,
        )

        # Step 2: Execute (the actual agent work)
        # Detect tools/techniques used from output
        tools_detected = []
        if any(w in output_lower for w in ["def ", "function", "class ", "import "]):
            tools_detected.append("code_generator")
        if any(w in output_lower for w in ["error", "exception", "bug", "fix"]):
            tools_detected.append("debugger")
        if any(w in output_lower for w in ["test", "assert", "expect"]):
            tools_detected.append("test_runner")
        if any(w in output_lower for w in ["select", "query", "index", "table"]):
            tools_detected.append("sql_analyzer")
        if any(w in output_lower for w in ["o(n", "complexity", "optimize", "cache"]):
            tools_detected.append("performance_analyzer")
        if not tools_detected:
            tools_detected.append("llm")

        self._client.record_step(
            action=task_action,
            tool_name=tools_detected[0],
            tool_input={"approach": task_action},
            tool_output=output[:300],
            tokens_used=len(output.split()) * 2,
            latency_ms=step_ms,
        )

        # Step 3: Validate output quality
        has_code = "```" in output or "def " in output or "function" in output
        has_explanation = len(output) > 100
        quality = "complete" if (has_code or has_explanation) else "partial"

        self._client.record_step(
            action="validate_output",
            tool_name="quality_checker",
            tool_output=f"Quality: {quality}, length: {len(output)} chars, tools: {tools_detected}",
            latency_ms=step_ms,
        )

    def _call_without_recovery(self, args: tuple, kwargs: dict) -> Any:
        """Execute without auto-recovery"""
        task_desc = str(args[0]) if args else "task"

        try:
            import time
            start = time.time()
            result = self._agent(*args, **kwargs)
            latency = int((time.time() - start) * 1000)

            self._record_rich_steps(task_desc, str(result or ""), latency)
            self._client.end_trajectory(TaskStatus.SUCCESS)
            return result
        except Exception as e:
            self._client.record_step(
                action="agent_error",
                tool_output=str(e),
            )
            self._client.end_trajectory(TaskStatus.FAILURE)
            raise

    def _wrap_method(self, method: Callable, method_name: str) -> Callable:
        """Wrap a method to record its execution + auto-recover from failures"""

        @wraps(method)
        def wrapper(*args, **kwargs):
            task_desc = kwargs.get("task", str(args[0]) if args else method_name)

            if self._client.config.auto_query:
                policy = self._client.query(task_desc)
                self._client._current_policy = policy

            self._client.start_trajectory(task_desc)

            if self._auto_recover:
                return self._method_with_recovery(method, method_name, task_desc, args, kwargs)
            else:
                return self._method_without_recovery(method, method_name, args, kwargs)

        @wraps(method)
        async def async_wrapper(*args, **kwargs):
            task_desc = kwargs.get("task", str(args[0]) if args else method_name)

            if self._client.config.auto_query:
                policy = await self._client.query_async(task_desc)
                self._client._current_policy = policy

            self._client.start_trajectory(task_desc)

            # Async recovery: run sync interceptor in executor
            # (native async interceptor)
            try:
                start_time = datetime.utcnow()
                result = await method(*args, **kwargs)
                latency = int((datetime.utcnow() - start_time).total_seconds() * 1000)

                self._client.record_step(
                    action=method_name,
                    tool_input={"args": str(args)[:200], "kwargs": str(kwargs)[:200]},
                    tool_output=str(result)[:1000],
                    latency_ms=latency,
                )
                self._client.end_trajectory(TaskStatus.SUCCESS)
                return result

            except Exception as e:
                self._client.record_step(
                    action=method_name,
                    tool_input={"args": str(args)[:200], "kwargs": str(kwargs)[:200]},
                    tool_output=f"ERROR: {str(e)}",
                )
                self._client.end_trajectory(TaskStatus.FAILURE)
                raise

        if asyncio.iscoroutinefunction(method):
            return async_wrapper
        return wrapper

    def _method_with_recovery(
        self, method: Callable, method_name: str, task_desc: str, args: tuple, kwargs: dict
    ) -> Any:
        """Execute a method with auto-recovery"""
        from zubbl.task_utils import normalize_task_type

        task_type = normalize_task_type(task_desc)

        recovery = self._client._interceptor.intercept_and_recover(
            func=method,
            args=args,
            kwargs=kwargs,
            task_description=task_desc,
            task_type=task_type,
        )

        self._client._last_recovery = recovery

        if recovery.success:
            self._client.record_step(
                action=method_name,
                tool_input={"args": str(args)[:200], "kwargs": str(kwargs)[:200]},
                tool_output=str(recovery.recovered_output)[:1000],
                latency_ms=recovery.latency_ms,
            )

            if recovery.fix_applied:
                self._client.record_step(
                    action="zubbl_auto_recovery",
                    tool_name="interceptor",
                    tool_input={
                        "fix_type": recovery.fix_applied.value,
                        "original_error": recovery.original_error,
                    },
                    tool_output="recovery_successful",
                    reasoning=f"Auto-recovered {method_name} from: {recovery.original_error}",
                )

            self._client.end_trajectory(TaskStatus.SUCCESS)
            return recovery.recovered_output
        else:
            self._client.record_step(
                action=method_name,
                tool_output=f"ERROR: {recovery.original_error}",
                reasoning=f"All {recovery.attempts} recovery attempts failed",
            )
            self._client.end_trajectory(TaskStatus.FAILURE)
            raise RuntimeError(
                f"{method_name} failed and auto-recovery exhausted "
                f"({recovery.attempts} attempts): {recovery.original_error}"
            )

    def _method_without_recovery(
        self, method: Callable, method_name: str, args: tuple, kwargs: dict
    ) -> Any:
        """Execute a method without auto-recovery (legacy behavior)"""
        try:
            start_time = datetime.utcnow()
            result = method(*args, **kwargs)
            latency = int((datetime.utcnow() - start_time).total_seconds() * 1000)

            self._client.record_step(
                action=method_name,
                tool_input={"args": str(args)[:200], "kwargs": str(kwargs)[:200]},
                tool_output=str(result)[:1000],
                latency_ms=latency,
            )
            self._client.end_trajectory(TaskStatus.SUCCESS)
            return result

        except Exception as e:
            self._client.record_step(
                action=method_name,
                tool_input={"args": str(args)[:200], "kwargs": str(kwargs)[:200]},
                tool_output=f"ERROR: {str(e)}",
            )
            self._client.end_trajectory(TaskStatus.FAILURE)
            raise

    @property
    def unwrapped(self) -> Any:
        """Get the original unwrapped agent"""
        return self._agent


class ZubblWrappedCrew:
    """
    Wrap an entire multi-agent crew (CrewAI, AutoGen, or custom).

    Each agent in the crew gets its own trajectory tracking.
    The crew-level execution is recorded as a parent trajectory
    with child trajectories for each agent.

    Example:
        from crewai import Crew
        zubbl = ZubblClient(api_key="your-key")
        smart_crew = zubbl.wrap_crew(crew, agent_names=["researcher", "writer"])
        result = smart_crew.kickoff()
    """

    def __init__(self, client: "ZubblClient", crew: Any, agent_names: Optional[List[str]] = None):
        self._client = client
        self._crew = crew
        self._agent_names = agent_names or []

    def __getattr__(self, name: str) -> Any:
        attr = getattr(self._crew, name)
        if callable(attr) and name in ("kickoff", "run", "invoke", "execute", "__call__"):
            return self._wrap_crew_method(attr, name)
        return attr

    def __call__(self, *args, **kwargs):
        return self._wrap_crew_method(self._crew, "crew_run")(*args, **kwargs)

    def _wrap_crew_method(self, method: Any, method_name: str):
        def wrapper(*args, **kwargs):
            task_desc = kwargs.get("task", kwargs.get("inputs", {}).get("task", str(args[0]) if args else "crew_execution"))
            if not isinstance(task_desc, str):
                task_desc = str(task_desc)[:200]

            # Start crew-level trajectory
            self._client.start_trajectory(
                task_description=f"[Crew] {task_desc}",
                metadata={
                    "type": "crew",
                    "agents": self._agent_names,
                    "method": method_name,
                },
            )

            self._client.record_step(
                action="crew_start",
                tool_name="crew",
                tool_input={"agents": self._agent_names, "task": task_desc[:200]},
                tool_output=f"Starting crew with {len(self._agent_names)} agents",
            )

            try:
                import time
                start = time.time()
                result = method(*args, **kwargs)
                latency = int((time.time() - start) * 1000)

                # Record crew completion
                result_str = str(result)[:1000] if result else ""
                self._client.record_step(
                    action="crew_complete",
                    tool_name="crew",
                    tool_output=result_str,
                    latency_ms=latency,
                )

                self._client.end_trajectory(TaskStatus.SUCCESS)
                return result

            except Exception as e:
                self._client.record_step(
                    action="crew_failed",
                    tool_name="crew",
                    tool_output=f"ERROR: {str(e)[:500]}",
                )
                self._client.end_trajectory(TaskStatus.FAILURE)
                raise

        return wrapper
