"""LangChain tools for the ihyee web intelligence API."""
from __future__ import annotations

import os
from typing import Any, Optional, Type

from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field, PrivateAttr

from ihyee import Ihyee
from ihyee_langchain.formatter import format_fetch_results, format_search_results


class SearchInput(BaseModel):
    """Input for the ihyee search tool."""
    query: str = Field(description="The search query string.")


class FetchInput(BaseModel):
    """Input for the ihyee fetch tool."""
    url: str = Field(
        description="URL to fetch content from. For multiple URLs, separate them with commas."
    )


class RenderInput(BaseModel):
    """Input for the ihyee render tool."""
    url: str = Field(
        description="URL to render with a headless browser and extract content from."
    )


class IhyeeSearchTool(BaseTool):
    """Search the web using ihyee and return extracted, summarized content.

    Wraps ``ihyee.Ihyee.search`` for use as a LangChain tool.
    """

    name: str = "ihyee_search"
    description: str = (
        "Search the web and return extracted, summarized content from top results. "
        "Input should be a search query string. Returns text summaries and "
        "full content from the most relevant web pages."
    )
    args_schema: Type[BaseModel] = SearchInput
    api_key: Optional[str] = Field(default=None, exclude=True)
    max_results: int = 5
    content_mode: str = "both"

    _client: Optional[Ihyee] = PrivateAttr(default=None)

    def _get_client(self) -> Ihyee:
        if self._client is None:
            self._client = Ihyee(api_key=self.api_key or os.environ.get("IHYEE_API_KEY"))
        return self._client

    def _run(self, query: str, **kwargs: Any) -> str:
        client = self._get_client()
        response = client.search(
            query,
            max_results=self.max_results,
            content_mode=self.content_mode,
        )
        return format_search_results(response)


class IhyeeFetchTool(BaseTool):
    """Fetch and extract content from specific URLs using ihyee.

    Wraps ``ihyee.Ihyee.fetch`` for use as a LangChain tool.
    """

    name: str = "ihyee_fetch"
    description: str = (
        "Fetch and extract content from one or more specific URLs. "
        "Input should be a URL, or multiple URLs separated by commas. "
        "Returns the extracted text content and metadata from each page."
    )
    args_schema: Type[BaseModel] = FetchInput
    api_key: Optional[str] = Field(default=None, exclude=True)
    content_mode: str = "both"

    _client: Optional[Ihyee] = PrivateAttr(default=None)

    def _get_client(self) -> Ihyee:
        if self._client is None:
            self._client = Ihyee(api_key=self.api_key or os.environ.get("IHYEE_API_KEY"))
        return self._client

    def _run(self, url: str, **kwargs: Any) -> str:
        client = self._get_client()
        urls = [u.strip() for u in url.split(",") if u.strip()]
        response = client.fetch(urls, content_mode=self.content_mode)
        return format_fetch_results(response)


class IhyeeRenderTool(BaseTool):
    """Render JavaScript-heavy pages using ihyee's headless browser.

    Wraps ``ihyee.Ihyee.render`` for use as a LangChain tool.
    Use this for single-page applications (SPAs) or pages that require
    JavaScript execution to display their content.
    """

    name: str = "ihyee_render"
    description: str = (
        "Render a JavaScript-heavy web page using a headless browser and "
        "extract its content. Use this for single-page applications (SPAs) "
        "or dynamic pages that need JavaScript to display content. "
        "Input should be a single URL."
    )
    args_schema: Type[BaseModel] = RenderInput
    api_key: Optional[str] = Field(default=None, exclude=True)
    content_mode: str = "both"
    wait_for: str = "networkidle"
    timeout_ms: int = 30000

    _client: Optional[Ihyee] = PrivateAttr(default=None)

    def _get_client(self) -> Ihyee:
        if self._client is None:
            self._client = Ihyee(api_key=self.api_key or os.environ.get("IHYEE_API_KEY"))
        return self._client

    def _run(self, url: str, **kwargs: Any) -> str:
        client = self._get_client()
        response = client.render(
            url,
            content_mode=self.content_mode,
            wait_for=self.wait_for,
            timeout_ms=self.timeout_ms,
        )
        return format_fetch_results(response)
