from __future__ import annotations

from typing import Any


def wrap_crewai(crew: Any) -> Any:
    """Wrap a CrewAI Crew instance to automatically trace kickoff() calls.

    Usage:
        from tokencat.wrappers.crewai import wrap_crewai

        crew = Crew(agents=[...], tasks=[...])
        crew = wrap_crewai(crew)
        result = crew.kickoff()  # automatically traced
    """
    try:
        original_kickoff = crew.kickoff
    except AttributeError:
        print("[ZappyBee] Could not wrap - no kickoff method found.")
        return crew

    original_kickoff_async = getattr(crew, "kickoff_async", None)

    def traced_kickoff(*args: Any, **kwargs: Any) -> Any:
        import asyncio
        from ..client import TokenCat

        async def _run() -> Any:
            crew_name = getattr(crew, "name", None) or "CrewAI"
            agents = getattr(crew, "agents", [])
            tasks = getattr(crew, "tasks", [])

            trace = await TokenCat.start_trace(
                f"crewai/{crew_name}",
                input={
                    "agents": [getattr(a, "role", str(a)) for a in agents[:10]],
                    "tasks": [getattr(t, "description", str(t))[:200] for t in tasks[:10]],
                    "inputs": kwargs.get("inputs", {}),
                },
            )

            try:
                result = original_kickoff(*args, **kwargs)

                output = None
                if hasattr(result, "raw"):
                    output = result.raw[:3000] if isinstance(result.raw, str) else str(result.raw)[:3000]
                elif isinstance(result, str):
                    output = result[:3000]

                await trace.end(output=output, status="success")
                return result

            except Exception as e:
                await trace.end(status="error", error_message=str(e))
                raise

        # If there's already a running event loop, schedule it
        try:
            loop = asyncio.get_running_loop()
            # We're inside an async context — can't use asyncio.run
            # Fall back to sync kickoff without tracing
            return original_kickoff(*args, **kwargs)
        except RuntimeError:
            return asyncio.run(_run())

    async def traced_kickoff_async(*args: Any, **kwargs: Any) -> Any:
        from ..client import TokenCat

        crew_name = getattr(crew, "name", None) or "CrewAI"
        agents = getattr(crew, "agents", [])
        tasks = getattr(crew, "tasks", [])

        trace = await TokenCat.start_trace(
            f"crewai/{crew_name}",
            input={
                "agents": [getattr(a, "role", str(a)) for a in agents[:10]],
                "tasks": [getattr(t, "description", str(t))[:200] for t in tasks[:10]],
                "inputs": kwargs.get("inputs", {}),
            },
        )

        try:
            result = await original_kickoff_async(*args, **kwargs)

            output = None
            if hasattr(result, "raw"):
                output = result.raw[:3000] if isinstance(result.raw, str) else str(result.raw)[:3000]
            elif isinstance(result, str):
                output = result[:3000]

            await trace.end(output=output, status="success")
            return result

        except Exception as e:
            await trace.end(status="error", error_message=str(e))
            raise

    crew.kickoff = traced_kickoff
    if original_kickoff_async:
        crew.kickoff_async = traced_kickoff_async

    return crew
