"""
Nexract — Web Intelligence API SDK for Python
Extract, understand, and structure any webpage with one API call.

Usage:
    from nexract import Nexract

    nx = Nexract("SK-LAB-YOUR-KEY")
    result = nx.extract("https://example.com/article")
    print(result.data)

    # With AI analysis
    result = nx.extract("https://example.com/article", ai=True)
    print(result.ai["summary"])
    print(result.ai["sentiment"])

    # With custom schema
    result = nx.extract("https://amazon.com/product", schema={
        "product_name": "string",
        "price": "number",
        "rating": "number"
    })
    print(result.structured_data)

    # Arabic content
    result = nx.extract("https://aljazeera.net/article", ai=True, language="ar")
    print(result.ai["summary_ar"])

AI Tool Integration:
    # OpenAI function calling
    tools = [nx.openai_tool()]

    # Anthropic Claude tool use
    tools = [nx.claude_tool()]

    # LangChain
    tool = nx.langchain_tool()

Install:
    pip install nexract
"""

import requests
import time
import json
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List, Callable


__version__ = "2.0.0"
__api_base__ = "https://www.nexract.ai"


# ══════════════════════════════════════════════════════════
# DATA MODELS
# ══════════════════════════════════════════════════════════

@dataclass
class ExtractionResult:
    """Result of a Nexract extraction."""
    success: bool
    data: str = ""
    format: str = "markdown"
    article: Dict[str, Any] = field(default_factory=dict)
    ai: Optional[Dict[str, Any]] = None
    structured_data: Optional[Dict[str, Any]] = None
    meta: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None
    error_code: Optional[str] = None
    raw: Dict[str, Any] = field(default_factory=dict)

    @property
    def title(self) -> str:
        return self.article.get("title", "")

    @property
    def author(self) -> str:
        return self.article.get("author", "")

    @property
    def word_count(self) -> int:
        return self.article.get("word_count", 0)

    @property
    def summary(self) -> str:
        if self.ai:
            return self.ai.get("summary", self.ai.get("summary_ar", ""))
        return ""

    @property
    def sentiment(self) -> str:
        if self.ai:
            return self.ai.get("sentiment", "")
        return ""

    @property
    def duration_ms(self) -> int:
        return self.meta.get("duration_ms", 0)

    @property
    def cached(self) -> bool:
        return self.meta.get("cached", False)


@dataclass
class BalanceResult:
    """Result of a balance check."""
    success: bool
    balance: int = 0
    total_used: int = 0
    error: Optional[str] = None


@dataclass
class BatchResult:
    """Result of a batch extraction."""
    results: List[ExtractionResult] = field(default_factory=list)
    total: int = 0
    successful: int = 0
    failed: int = 0

    def __iter__(self):
        return iter(self.results)

    def __len__(self):
        return len(self.results)


# ══════════════════════════════════════════════════════════
# EXCEPTIONS
# ══════════════════════════════════════════════════════════

class NexractError(Exception):
    """Nexract API error."""
    def __init__(self, message: str, code: str = None, status: int = None):
        self.message = message
        self.code = code
        self.status = status
        super().__init__(message)


class InsufficientCreditsError(NexractError):
    """Raised when wallet has no credits."""
    pass


class RateLimitError(NexractError):
    """Raised when rate limit is exceeded."""
    pass


# ══════════════════════════════════════════════════════════
# TOOL DEFINITIONS — Used by OpenAI, Claude, LangChain
# ══════════════════════════════════════════════════════════

# Shared description used across all tool formats
_TOOL_NAME = "nexract_extract"
_TOOL_DESCRIPTION = (
    "Extract clean text content from any webpage URL. "
    "Returns the page content as clean markdown with title, author, and word count. "
    "Use this whenever you need to read, summarize, analyze, or extract data from a live webpage. "
    "Supports any public URL including news articles, blogs, documentation, and product pages."
)
_TOOL_PARAMETERS = {
    "type": "object",
    "properties": {
        "url": {
            "type": "string",
            "description": "The full URL of the webpage to extract (must include https://)",
        },
        "ai": {
            "type": "boolean",
            "description": "Set to true to get AI analysis (summary, sentiment, entities, category). Costs 2 extra credits.",
            "default": False,
        },
        "format": {
            "type": "string",
            "enum": ["markdown", "text", "html"],
            "description": "Output format for the extracted content.",
            "default": "markdown",
        },
    },
    "required": ["url"],
}


# ══════════════════════════════════════════════════════════
# MAIN CLIENT
# ══════════════════════════════════════════════════════════

class Nexract:
    """
    Nexract Web Intelligence API client.

    Args:
        api_key: Your Nexract API key (SK-LAB-...)
        base_url: API base URL (default: https://www.nexract.ai)
        timeout: Request timeout in seconds (default: 120)
        max_retries: Max retry attempts for 503/429 errors (default: 3)
    """

    def __init__(self, api_key: str, base_url: str = None, timeout: int = 120, max_retries: int = 3):
        if not api_key or not api_key.startswith("SK-LAB-"):
            raise ValueError("Invalid API key. Must start with SK-LAB-")
        self.api_key = api_key
        self.base_url = (base_url or __api_base__).rstrip("/")
        self.timeout = timeout
        self.max_retries = max_retries
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": f"nexract-python/{__version__}",
        })

    # ──────────────────────────────────────────────────────
    # CORE: Extract
    # ──────────────────────────────────────────────────────

    def extract(
        self,
        url: str,
        format: str = "markdown",
        ai: bool = False,
        schema: Dict[str, Any] = None,
        language: str = "auto",
        selector: str = "auto",
    ) -> ExtractionResult:
        """
        Extract and analyze a webpage.

        Args:
            url: The webpage URL to extract
            format: Output format — "markdown", "html", or "text"
            ai: Enable AI analysis (summary, entities, sentiment) — costs 2 extra credits
            schema: JSON schema for structured extraction — costs 2 extra credits
            language: Language hint ("auto", "ar", "en", etc.)
            selector: CSS selector for targeted extraction (default: "auto")

        Returns:
            ExtractionResult with extracted data and optional AI analysis

        Raises:
            NexractError: If the API returns an error
        """
        body = {"url": url, "format": format, "selector": selector}
        if ai:
            body["ai"] = True
            body["language"] = language
        if schema:
            body["schema"] = schema

        data = self._request_with_retry("POST", "/api/v2/extract", body)

        return ExtractionResult(
            success=True,
            data=data.get("data", ""),
            format=data.get("format", format),
            article=data.get("article", {}),
            ai=data.get("ai"),
            structured_data=data.get("structured_data"),
            meta=data.get("meta", {}),
            raw=data,
        )

    # ──────────────────────────────────────────────────────
    # CORE: Batch Extract
    # ──────────────────────────────────────────────────────

    def batch_extract(
        self,
        urls: List[str],
        format: str = "markdown",
        ai: bool = False,
        schema: Dict[str, Any] = None,
        language: str = "auto",
        delay: float = 1.0,
    ) -> BatchResult:
        """
        Extract multiple URLs sequentially.

        Args:
            urls: List of URLs to extract
            format: Output format (applied to all URLs)
            ai: Enable AI analysis for all URLs
            schema: JSON schema (applied to all URLs)
            language: Language hint
            delay: Seconds between requests (default: 1.0)

        Returns:
            BatchResult (iterable)
        """
        batch = BatchResult(total=len(urls))

        for i, url in enumerate(urls):
            try:
                result = self.extract(url, format=format, ai=ai, schema=schema, language=language)
                batch.results.append(result)
                batch.successful += 1
            except NexractError as e:
                batch.results.append(ExtractionResult(
                    success=False, error=e.message, error_code=e.code
                ))
                batch.failed += 1

            if i < len(urls) - 1:
                time.sleep(delay)

        return batch

    # ──────────────────────────────────────────────────────
    # CORE: Balance
    # ──────────────────────────────────────────────────────

    def balance(self) -> BalanceResult:
        """Check your credit balance."""
        resp = self.session.post(
            f"{self.base_url}/api/balance",
            json={},
            timeout=30,
        )
        data = resp.json()
        return BalanceResult(
            success=data.get("success", False),
            balance=data.get("balance", 0),
            total_used=data.get("totalUsed", 0),
            error=data.get("error"),
        )

    # ══════════════════════════════════════════════════════
    # AI TOOL INTEGRATIONS
    # ══════════════════════════════════════════════════════

    # ──────────────────────────────────────────────────────
    # OpenAI Function Calling
    # ──────────────────────────────────────────────────────

    @staticmethod
    def openai_tool() -> dict:
        """
        Returns an OpenAI-compatible tool definition for function calling.

        Usage with OpenAI:
            from openai import OpenAI
            from nexract import Nexract

            nx = Nexract("SK-LAB-...")
            client = OpenAI()

            response = client.chat.completions.create(
                model="gpt-4o",
                messages=[{"role": "user", "content": "Summarize https://example.com"}],
                tools=[nx.openai_tool()],
            )

            # Handle the tool call
            for call in response.choices[0].message.tool_calls or []:
                if call.function.name == "nexract_extract":
                    args = json.loads(call.function.arguments)
                    result = nx.handle_tool_call(args)
                    # Send result back to OpenAI as tool message
        """
        return {
            "type": "function",
            "function": {
                "name": _TOOL_NAME,
                "description": _TOOL_DESCRIPTION,
                "parameters": _TOOL_PARAMETERS,
            },
        }

    # ──────────────────────────────────────────────────────
    # Anthropic Claude Tool Use
    # ──────────────────────────────────────────────────────

    @staticmethod
    def claude_tool() -> dict:
        """
        Returns an Anthropic Claude-compatible tool definition.

        Usage with Claude:
            import anthropic
            from nexract import Nexract

            nx = Nexract("SK-LAB-...")
            client = anthropic.Anthropic()

            response = client.messages.create(
                model="claude-sonnet-4-20250514",
                max_tokens=1024,
                tools=[nx.claude_tool()],
                messages=[{"role": "user", "content": "Extract https://example.com"}],
            )

            # Handle the tool call
            for block in response.content:
                if block.type == "tool_use" and block.name == "nexract_extract":
                    result = nx.handle_tool_call(block.input)
                    # Send result back to Claude as tool_result
        """
        return {
            "name": _TOOL_NAME,
            "description": _TOOL_DESCRIPTION,
            "input_schema": _TOOL_PARAMETERS,
        }

    # ──────────────────────────────────────────────────────
    # LangChain Tool
    # ──────────────────────────────────────────────────────

    def langchain_tool(self):
        """
        Returns a LangChain-compatible Tool object.

        Usage with LangChain:
            from langchain.agents import initialize_agent, AgentType
            from langchain_openai import ChatOpenAI
            from nexract import Nexract

            nx = Nexract("SK-LAB-...")
            llm = ChatOpenAI(model="gpt-4o")

            agent = initialize_agent(
                tools=[nx.langchain_tool()],
                llm=llm,
                agent=AgentType.OPENAI_FUNCTIONS,
            )

            result = agent.invoke("Summarize https://example.com")

        Requires: pip install langchain langchain-core
        """
        try:
            from langchain_core.tools import StructuredTool
            from pydantic import BaseModel, Field
        except ImportError:
            raise ImportError(
                "LangChain integration requires langchain-core and pydantic. "
                "Install with: pip install langchain-core pydantic"
            )

        client = self

        class NexractInput(BaseModel):
            url: str = Field(description="The full URL of the webpage to extract (must include https://)")
            ai: bool = Field(default=False, description="Enable AI analysis (summary, sentiment, entities)")
            format: str = Field(default="markdown", description="Output format: markdown, text, or html")

        def _run_extract(url: str, ai: bool = False, format: str = "markdown") -> str:
            try:
                result = client.extract(url, format=format, ai=ai)
                output = f"# {result.title}\n\n{result.data}"
                if result.ai:
                    output += f"\n\n---\nSummary: {result.summary}"
                    output += f"\nSentiment: {result.sentiment}"
                return output
            except NexractError as e:
                return f"[Extraction failed: {e.message}]"

        return StructuredTool.from_function(
            func=_run_extract,
            name=_TOOL_NAME,
            description=_TOOL_DESCRIPTION,
            args_schema=NexractInput,
        )

    # ──────────────────────────────────────────────────────
    # Universal Tool Call Handler
    # ──────────────────────────────────────────────────────

    def handle_tool_call(self, arguments: dict) -> str:
        """
        Handle a tool call from OpenAI, Claude, or any AI framework.
        Accepts the arguments dict and returns extracted content as string.

        Args:
            arguments: Dict with at least {"url": "..."}, optionally {"ai": true, "format": "..."}

        Returns:
            Extracted content as string (ready to send back to AI as tool result)
        """
        url = arguments.get("url", "")
        if not url:
            return "[Error: No URL provided in tool call]"

        ai = arguments.get("ai", False)
        fmt = arguments.get("format", "markdown")

        try:
            result = self.extract(url, format=fmt, ai=ai)

            output = result.data
            if result.ai:
                ai_section = "\n\n--- AI Analysis ---"
                if result.summary:
                    ai_section += f"\nSummary: {result.summary}"
                if result.sentiment:
                    ai_section += f"\nSentiment: {result.sentiment}"
                if result.ai.get("category"):
                    ai_section += f"\nCategory: {result.ai['category']}"
                if result.ai.get("key_points"):
                    ai_section += "\nKey Points:"
                    for point in result.ai["key_points"][:5]:
                        ai_section += f"\n  - {point}"
                output += ai_section

            return output

        except NexractError as e:
            return f"[Extraction failed: {e.message}]"

    # ══════════════════════════════════════════════════════
    # INTERNAL: Request with Retry
    # ══════════════════════════════════════════════════════

    def _request_with_retry(self, method: str, path: str, body: dict) -> dict:
        """Make an API request with exponential backoff retry for 429/503."""
        last_error = None

        for attempt in range(self.max_retries + 1):
            try:
                resp = self.session.request(
                    method=method,
                    url=f"{self.base_url}{path}",
                    json=body,
                    timeout=self.timeout,
                )

                data = resp.json()

                # Success
                if resp.ok and data.get("success"):
                    return data

                # Non-retryable errors
                error_msg = data.get("error", "Request failed")
                error_code = data.get("error_code")

                if resp.status_code == 402:
                    raise InsufficientCreditsError(error_msg, error_code, 402)
                if resp.status_code == 401:
                    raise NexractError(error_msg, error_code, 401)
                if resp.status_code == 400:
                    raise NexractError(error_msg, error_code, 400)
                if resp.status_code == 403:
                    raise NexractError(error_msg, error_code, 403)
                if resp.status_code == 404:
                    raise NexractError(error_msg, error_code, 404)
                if resp.status_code == 422:
                    raise NexractError(error_msg, error_code, 422)

                # Retryable errors (429, 503)
                if resp.status_code in (429, 503) and attempt < self.max_retries:
                    wait = min(2 ** attempt * 1.5, 30)  # 1.5s, 3s, 6s... max 30s
                    time.sleep(wait)
                    last_error = NexractError(error_msg, error_code, resp.status_code)
                    continue

                # Other errors — no retry
                raise NexractError(error_msg, error_code, resp.status_code)

            except requests.exceptions.Timeout:
                if attempt < self.max_retries:
                    wait = min(2 ** attempt * 1.5, 30)
                    time.sleep(wait)
                    last_error = NexractError("Request timeout", "TIMEOUT", 408)
                    continue
                raise NexractError("Request timeout", "TIMEOUT", 408)

            except requests.exceptions.ConnectionError:
                if attempt < self.max_retries:
                    wait = min(2 ** attempt * 1.5, 30)
                    time.sleep(wait)
                    last_error = NexractError("Connection failed", "NETWORK_ERROR", 0)
                    continue
                raise NexractError("Could not connect to Nexract API", "NETWORK_ERROR", 0)

            except NexractError:
                raise

            except Exception as e:
                raise NexractError(str(e), "UNKNOWN_ERROR", 0)

        # All retries exhausted
        if last_error:
            raise last_error
        raise NexractError("Request failed after retries", "RETRY_EXHAUSTED", 0)

    def __repr__(self):
        masked = self.api_key[:10] + "..." + self.api_key[-4:]
        return f"Nexract(key={masked})"
