"""CrewAI integration for Sentrial.

Automatically captures agent tasks, tool calls, and crew execution from CrewAI.

Supports CrewAI >= 0.28.0 (modern callback API).

Usage:

    from crewai import Agent, Task, Crew
    from sentrial import SentrialClient
    from sentrial.crewai import SentrialCrewHandler
    
    client = SentrialClient()
    session_id = client.create_session(name="Research Crew", agent_name="crewai-research")
    handler = SentrialCrewHandler(client, session_id)
    
    # Create your crew
    researcher = Agent(role="Researcher", goal="Find information", ...)
    writer = Agent(role="Writer", goal="Write content", ...)
    
    research_task = Task(description="Research topic X", agent=researcher)
    write_task = Task(description="Write about X", agent=writer)
    
    crew = Crew(
        agents=[researcher, writer],
        tasks=[research_task, write_task],
        step_callback=handler.on_step,
        task_callback=handler.on_task_complete,
    )
    
    result = crew.kickoff(inputs={"topic": "AI agents"})
    
    handler.finish(success=True, output=str(result))


Alternative - wrap entire crew execution:

    from sentrial.crewai import track_crew
    
    @track_crew(client, agent_name="research-crew")
    def run_research(topic: str):
        crew = Crew(agents=[...], tasks=[...])
        return crew.kickoff(inputs={"topic": topic})
    
    result = run_research("AI agents")
"""

from typing import Any, Dict, List, Optional, Callable, Union
from functools import wraps
import time
import json

from .client import SentrialClient
from .pricing import MODEL_PRICING, calculate_cost


# Version detection
CREWAI_VERSION = "none"

try:
    from crewai import Crew, Agent, Task
    from crewai.tasks.task_output import TaskOutput
    CREWAI_VERSION = "modern"
except ImportError:
    pass

# Try older import path
if CREWAI_VERSION == "none":
    try:
        from crewai.crew import Crew
        from crewai.agent import Agent
        from crewai.task import Task
        CREWAI_VERSION = "legacy"
    except ImportError:
        pass


class SentrialCrewHandler:
    """
    CrewAI handler for Sentrial performance monitoring.
    
    Automatically tracks:
    - Task executions and outputs
    - Agent steps and reasoning
    - Tool calls
    - Token usage and costs (when available)
    
    Use with CrewAI's callback system:
    - step_callback: Called on each agent step
    - task_callback: Called when a task completes
    """
    
    def __init__(
        self,
        client: SentrialClient,
        session_id: str,
        verbose: bool = False,
    ):
        """
        Initialize Sentrial CrewAI handler.
        
        Args:
            client: SentrialClient instance
            session_id: Active session ID
            verbose: Print tracking info
        """
        self.client = client
        self.session_id = session_id
        self.verbose = verbose
        
        # Tracking state
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.step_count = 0
        self.task_count = 0
        self.tool_calls = 0
        
        # Token and cost tracking
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        self.total_tokens = 0
        self.total_cost = 0.0
        self.model_name: Optional[str] = None
        
        # User input/output
        self.user_input: Optional[str] = None
        self.assistant_output: Optional[str] = None
        
        # Task tracking
        self.task_outputs: List[Dict[str, Any]] = []
    
    def _log(self, message: str):
        """Log if verbose mode is enabled."""
        if self.verbose:
            print(f"[Sentrial/CrewAI] {message}")
    
    @property
    def duration_ms(self) -> int:
        """Get duration in milliseconds."""
        if self.start_time is None:
            return 0
        end = self.end_time if self.end_time else time.time()
        return int((end - self.start_time) * 1000)
    
    def set_input(self, user_input: str) -> None:
        """Set the user's original input for this session."""
        self.user_input = user_input
        if self.start_time is None:
            self.start_time = time.time()
        self._log(f"User input set: {user_input[:50]}...")
    
    # ===== CrewAI Callbacks =====
    
    def on_step(self, step_output: Any) -> None:
        """
        Callback for each agent step.
        
        Pass this to Crew's step_callback parameter.
        
        Args:
            step_output: The output from the agent step (AgentAction or similar)
        """
        if self.start_time is None:
            self.start_time = time.time()
        
        self.step_count += 1
        
        # Extract step info based on type
        step_type = type(step_output).__name__
        self._log(f"Step {self.step_count}: {step_type}")
        
        # Handle different step output types
        if hasattr(step_output, 'tool'):
            # Tool call
            self._track_tool_step(step_output)
        elif hasattr(step_output, 'thought'):
            # Agent thought/reasoning
            self._track_thought_step(step_output)
        elif hasattr(step_output, 'text') or hasattr(step_output, 'output'):
            # Text output
            self._track_output_step(step_output)
        else:
            # Generic step
            self._track_generic_step(step_output)
    
    def _track_tool_step(self, step: Any) -> None:
        """Track a tool execution step."""
        self.tool_calls += 1
        
        tool_name = getattr(step, 'tool', 'unknown')
        tool_input = getattr(step, 'tool_input', {})
        tool_output = getattr(step, 'result', getattr(step, 'output', None))
        
        # Serialize inputs
        if isinstance(tool_input, str):
            try:
                input_dict = json.loads(tool_input)
            except json.JSONDecodeError:
                input_dict = {"input": tool_input}
        elif isinstance(tool_input, dict):
            input_dict = tool_input
        else:
            input_dict = {"input": str(tool_input)}
        
        # Serialize output
        output_dict = {"output": str(tool_output)[:500]} if tool_output else {"status": "executed"}
        
        self._log(f"Tool: {tool_name}")
        
        self.client.track_tool_call(
            session_id=self.session_id,
            tool_name=tool_name,
            tool_input=input_dict,
            tool_output=output_dict,
        )
    
    def _track_thought_step(self, step: Any) -> None:
        """Track an agent thought/reasoning step."""
        thought = getattr(step, 'thought', str(step))
        
        self.client.track_decision(
            session_id=self.session_id,
            reasoning=f"Agent thought: {str(thought)[:300]}",
            confidence=1.0
        )
    
    def _track_output_step(self, step: Any) -> None:
        """Track an output step."""
        output = getattr(step, 'text', getattr(step, 'output', str(step)))
        
        # Update assistant output to latest
        self.assistant_output = str(output)[:500]
        
        self.client.track_decision(
            session_id=self.session_id,
            reasoning=f"Agent output: {str(output)[:200]}",
            confidence=1.0
        )
    
    def _track_generic_step(self, step: Any) -> None:
        """Track a generic step."""
        # Try to extract useful info
        step_str = str(step)[:300]
        
        self.client.track_decision(
            session_id=self.session_id,
            reasoning=f"Step: {step_str}",
            confidence=1.0
        )
    
    def on_task_complete(self, task_output: Any) -> None:
        """
        Callback when a task completes.
        
        Pass this to Crew's task_callback parameter.
        
        Args:
            task_output: TaskOutput from the completed task
        """
        self.task_count += 1
        
        # Extract task info
        if hasattr(task_output, 'description'):
            description = task_output.description
        elif hasattr(task_output, 'task') and hasattr(task_output.task, 'description'):
            description = task_output.task.description
        else:
            description = f"Task {self.task_count}"
        
        # Get output
        if hasattr(task_output, 'raw'):
            raw_output = task_output.raw
        elif hasattr(task_output, 'output'):
            raw_output = task_output.output
        else:
            raw_output = str(task_output)
        
        # Get agent info
        agent_name = "unknown"
        if hasattr(task_output, 'agent'):
            agent_name = getattr(task_output.agent, 'role', str(task_output.agent))
        
        self._log(f"Task completed by {agent_name}: {description[:50]}...")
        
        # Store task output
        self.task_outputs.append({
            "task": description[:200],
            "agent": agent_name,
            "output": str(raw_output)[:500],
        })
        
        # Update assistant output to latest task output
        self.assistant_output = str(raw_output)[:500]
        
        # Track as tool call (task execution)
        self.client.track_tool_call(
            session_id=self.session_id,
            tool_name=f"task_{self.task_count}",
            tool_input={"description": description[:200], "agent": agent_name},
            tool_output={"result": str(raw_output)[:500]},
            reasoning=f"Task completed by {agent_name}",
        )
    
    # ===== Crew Wrapper =====
    
    def wrap_crew(self, crew: Any) -> Any:
        """
        Wrap a Crew instance to automatically track execution.
        
        This adds the callbacks automatically and wraps kickoff().
        
        Args:
            crew: CrewAI Crew instance
            
        Returns:
            The wrapped Crew instance
        """
        # Add callbacks if not already set
        if crew.step_callback is None:
            crew.step_callback = self.on_step
        if crew.task_callback is None:
            crew.task_callback = self.on_task_complete
        
        # Wrap kickoff method
        original_kickoff = crew.kickoff
        
        @wraps(original_kickoff)
        def tracked_kickoff(inputs: Optional[Dict[str, Any]] = None):
            self.start_time = time.time()
            
            # Capture input
            if inputs:
                self.user_input = json.dumps(inputs)[:500]
            
            try:
                result = original_kickoff(inputs)
                
                self.end_time = time.time()
                
                # Capture final output
                if hasattr(result, 'raw'):
                    self.assistant_output = str(result.raw)[:500]
                else:
                    self.assistant_output = str(result)[:500]
                
                return result
            except Exception as e:
                self.end_time = time.time()
                raise
        
        crew.kickoff = tracked_kickoff
        
        # Also wrap kickoff_async if it exists
        if hasattr(crew, 'kickoff_async'):
            original_kickoff_async = crew.kickoff_async
            
            @wraps(original_kickoff_async)
            async def tracked_kickoff_async(inputs: Optional[Dict[str, Any]] = None):
                self.start_time = time.time()
                
                if inputs:
                    self.user_input = json.dumps(inputs)[:500]
                
                try:
                    result = await original_kickoff_async(inputs)
                    
                    self.end_time = time.time()
                    
                    if hasattr(result, 'raw'):
                        self.assistant_output = str(result.raw)[:500]
                    else:
                        self.assistant_output = str(result)[:500]
                    
                    return result
                except Exception as e:
                    self.end_time = time.time()
                    raise
            
            crew.kickoff_async = tracked_kickoff_async
        
        self._log("Crew wrapped with Sentrial tracking")
        return crew
    
    # ===== Token/Cost Tracking =====
    
    def add_token_usage(
        self,
        prompt_tokens: int,
        completion_tokens: int,
        model: Optional[str] = None,
    ) -> None:
        """
        Manually add token usage (call this if you have access to LLM stats).
        
        CrewAI doesn't always expose token usage, so you may need to track
        it separately from your LLM provider.
        
        Args:
            prompt_tokens: Number of input tokens
            completion_tokens: Number of output tokens
            model: Model name for cost calculation
        """
        self.total_prompt_tokens += prompt_tokens
        self.total_completion_tokens += completion_tokens
        self.total_tokens += prompt_tokens + completion_tokens
        
        if model:
            self.model_name = model
            cost = calculate_cost(model, prompt_tokens, completion_tokens)
            self.total_cost += cost
        
        self._log(f"Added usage: {prompt_tokens} in, {completion_tokens} out")
    
    # ===== Session Completion =====
    
    def finish(
        self,
        success: bool = True,
        failure_reason: Optional[str] = None,
        output: Optional[str] = None,
        custom_metrics: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Finish the session and record final metrics.
        
        Args:
            success: Whether the crew run succeeded
            failure_reason: Reason for failure if success=False
            output: Final output (overrides auto-captured output)
            custom_metrics: Additional custom metrics
            
        Returns:
            Updated session data
        """
        if self.end_time is None:
            self.end_time = time.time()
        
        if output:
            self.assistant_output = output[:500]
        
        # Add CrewAI-specific metrics
        metrics = custom_metrics or {}
        metrics.update({
            "crewai_steps": self.step_count,
            "crewai_tasks": self.task_count,
            "crewai_tool_calls": self.tool_calls,
            "crewai_version": CREWAI_VERSION,
        })
        
        return self.client.complete_session(
            session_id=self.session_id,
            success=success,
            failure_reason=failure_reason,
            estimated_cost=self.total_cost,
            custom_metrics=metrics,
            duration_ms=self.duration_ms,
            prompt_tokens=self.total_prompt_tokens,
            completion_tokens=self.total_completion_tokens,
            total_tokens=self.total_tokens,
            user_input=self.user_input,
            assistant_output=self.assistant_output,
        )
    
    def get_usage_summary(self) -> Dict[str, Any]:
        """Get summary of usage for this session."""
        return {
            "steps": self.step_count,
            "tasks": self.task_count,
            "tool_calls": self.tool_calls,
            "task_outputs": self.task_outputs,
            "total_prompt_tokens": self.total_prompt_tokens,
            "total_completion_tokens": self.total_completion_tokens,
            "total_tokens": self.total_tokens,
            "total_cost": self.total_cost,
            "model": self.model_name,
            "duration_ms": self.duration_ms,
            "crewai_version": CREWAI_VERSION,
        }


# ===== Decorator for easy crew tracking =====

def track_crew(
    client: SentrialClient,
    agent_name: str = "crewai",
    user_id: Optional[str] = None,
):
    """
    Decorator to track a function that runs a CrewAI crew.
    
    Usage:
        @track_crew(client, agent_name="research-crew")
        def run_research(topic: str):
            crew = Crew(agents=[...], tasks=[...])
            return crew.kickoff(inputs={"topic": topic})
        
        result = run_research("AI agents")
    
    Args:
        client: SentrialClient instance
        agent_name: Name of the agent for tracking
        user_id: Optional user ID
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Create session
            session_id = client.create_session(
                name=f"{agent_name}: {func.__name__}",
                agent_name=agent_name,
                user_id=user_id,
            )
            
            handler = SentrialCrewHandler(client, session_id)
            handler.start_time = time.time()
            
            # Capture input
            input_str = json.dumps({"args": str(args), "kwargs": str(kwargs)})[:500]
            handler.user_input = input_str
            
            try:
                result = func(*args, **kwargs)
                
                # Capture output
                if hasattr(result, 'raw'):
                    handler.assistant_output = str(result.raw)[:500]
                else:
                    handler.assistant_output = str(result)[:500]
                
                handler.finish(success=True)
                return result
                
            except Exception as e:
                handler.finish(success=False, failure_reason=str(e))
                raise
        
        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            session_id = client.create_session(
                name=f"{agent_name}: {func.__name__}",
                agent_name=agent_name,
                user_id=user_id,
            )
            
            handler = SentrialCrewHandler(client, session_id)
            handler.start_time = time.time()
            
            input_str = json.dumps({"args": str(args), "kwargs": str(kwargs)})[:500]
            handler.user_input = input_str
            
            try:
                result = await func(*args, **kwargs)
                
                if hasattr(result, 'raw'):
                    handler.assistant_output = str(result.raw)[:500]
                else:
                    handler.assistant_output = str(result)[:500]
                
                handler.finish(success=True)
                return result
                
            except Exception as e:
                handler.finish(success=False, failure_reason=str(e))
                raise
        
        import asyncio
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        return wrapper
    
    return decorator
