"""
Agent Status integration for LangChain.

Provides automatic monitoring for LangChain chains, agents, and LLMs.

Usage:

    # Option 1: Expose for outside-in monitoring (RECOMMENDED)
    from agent_status.integrations.langchain import expose

    # Your LangChain agent
    chain = prompt | llm | parser

    # Expose to the distributed monitoring network
    expose(
        chain=chain,
        api_key="fb_live_xxx",
        agent_name="My LangChain Agent",
    )
    # Now distributed nodes can probe your agent from anywhere in the world!

    # Option 2: Local callback handler for metrics tracking
    from agent_status.integrations.langchain import AgentStatusCallbackHandler

    handler = AgentStatusCallbackHandler(
        api_key="fb_live_xxx",
        agent_name="My LangChain Agent",
    )

    llm = ChatOpenAI(callbacks=[handler])
    result = llm.invoke("Hello!")
"""

from __future__ import annotations

import asyncio
import json
import threading
import time
import logging
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List, Optional, Union, TYPE_CHECKING
from uuid import UUID

LANGCHAIN_AVAILABLE = False
BaseCallbackHandler = object

try:
    from langchain_core.callbacks import BaseCallbackHandler as _BaseCallbackHandler
    BaseCallbackHandler = _BaseCallbackHandler
    LANGCHAIN_AVAILABLE = True
except ImportError:
    pass

if TYPE_CHECKING:
    from langchain_core.outputs import LLMResult
    from langchain_core.agents import AgentAction, AgentFinish
    from langchain_core.messages import BaseMessage
    from langchain_core.runnables import Runnable

import agent_status
from agent_status.client import AgentStatusClient

logger = logging.getLogger(__name__)

TUNNEL_RELAY_URL = "wss://rora-tunnel.carmel.so/ws/connect"


def expose(
    chain: Any,
    api_key: str,
    agent_name: str = "LangChain Agent",
    *,
    input_key: str = "input",
    tunnel_url: str = TUNNEL_RELAY_URL,
    auto_register: bool = True,
    blocking: bool = True,
) -> Optional[str]:
    """
    Expose a LangChain chain/agent for outside-in monitoring.

    This connects your local LangChain agent to the distributed monitoring network,
    allowing nodes around the world to probe your agent without you
    needing to set up any infrastructure.

    Args:
        chain: A LangChain Runnable (chain, agent, or LLM)
        api_key: Your API key
        agent_name: Name for this agent in the dashboard
        input_key: The input key to use when invoking the chain (default: "input")
        tunnel_url: WebSocket URL for the tunnel relay
        auto_register: Auto-register agent (default: True)
        blocking: If True, blocks until interrupted. If False, runs in background thread.

    Returns:
        The public URL where the agent can be reached (if non-blocking)

    Example:
        from langchain_openai import ChatOpenAI
        from langchain_core.prompts import ChatPromptTemplate
        from agent_status.integrations.langchain import expose

        prompt = ChatPromptTemplate.from_template("You are a helpful assistant. {input}")
        llm = ChatOpenAI()
        chain = prompt | llm

        expose(chain, api_key="fb_live_xxx", agent_name="My Support Bot")

    Note: While designed for LangChain, this works with any object that has an
    `invoke(dict) -> Any` method or is a simple callable.
    """
    handler = _TunnelHandler(
        chain=chain,
        api_key=api_key,
        agent_name=agent_name,
        input_key=input_key,
        tunnel_url=tunnel_url,
        auto_register=auto_register,
    )

    if blocking:
        handler.run()
        return None
    else:
        thread = threading.Thread(target=handler.run, daemon=True)
        thread.start()
        return handler.public_url


class _TunnelHandler:
    """Internal handler for tunnel WebSocket connection."""

    def __init__(
        self,
        chain: Any,
        api_key: str,
        agent_name: str,
        input_key: str,
        tunnel_url: str,
        auto_register: bool,
    ):
        self.chain = chain
        self.api_key = api_key
        self.agent_name = agent_name
        self.input_key = input_key
        self.tunnel_url = tunnel_url
        self.auto_register = auto_register
        self.public_url: Optional[str] = None
        self.agent_id: Optional[str] = None
        self._running = False

    def run(self):
        """Run the tunnel handler (blocking)."""
        try:
            import websockets
        except ImportError:
            raise ImportError(
                "websockets is required for expose(). Install with: pip install websockets"
            )

        asyncio.run(self._run_async())

    async def _run_async(self):
        """Async main loop."""
        import websockets

        if self.auto_register:
            await self._register_agent()

        ws_url = f"{self.tunnel_url}?agent_id={self.agent_id or self.agent_name}&token={self.api_key}"

        logger.info(f"Connecting to tunnel relay: {self.tunnel_url}")

        self._running = True
        while self._running:
            try:
                async with websockets.connect(ws_url) as websocket:
                    msg = await websocket.recv()
                    data = json.loads(msg)

                    if "error" in data:
                        logger.error(f"Tunnel error: {data['error']}")
                        return

                    self.public_url = data.get("url")
                    logger.info(f"Agent exposed at: {self.public_url}")
                    print(f"\n{'='*60}")
                    print(f"  Agent Status - Agent Exposed!")
                    print(f"  Name: {self.agent_name}")
                    print(f"  URL:  {self.public_url}")
                    print(f"  Press Ctrl+C to stop")
                    print(f"{'='*60}\n")

                    await self._handle_messages(websocket)

            except websockets.ConnectionClosed:
                logger.warning("Connection closed, reconnecting in 5s...")
                await asyncio.sleep(5)
            except Exception as e:
                logger.error(f"Connection error: {e}, reconnecting in 5s...")
                await asyncio.sleep(5)

    async def _register_agent(self):
        """Register agent with the API."""
        try:
            agent_status.init(api_key=self.api_key)
            logger.debug(f"Agent Status initialized for {self.agent_name}")
        except Exception as e:
            logger.warning(f"Failed to initialize Agent Status: {e}")

    async def _handle_messages(self, websocket):
        """Handle incoming WebSocket messages (probes)."""
        import websockets

        async for message in websocket:
            try:
                data = json.loads(message)

                if data.get("type") == "probe":
                    response = await self._handle_probe(data)
                    await websocket.send(json.dumps(response))

            except Exception as e:
                logger.error(f"Error handling message: {e}")

    async def _handle_probe(self, probe: Dict[str, Any]) -> Dict[str, Any]:
        """Handle a probe request."""
        request_id = probe.get("request_id")
        body = probe.get("body", {})

        try:
            prompt = body.get("prompt") or body.get("message") or body.get(self.input_key, "")

            if not prompt:
                return {
                    "request_id": request_id,
                    "response": {
                        "status_code": 400,
                        "body": {"error": "No prompt provided"},
                    }
                }

            start_time = time.time()

            if hasattr(self.chain, "invoke"):
                result = await asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: self.chain.invoke({self.input_key: prompt})
                )
            elif callable(self.chain):
                result = await asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: self.chain(prompt)
                )
            else:
                raise ValueError("Chain must be a Runnable or callable")

            latency_ms = (time.time() - start_time) * 1000

            if hasattr(result, "content"):
                response_text = result.content
            elif isinstance(result, dict):
                response_text = result.get("output") or result.get("text") or str(result)
            else:
                response_text = str(result)

            logger.info(f"Probe handled: latency={latency_ms:.0f}ms")

            return {
                "request_id": request_id,
                "response": {
                    "status_code": 200,
                    "body": {
                        "response": response_text,
                        "latency_ms": latency_ms,
                    },
                }
            }

        except Exception as e:
            logger.error(f"Probe error: {e}")
            return {
                "request_id": request_id,
                "response": {
                    "status_code": 500,
                    "body": {"error": str(e)},
                }
            }


class AgentStatusCallbackHandler(BaseCallbackHandler if LANGCHAIN_AVAILABLE else object):
    """
    LangChain callback handler that reports metrics to Agent Status.

    Automatically tracks:
    - LLM calls (latency, tokens, model)
    - Chain executions
    - Agent actions and tool calls
    - Errors and exceptions

    Args:
        api_key: Your API key
        agent_name: Name for this agent in the dashboard
        agent_id: Existing agent ID (if already registered)
        auto_register: Auto-register agent on first call (default: True)
        base_url: API base URL

    Example:
        handler = AgentStatusCallbackHandler(
            api_key="fb_live_xxx",
            agent_name="Support Bot",
        )

        llm = ChatOpenAI(callbacks=[handler])
        chain = prompt | llm | parser
        chain.invoke({"question": "Hello"})

        print(handler.metrics)
    """

    def __init__(
        self,
        api_key: str,
        agent_name: str = "LangChain Agent",
        agent_id: Optional[str] = None,
        auto_register: bool = True,
        base_url: str = "https://api.fabric.carmel.so",
        **kwargs,
    ):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is not installed. Install with: pip install langchain-core"
            )

        super().__init__(**kwargs)

        self.api_key = api_key
        self.agent_name = agent_name
        self.agent_id = agent_id
        self.auto_register = auto_register
        self.base_url = base_url

        self._client = AgentStatusClient(api_key=api_key, base_url=base_url)

        self._call_stack: List[Dict[str, Any]] = []
        self._metrics: Dict[str, Any] = {
            "total_calls": 0,
            "total_tokens": 0,
            "total_latency_ms": 0,
            "errors": 0,
            "llm_calls": 0,
            "chain_runs": 0,
            "tool_calls": 0,
        }

        self._current_run: Optional[Dict[str, Any]] = None
        self._run_start: Optional[float] = None

    @property
    def metrics(self) -> Dict[str, Any]:
        """Get current metrics summary."""
        return {
            **self._metrics,
            "avg_latency_ms": (
                self._metrics["total_latency_ms"] / self._metrics["total_calls"]
                if self._metrics["total_calls"] > 0
                else 0
            ),
        }

    def _ensure_registered(self) -> None:
        """Ensure agent is registered."""
        if self.agent_id:
            return

        if not self.auto_register:
            return

        try:
            agent_status.init(api_key=self.api_key, base_url=self.base_url)
            logger.debug(f"Agent Status handler initialized for {self.agent_name}")
        except Exception as e:
            logger.warning(f"Failed to initialize Agent Status: {e}")

    def __enter__(self):
        """Context manager entry."""
        self._ensure_registered()
        self._run_start = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - report metrics."""
        if self._run_start:
            duration = (time.time() - self._run_start) * 1000
            self._metrics["total_latency_ms"] += duration
        return False

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when LLM starts."""
        self._call_stack.append({
            "type": "llm",
            "run_id": str(run_id),
            "start_time": time.time(),
            "model": serialized.get("kwargs", {}).get("model_name", "unknown"),
            "prompts": prompts,
        })

    def on_llm_end(
        self,
        response: "LLMResult",
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when LLM ends."""
        call = self._pop_call(str(run_id))
        if not call:
            return

        latency_ms = (time.time() - call["start_time"]) * 1000

        tokens = 0
        if response.llm_output:
            usage = response.llm_output.get("token_usage", {})
            tokens = usage.get("total_tokens", 0)

        self._metrics["total_calls"] += 1
        self._metrics["llm_calls"] += 1
        self._metrics["total_latency_ms"] += latency_ms
        self._metrics["total_tokens"] += tokens

        logger.debug(
            f"LLM call completed: {call['model']}, "
            f"latency={latency_ms:.0f}ms, tokens={tokens}"
        )

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when LLM errors."""
        self._pop_call(str(run_id))
        self._metrics["errors"] += 1
        logger.error(f"LLM error: {error}")

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List["BaseMessage"]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when chat model starts."""
        self._call_stack.append({
            "type": "chat",
            "run_id": str(run_id),
            "start_time": time.time(),
            "model": serialized.get("kwargs", {}).get("model_name", "unknown"),
            "message_count": sum(len(m) for m in messages),
        })

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when chain starts."""
        self._call_stack.append({
            "type": "chain",
            "run_id": str(run_id),
            "start_time": time.time(),
            "name": serialized.get("name", "unknown"),
        })

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when chain ends."""
        call = self._pop_call(str(run_id))
        if not call:
            return

        latency_ms = (time.time() - call["start_time"]) * 1000
        self._metrics["chain_runs"] += 1

        logger.debug(f"Chain completed: {call['name']}, latency={latency_ms:.0f}ms")

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when chain errors."""
        self._pop_call(str(run_id))
        self._metrics["errors"] += 1
        logger.error(f"Chain error: {error}")

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when tool starts."""
        self._call_stack.append({
            "type": "tool",
            "run_id": str(run_id),
            "start_time": time.time(),
            "name": serialized.get("name", "unknown"),
        })

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when tool ends."""
        call = self._pop_call(str(run_id))
        if not call:
            return

        latency_ms = (time.time() - call["start_time"]) * 1000
        self._metrics["tool_calls"] += 1

        logger.debug(f"Tool completed: {call['name']}, latency={latency_ms:.0f}ms")

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when tool errors."""
        self._pop_call(str(run_id))
        self._metrics["errors"] += 1
        logger.error(f"Tool error: {error}")

    def on_agent_action(
        self,
        action: "AgentAction",
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when agent takes action."""
        logger.debug(f"Agent action: {action.tool} - {action.tool_input}")

    def on_agent_finish(
        self,
        finish: "AgentFinish",
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when agent finishes."""
        logger.debug(f"Agent finished: {finish.return_values}")

    def _pop_call(self, run_id: str) -> Optional[Dict[str, Any]]:
        """Pop a call from the stack by run_id."""
        for i, call in enumerate(self._call_stack):
            if call["run_id"] == run_id:
                return self._call_stack.pop(i)
        return None

    def reset_metrics(self) -> None:
        """Reset all tracked metrics."""
        self._metrics = {
            "total_calls": 0,
            "total_tokens": 0,
            "total_latency_ms": 0,
            "errors": 0,
            "llm_calls": 0,
            "chain_runs": 0,
            "tool_calls": 0,
        }
        self._call_stack = []

    def report(self) -> Dict[str, Any]:
        """Generate a summary report of metrics."""
        return {
            "agent_name": self.agent_name,
            "agent_id": self.agent_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "metrics": self.metrics,
        }


# Backward compatibility alias
RoraCallbackHandler = AgentStatusCallbackHandler
