from __future__ import annotations

import time
from typing import Any

from leshy.core.models import TargetConfig, TargetResponse
from leshy.core.registry import target_registry
from leshy.targets.base import Target


@target_registry.register("langgraph")
class LangGraphTarget(Target):
    """LangGraph target for testing deployed LangGraph agents via langgraph-sdk."""

    name = "langgraph"

    def __init__(self, config: TargetConfig, **kwargs: Any) -> None:
        super().__init__(config, **kwargs)
        self._client: Any = None
        self._assistant_id: str | None = config.headers.get("x-assistant-id")

    def _get_client(self) -> Any:
        try:
            from langgraph_sdk import get_client
        except ImportError:
            raise ImportError(
                "The 'langgraph-sdk' package is required for LangGraph targets. "
                "Install it with: pip install 'leshy[langgraph]'"
            )
        return get_client(url=self.config.url)

    async def setup(self) -> None:
        self._client = self._get_client()
        if not self._assistant_id:
            assistants = await self._client.assistants.search()
            if assistants:
                self._assistant_id = assistants[0]["assistant_id"]
            else:
                raise RuntimeError("No assistants found on LangGraph server")

    async def send(self, payload: str, *, thread_id: str | None = None) -> TargetResponse:
        if self._client is None:
            await self.setup()

        start = time.monotonic()

        try:
            if thread_id is None:
                thread = await self._client.threads.create()
                thread_id = thread["thread_id"]

            input_data = {"messages": [{"role": "human", "content": payload}]}

            final_content = ""
            raw_data: dict[str, Any] = {}

            async for chunk in self._client.runs.stream(
                thread_id,
                self._assistant_id,
                input=input_data,
                stream_mode="updates",
            ):
                raw_data = chunk.data if hasattr(chunk, "data") else {}
                if isinstance(raw_data, dict) and "messages" in raw_data:
                    messages = raw_data["messages"]
                    if messages:
                        last_msg = messages[-1]
                        if isinstance(last_msg, dict):
                            final_content = last_msg.get("content", "")
                        elif hasattr(last_msg, "content"):
                            final_content = last_msg.content

            latency = (time.monotonic() - start) * 1000
            return TargetResponse(
                content=final_content,
                raw=raw_data if isinstance(raw_data, dict) else {},
                status_code=200,
                latency_ms=latency,
                thread_id=thread_id,
            )

        except Exception as e:
            latency = (time.monotonic() - start) * 1000
            return TargetResponse(
                content=f"LangGraph Error: {e}",
                raw={"error": str(e)},
                status_code=0,
                latency_ms=latency,
                thread_id=thread_id,
            )

    async def teardown(self) -> None:
        self._client = None
