"""
CrewAI patches for Crew, Agent, and Task execution tracing.

Wraps Crew.kickoff() in an AGENT span (orchestrator role), agent execution
in AGENT spans (worker role), and task execution in INTERNAL spans.

Does NOT suppress provider instrumentation — CrewAI doesn't create its own
LLM spans, so we let the provider patches (OpenAI, Anthropic, etc.) handle
LLM calls naturally as children of the CrewAI agent span.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Dict, TypeVar

import wrapt

from risicare.integrations._base import (
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
    create_framework_attributes,
    get_framework_version,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


def _extract_crew_name(instance: Any) -> str:
    """Extract crew name from Crew instance."""
    name = getattr(instance, "name", None)
    if name:
        return name
    # Fallback to class name or default
    return getattr(type(instance), "__name__", "Crew")


def _extract_agent_info(instance: Any) -> Dict[str, Any]:
    """Extract agent info from CrewAI Agent instance."""
    return {
        "name": getattr(instance, "role", None) or getattr(instance, "name", "agent"),
        "role": getattr(instance, "role", ""),
        "goal": getattr(instance, "goal", ""),
        "backstory_length": len(getattr(instance, "backstory", "") or ""),
        "allow_delegation": getattr(instance, "allow_delegation", False),
    }


def _extract_task_info(task: Any) -> Dict[str, Any]:
    """Extract task info from CrewAI Task instance."""
    description = getattr(task, "description", "") or ""
    return {
        "description": description[:100] if description else "",
        "expected_output": getattr(task, "expected_output", "")[:100] if getattr(task, "expected_output", "") else "",
        "agent_role": getattr(getattr(task, "agent", None), "role", ""),
    }


def _wrap_kickoff(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Crew.kickoff (sync)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    crew_name = _extract_crew_name(instance)
    version = get_framework_version("crewai")
    process_type = getattr(instance, "process", None)
    process_str = str(process_type.value) if hasattr(process_type, "value") else str(process_type or "sequential")

    attrs = create_framework_attributes(
        "crewai", version,
        crew_name=crew_name,
        process_type=process_str,
    )
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "crewai"

    # Count agents and tasks
    agents = getattr(instance, "agents", [])
    tasks = getattr(instance, "tasks", [])
    attrs["framework.crewai.agent_count"] = len(agents)
    attrs["framework.crewai.task_count"] = len(tasks)

    if agents:
        attrs["framework.crewai.agent_roles"] = [
            getattr(a, "role", "unknown") for a in agents[:10]
        ]

    with tracer.start_span(
        name=f"crewai.crew/{crew_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Capture result summary
            if should_trace_content() and result is not None:
                result_str = str(result)
                safe_set_attribute(
                    span,
                    "framework.crewai.result",
                    truncate_content(result_str, 5000),
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


async def _wrap_kickoff_async(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Crew.kickoff_async."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    crew_name = _extract_crew_name(instance)
    version = get_framework_version("crewai")
    process_type = getattr(instance, "process", None)
    process_str = str(process_type.value) if hasattr(process_type, "value") else str(process_type or "sequential")

    attrs = create_framework_attributes(
        "crewai", version,
        crew_name=crew_name,
        process_type=process_str,
    )
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "crewai"

    agents = getattr(instance, "agents", [])
    tasks = getattr(instance, "tasks", [])
    attrs["framework.crewai.agent_count"] = len(agents)
    attrs["framework.crewai.task_count"] = len(tasks)

    with tracer.start_span(
        name=f"crewai.crew/{crew_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and result is not None:
                result_str = str(result)
                safe_set_attribute(
                    span,
                    "framework.crewai.result",
                    truncate_content(result_str, 5000),
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


def _wrap_agent_execute_task(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for CrewAI Agent.execute_task."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind, SemanticPhase

    agent_info = _extract_agent_info(instance)
    version = get_framework_version("crewai")

    attrs = create_framework_attributes(
        "crewai", version,
        agent_role=agent_info["role"],
        allow_delegation=agent_info["allow_delegation"],
    )
    attrs["agent.role"] = "worker"
    attrs["agent.type"] = "crewai"
    attrs["agent.name"] = agent_info["name"]

    if should_trace_content() and agent_info["goal"]:
        attrs["framework.crewai.agent_goal"] = truncate_content(
            agent_info["goal"], 2000
        )

    # Extract task info from first positional arg
    task = args[0] if args else kwargs.get("task")
    if task:
        task_info = _extract_task_info(task)
        if task_info["description"]:
            attrs["framework.crewai.task_description"] = task_info["description"]

    with tracer.start_span(
        name=f"crewai.agent/{agent_info['name']}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        span.semantic_phase = SemanticPhase.ACT
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and result is not None:
                result_str = str(result)
                safe_set_attribute(
                    span,
                    "framework.crewai.agent_output",
                    truncate_content(result_str, 5000),
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


def patch_crewai(module: Any) -> None:
    """Apply wrapt patches to CrewAI classes."""
    # Patch Crew.kickoff
    try:
        wrapt.wrap_function_wrapper(module, "Crew.kickoff", _wrap_kickoff)
    except Exception as e:
        logger.debug(f"Failed to patch Crew.kickoff: {e}")

    # Patch Crew.kickoff_async
    try:
        wrapt.wrap_function_wrapper(module, "Crew.kickoff_async", _wrap_kickoff_async)
    except Exception as e:
        logger.debug(f"Failed to patch Crew.kickoff_async: {e}")

    # Patch agent execution
    try:
        wrapt.wrap_function_wrapper(
            "crewai.agent", "CrewAgent.execute_task", _wrap_agent_execute_task
        )
    except Exception as e:
        logger.debug(f"Failed to patch CrewAgent.execute_task: {e}")
        # Try alternative class path
        try:
            wrapt.wrap_function_wrapper(
                module, "Agent.execute_task", _wrap_agent_execute_task
            )
        except Exception:
            logger.debug("Failed to patch Agent.execute_task (fallback)")
