"""LangChain integration for Sentrial.

Automatically captures agent traces, tool calls, and token usage.
This data powers signal detection, root cause analysis, and AI-suggested fixes.

Supports ALL LangChain versions:
- langchain 1.0.0+ (LangGraph era - AgentExecutor still works!)
- langchain-core >= 0.1.0 (recommended, uses langchain_core imports)
- langchain 0.0.x - 0.2.x (legacy, uses langchain imports)

The callback handler automatically detects which version is installed
and uses the appropriate imports.

NOTE: LangChain 1.0+ deprecated AgentExecutor in favor of LangGraph,
but AgentExecutor STILL WORKS. Our callback handler works with both.
"""

from typing import Any, Dict, List, Optional, TYPE_CHECKING, Callable
from uuid import UUID
import logging
import time

_logger = logging.getLogger("sentrial")

# Version-compatible imports - support ALL LangChain versions
LANGCHAIN_VERSION = "unknown"
LANGCHAIN_MAJOR_VERSION = 0
BaseCallbackHandler = None
AgentAction = None
AgentFinish = None
LLMResult = None

# Detect actual langchain version
try:
    import langchain
    _version_str = getattr(langchain, "__version__", "0.0.0")
    LANGCHAIN_MAJOR_VERSION = int(_version_str.split(".")[0])
except:
    pass

# Try langchain-core first (works for 0.1.0+ and 1.0.0+)
try:
    from langchain_core.callbacks.base import BaseCallbackHandler
    from langchain_core.agents import AgentAction, AgentFinish
    from langchain_core.outputs import LLMResult
    LANGCHAIN_VERSION = "core-1.x" if LANGCHAIN_MAJOR_VERSION >= 1 else "core"
except ImportError:
    pass

# Fall back to langchain imports (0.0.x - 0.2.x)
if BaseCallbackHandler is None:
    try:
        from langchain.callbacks.base import BaseCallbackHandler
        from langchain.agents import AgentAction, AgentFinish
        from langchain.schema import LLMResult
        LANGCHAIN_VERSION = "legacy"
    except ImportError:
        pass

# Final fallback - try langchain.schema for even older versions
if BaseCallbackHandler is None:
    try:
        from langchain.callbacks.base import BaseCallbackHandler
        from langchain.schema.agent import AgentAction, AgentFinish
        from langchain.schema.output import LLMResult
        LANGCHAIN_VERSION = "legacy-schema"
    except ImportError:
        raise ImportError(
            "LangChain is required for this integration. "
            "Install it with: pip install langchain-core>=0.1.0 (recommended) "
            "or pip install langchain>=0.0.200 (legacy)"
        )


def get_agent_executor():
    """
    Get AgentExecutor class that works with your LangChain version.
    
    NOTE: LangChain 1.0+ REMOVED AgentExecutor! Use get_react_agent() instead,
    which uses LangGraph under the hood.
    
    Returns:
        AgentExecutor class (for LangChain < 1.0)
        
    Raises:
        ImportError: If AgentExecutor not found (LangChain 1.0+)
        
    Example (LangChain < 1.0):
        AgentExecutor = get_agent_executor()
        executor = AgentExecutor(agent=agent, tools=tools)
    """
    # Try multiple import paths (LangChain moves things around between versions)
    try:
        from langchain.agents import AgentExecutor
        return AgentExecutor
    except ImportError:
        pass
    
    try:
        from langchain_core.agents import AgentExecutor
        return AgentExecutor
    except ImportError:
        pass
    
    try:
        from langchain.agents.agent import AgentExecutor
        return AgentExecutor
    except ImportError:
        pass
    
    raise ImportError(
        "AgentExecutor not found. In LangChain 1.0+, AgentExecutor was REMOVED.\n"
        "Use get_react_agent() instead, or downgrade: pip install 'langchain<1.0'"
    )


def get_react_agent():
    """
    Get the LangGraph-based react agent for LangChain 1.0+.
    
    This is the NEW way to create agents in LangChain 1.x.
    
    Returns:
        create_react_agent function from langgraph
        
    Example:
        create_react_agent = get_react_agent()
        agent = create_react_agent(model, tools)  # Note: different API than AgentExecutor!
        
        # Use with Sentrial:
        handler = SentrialCallbackHandler(client, session_id)
        result = agent.invoke({"messages": [("user", query)]}, config={"callbacks": [handler]})
    """
    try:
        from langgraph.prebuilt import create_react_agent
        return create_react_agent
    except ImportError:
        raise ImportError(
            "LangGraph not found. For LangChain 1.0+, install it:\n"
            "pip install langgraph"
        )


def get_openai_tools_agent_creator():
    """
    Get create_openai_tools_agent function that works with your LangChain version.
    
    NOTE: This was REMOVED in LangChain 1.0+. Use get_react_agent() instead.
    
    Returns:
        create_openai_tools_agent function (for LangChain < 1.0)
        
    Example (LangChain < 1.0):
        create_agent = get_openai_tools_agent_creator()
        agent = create_agent(llm, tools, prompt)
    """
    try:
        from langchain.agents import create_openai_tools_agent
        return create_openai_tools_agent
    except ImportError:
        pass
    
    try:
        from langchain_openai.agents import create_openai_tools_agent
        return create_openai_tools_agent
    except ImportError:
        pass
    
    raise ImportError(
        "create_openai_tools_agent not found. In LangChain 1.0+, this was REMOVED.\n"
        "Use get_react_agent() instead, or downgrade: pip install 'langchain<1.0'"
    )


def create_agent_with_sentrial(
    llm,
    tools: list,
    client: "SentrialClient",
    agent_name: str,
    user_id: str = "anonymous",
    verbose: bool = False,
):
    """
    Create an agent with Sentrial tracking that works with ANY LangChain version.
    
    This is the RECOMMENDED way to use Sentrial with LangChain - it handles
    all version differences automatically.
    
    Args:
        llm: The LLM to use (ChatOpenAI, ChatAnthropic, etc.)
        tools: List of tools for the agent
        client: SentrialClient instance
        agent_name: Name for this agent (appears in dashboard)
        user_id: User identifier for session grouping
        verbose: Print debug info
        
    Returns:
        A callable that takes a query string and returns the response.
        
    Example:
        from sentrial import SentrialClient, create_agent_with_sentrial
        from langchain_openai import ChatOpenAI
        
        client = SentrialClient()
        llm = ChatOpenAI(model="gpt-4o")
        
        agent = create_agent_with_sentrial(
            llm=llm,
            tools=[my_tool1, my_tool2],
            client=client,
            agent_name="my-agent",
        )
        
        result = agent("What's the weather?")
        print(result)
    """
    # Detect LangChain version and use appropriate method
    if LANGCHAIN_MAJOR_VERSION >= 1:
        # LangChain 1.x - use LangGraph
        try:
            from langgraph.prebuilt import create_react_agent
        except ImportError:
            raise ImportError(
                "LangChain 1.0+ requires LangGraph for agents.\n"
                "Install it: pip install langgraph"
            )
        
        react_agent = create_react_agent(llm, tools)
        
        def run_agent(query: str) -> str:
            session_id = client.create_session(
                name=f"{agent_name}: {query[:50]}",
                agent_name=agent_name,
                user_id=user_id,
            )
            handler = SentrialCallbackHandler(
                client=client,
                session_id=session_id,
                verbose=verbose,
            )
            handler.set_input(query)
            
            try:
                result = react_agent.invoke(
                    {"messages": [("user", query)]},
                    config={"callbacks": [handler]},
                )
                output = result["messages"][-1].content if result.get("messages") else str(result)
                handler.finish(success=True)
                return output
            except Exception as e:
                handler.finish(success=False, failure_reason=str(e))
                raise
        
        return run_agent
    
    else:
        # LangChain < 1.0 - use AgentExecutor
        from langchain.agents import AgentExecutor, create_openai_tools_agent
        from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
        
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant. Use tools when needed."),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])
        
        agent = create_openai_tools_agent(llm, tools, prompt)
        executor = AgentExecutor(agent=agent, tools=tools, verbose=verbose)
        
        def run_agent(query: str) -> str:
            session_id = client.create_session(
                name=f"{agent_name}: {query[:50]}",
                agent_name=agent_name,
                user_id=user_id,
            )
            handler = SentrialCallbackHandler(
                client=client,
                session_id=session_id,
                verbose=verbose,
            )
            handler.set_input(query)
            
            try:
                result = executor.invoke(
                    {"input": query},
                    config={"callbacks": [handler]},
                )
                output = result.get("output", str(result))
                handler.finish(success=True)
                return output
            except Exception as e:
                handler.finish(success=False, failure_reason=str(e))
                raise
        
        return run_agent

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


class SentrialCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler for Sentrial performance monitoring.
    
    Automatically tracks:
    - Agent reasoning (Chain of Thought)
    - Tool executions (with inputs/outputs)
    - Tool errors
    - LLM calls with REAL token usage and cost
    
    Usage:
        from sentrial import SentrialClient
        from sentrial.langchain import SentrialCallbackHandler
        
        client = SentrialClient(api_url="...", project_id="...")
        session_id = client.create_session(name="My Agent Run", agent_name="my_agent")
        handler = SentrialCallbackHandler(client, session_id)
        
        # Pass to LangChain agent
        agent_executor = AgentExecutor(
            agent=agent,
            tools=tools,
            callbacks=[handler],
            verbose=True
        )
        
        # After agent finishes, get actual costs
        print(f"Total cost: ${handler.total_cost:.4f}")
        print(f"Tokens used: {handler.total_tokens}")
    """
    
    def __init__(
        self, 
        client: SentrialClient, 
        session_id: str,
        track_llm_calls: bool = True,  # Default to True now
        verbose: bool = False
    ):
        """
        Initialize Sentrial callback handler.
        
        Args:
            client: SentrialClient instance
            session_id: Active session ID
            track_llm_calls: Whether to track individual LLM API calls (default: True)
            verbose: Print tracking info (default: False)
        """
        super().__init__()
        self.client = client
        self.session_id = session_id
        self.track_llm_calls = track_llm_calls
        self.verbose = verbose
        
        # Track tool runs in flight (run_id -> tool data)
        self.tool_runs: Dict[str, Dict[str, Any]] = {}
        
        # Track pending agent actions (for correlating with tool outputs)
        self.pending_actions: Dict[str, Dict[str, Any]] = {}
        
        # Track agent steps
        self.step_count = 0
        
        # Track last LLM output for reasoning
        self.last_llm_reasoning: Optional[str] = None
        
        # Token and cost tracking (REAL data from LLM responses)
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        self.total_tokens = 0
        self.total_cost = 0.0
        self.llm_calls = 0
        self.model_name: Optional[str] = None
        
        # Automatic duration tracking
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        
        # User input/output tracking (for Raindrop-style display)
        self.user_input: Optional[str] = None
        self.assistant_output: Optional[str] = None
        self._last_llm_output: Optional[str] = None
    
    def set_input(self, user_input: str) -> None:
        """
        Set the user's original input/query for this session.
        
        Call this before running the agent to track the user's query.
        
        Args:
            user_input: The user's original question/query
            
        Example:
            handler = SentrialCallbackHandler(client, session_id)
            handler.set_input("What's the weather in San Francisco?")
            agent_executor.invoke({"input": user_input}, callbacks=[handler])
        """
        self.user_input = user_input
        self._log(f"User input set: {user_input[:50]}...")
    
    def _log(self, message: str):
        """Log if verbose mode is enabled."""
        if self.verbose:
            print(f"[Sentrial] {message}")
    
    @property
    def duration_ms(self) -> int:
        """Get duration in milliseconds (automatically tracked)."""
        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 get_usage_summary(self) -> Dict[str, Any]:
        """Get summary of token usage, costs, and duration for this session."""
        return {
            "llm_calls": self.llm_calls,
            "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,
        }
    
    # ===== Agent Reasoning (Chain of Thought) =====
    
    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """
        Capture agent reasoning. Does NOT track the tool call yet — that
        happens in on_tool_end once the real output is available.
        """
        # Start timing on first action if not already started
        if self.start_time is None:
            self.start_time = time.time()

        self.step_count += 1
        reasoning = action.log if action.log else f"Action: {action.tool}"
        tool_name = action.tool
        tool_input = action.tool_input

        self._log(f"Step {self.step_count}: Agent action - {tool_name}")

        # Store the pending action for correlation with tool output in on_tool_end
        self.pending_actions[str(run_id)] = {
            "tool": tool_name,
            "input": tool_input,
            "reasoning": reasoning,
            "step_number": self.step_count,
        }
    
    def on_agent_finish(
        self,
        finish: AgentFinish,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """Track agent completion and capture assistant output."""
        try:
            self._log("Agent finished")

            # Capture assistant output for session summary
            output = finish.return_values.get('output', 'No output')
            self.assistant_output = output

            if not self.session_id:
                return

            self.client.track_decision(
                session_id=self.session_id,
                reasoning=f"Agent completed: {output}",
                confidence=1.0,
            )
        except Exception:
            _logger.debug("Sentrial: Error in on_agent_finish callback", exc_info=True)
    
    # ===== Tool Execution =====
    
    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """Buffer tool inputs for correlation with outputs."""
        tool_name = serialized.get("name", "unknown_tool")
        
        self._log(f"Tool started: {tool_name}")
        
        # Check if we have a pending action for this tool
        parent_id = str(parent_run_id) if parent_run_id else None
        pending = None
        if parent_id and parent_id in self.pending_actions:
            pending = self.pending_actions.get(parent_id)
        
        self.tool_runs[str(run_id)] = {
            "name": tool_name,
            "input": input_str,
            "pending_action": pending,
            "parent_run_id": parent_id
        }
    
    def on_tool_end(
        self,
        output: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """Track successful tool execution with full input/output."""
        try:
            run_data = self.tool_runs.pop(str(run_id), None)

            # Safely serialize output (might be ToolMessage or other object)
            if isinstance(output, str):
                output_value = output
            elif hasattr(output, 'content'):
                output_value = str(output.content)
            else:
                output_value = str(output)

            # Try to correlate with pending agent action for reasoning
            parent_id = str(parent_run_id) if parent_run_id else None
            pending = self.pending_actions.pop(parent_id, None) if parent_id else None

            if pending:
                tool_name = pending["tool"]
                tool_input = pending["input"]
                reasoning = pending.get("reasoning")
                input_dict = tool_input if isinstance(tool_input, dict) else {"input": str(tool_input) if tool_input else ""}
                self._log(f"Tool completed: {tool_name}")
            elif run_data:
                tool_name = run_data["name"]
                tool_input = run_data["input"]
                reasoning = self.last_llm_reasoning
                self.last_llm_reasoning = None
                input_dict = {"input": tool_input} if isinstance(tool_input, str) else tool_input
                self._log(f"Tool completed (no agent_action): {tool_name}")
            else:
                return

            if not self.session_id:
                return

            self.client.track_tool_call(
                session_id=self.session_id,
                tool_name=tool_name,
                tool_input=input_dict,
                tool_output={"output": output_value},
                reasoning=reasoning,
            )
        except Exception:
            _logger.debug("Sentrial: Error in on_tool_end callback", exc_info=True)
    
    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """Track tool errors."""
        try:
            run_data = self.tool_runs.pop(str(run_id), None)

            if not run_data:
                return

            tool_name = run_data["name"]
            tool_input = run_data["input"]

            self._log(f"Tool error: {tool_name} - {error}")

            if not self.session_id:
                return

            input_dict = {"input": tool_input} if isinstance(tool_input, str) else tool_input

            self.client.track_tool_call(
                session_id=self.session_id,
                tool_name=tool_name,
                tool_input=input_dict,
                tool_output={"error": str(error), "error_type": type(error).__name__},
            )
        except Exception:
            _logger.debug("Sentrial: Error in on_tool_error callback", exc_info=True)
    
    # ===== LLM Calls - Now tracks REAL token usage and costs =====
    
    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """Track LLM call start and auto-track duration."""
        # Start timing on first LLM call
        if self.start_time is None:
            self.start_time = time.time()

        # Check for experiment system prompt override (set by CLI)
        import os
        exp_prompt = os.environ.get("SENTRIAL_EXPERIMENT_SYSTEM_PROMPT")
        if exp_prompt and prompts:
            # Prepend experiment system prompt to first prompt
            # This modifies the prompts in-place for the LLM
            prompts[0] = f"{exp_prompt}\n\n{prompts[0]}"
            self._log(f"Injected experiment system prompt")

        self._log(f"LLM call started: {len(prompts)} prompts")
    
    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """
        Track LLM responses with REAL token usage and costs.

        Extracts actual usage data from the LLM response.
        """
        try:
            self.llm_calls += 1

            # Extract token usage - check multiple locations since different providers structure it differently
            prompt_tokens = 0
            completion_tokens = 0

            # 1. Check llm_output (OpenAI style)
            llm_output = response.llm_output or {}
            token_usage = llm_output.get("token_usage", {}) or llm_output.get("usage", {})
            if token_usage:
                prompt_tokens = token_usage.get("prompt_tokens", 0) or token_usage.get("input_tokens", 0) or 0
                completion_tokens = token_usage.get("completion_tokens", 0) or token_usage.get("output_tokens", 0) or 0

            # 2. Check generation.message.usage_metadata (Gemini, newer LangChain)
            if not (prompt_tokens or completion_tokens) and response.generations:
                for gen_list in response.generations:
                    for gen in gen_list:
                        if hasattr(gen, 'message') and hasattr(gen.message, 'usage_metadata'):
                            usage_meta = gen.message.usage_metadata
                            if usage_meta:
                                prompt_tokens = usage_meta.get('input_tokens', 0) or usage_meta.get('prompt_tokens', 0) or 0
                                completion_tokens = usage_meta.get('output_tokens', 0) or usage_meta.get('completion_tokens', 0) or 0
                                break
                    if prompt_tokens or completion_tokens:
                        break

            # Get model name from llm_output or generation_info
            model_name = llm_output.get("model_name") or llm_output.get("model") or ""
            if not model_name and response.generations:
                for gen_list in response.generations:
                    for gen in gen_list:
                        if hasattr(gen, 'generation_info') and gen.generation_info:
                            model_name = gen.generation_info.get('model_name', '')
                            break
            if model_name:
                self.model_name = model_name

            # Calculate cost for this call
            call_cost = calculate_cost(model_name, prompt_tokens, completion_tokens)

            # Accumulate totals
            self.total_prompt_tokens += prompt_tokens
            self.total_completion_tokens += completion_tokens
            self.total_tokens += prompt_tokens + completion_tokens
            self.total_cost += call_cost

            # Update end time on every LLM call (last one will be the final end time)
            self.end_time = time.time()

            self._log(
                f"LLM call completed: {prompt_tokens} in, {completion_tokens} out, "
                f"${call_cost:.4f} (total: ${self.total_cost:.4f})"
            )

            # Extract reasoning and output from the LLM response
            if response.generations and len(response.generations) > 0:
                generation = response.generations[0][0]

                # Capture text for assistant_output fallback (LangGraph never fires on_agent_finish)
                gen_text = None
                if hasattr(generation, 'message') and hasattr(generation.message, 'content'):
                    gen_text = generation.message.content
                elif hasattr(generation, 'text'):
                    gen_text = generation.text

                if gen_text:
                    self._last_llm_output = gen_text

                # Try to extract "Thought:" reasoning for the next tool call
                text = gen_text or ""
                if 'Thought:' in text:
                    thought_start = text.find('Thought:') + len('Thought:')
                    thought_end = text.find('Action:', thought_start)
                    if thought_end == -1:
                        thought_end = text.find('Final Answer:', thought_start)
                    if thought_end == -1:
                        thought_end = len(text)
                    reasoning = text[thought_start:thought_end].strip()
                    self.last_llm_reasoning = reasoning
        except Exception:
            _logger.debug("Sentrial: Error in on_llm_end callback", exc_info=True)
    
    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any
    ) -> None:
        """Track LLM errors."""
        self._log(f"LLM error: {error}")

    # ===== Session Completion =====

    def finish(
        self,
        success: bool = True,
        failure_reason: Optional[str] = None,
        custom_metrics: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Finish the session and record final metrics.
        
        Automatically includes tracked token usage and costs.
        
        Args:
            success: Whether the agent run succeeded
            failure_reason: Reason for failure if success=False
            custom_metrics: Additional custom metrics

        Returns:
            Updated session data

        Example:
            handler = SentrialCallbackHandler(client, session_id)
            # ... run agent with handler ...
            result = handler.finish(
                success=True,
                custom_metrics={"quality": 4.5}
            )
            print(f"Total cost: ${handler.total_cost:.4f}")
        """
        # Set end time if not already set
        if self.end_time is None:
            self.end_time = time.time()

        if not self.session_id:
            return {}

        # Use last LLM output as fallback for assistant_output
        # (on_agent_finish never fires in LangGraph)
        final_output = self.assistant_output or self._last_llm_output

        return self.client.complete_session(
            session_id=self.session_id,
            success=success,
            failure_reason=failure_reason,
            estimated_cost=self.total_cost,
            custom_metrics=custom_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=final_output,
        )
