"""
Nexract Python SDK
==================
Web extraction API for AI agents and LLMs.

Usage:
    from nexract import Nexract

    client = Nexract(api_key="SK-LAB-YOUR-KEY")
    result = client.extract("https://example.com")
    print(result.markdown)

Docs: https://www.nexract.ai/integrations.html
"""

import requests
import time
from dataclasses import dataclass, field
from typing import Optional, List

__version__ = "1.0.0"
__author__  = "Nexract"
__url__     = "https://www.nexract.ai"

BASE_URL = "https://www.nexract.ai"


# ──────────────────────────────────────────
# Data Models
# ──────────────────────────────────────────

@dataclass
class ExtractionResult:
    """Result of a single web extraction."""
    success    : bool
    markdown   : str = ""
    title      : str = ""
    author     : str = ""
    excerpt    : str = ""
    word_count : int = 0
    duration_ms: int = 0
    error      : str = ""

    def to_dict(self) -> dict:
        return {
            "success"    : self.success,
            "title"      : self.title,
            "author"     : self.author,
            "excerpt"    : self.excerpt,
            "word_count" : self.word_count,
            "duration_ms": self.duration_ms,
            "markdown"   : self.markdown,
            "error"      : self.error,
        }

    def save(self, filename: Optional[str] = None) -> str:
        """Save extracted content to a text file."""
        if not filename:
            filename = f"nexract_{time.strftime('%Y-%m-%d_%H-%M-%S')}.txt"
        with open(filename, "w", encoding="utf-8") as f:
            f.write(f"Title:  {self.title}\n")
            f.write(f"Author: {self.author}\n")
            f.write(f"Words:  {self.word_count}\n")
            f.write("─" * 60 + "\n\n")
            f.write(self.markdown)
        print(f"Saved → {filename}")
        return filename


@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)


# ──────────────────────────────────────────
# Exceptions
# ──────────────────────────────────────────

class NexractError(Exception):
    """Base exception for Nexract SDK errors."""

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

class AuthenticationError(NexractError):
    """Raised when API key is invalid."""

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


# ──────────────────────────────────────────
# Main Client
# ──────────────────────────────────────────

class Nexract:
    """
    Nexract web extraction client.

    Args:
        api_key    : Your API key (starts with SK-LAB-)
        timeout    : Request timeout in seconds (default: 120)
        base_url   : API base URL (default: https://www.nexract.ai)

    Example:
        client = Nexract(api_key="SK-LAB-YOUR-KEY")
        result = client.extract("https://news.ycombinator.com")
        print(result.markdown)
    """

    def __init__(
        self,
        api_key : str,
        timeout : int = 120,
        base_url: str = BASE_URL,
    ):
        if not api_key or not api_key.startswith("SK-LAB-"):
            raise AuthenticationError(
                "Invalid API key format. Keys must start with 'SK-LAB-'.\n"
                "Get your key at: https://www.nexract.ai/client.html"
            )

        self.api_key  = api_key
        self.timeout  = timeout
        self.base_url = base_url.rstrip("/")

        self._session = requests.Session()
        self._session.headers.update({
            "Content-Type": "application/json",
            "User-Agent"  : f"nexract-python/{__version__}",
        })

    # ──────────────────────────────────────
    # Public: Check Balance
    # ──────────────────────────────────────

    def balance(self) -> dict:
        """
        Check your wallet balance.

        Returns:
            dict with keys: clientName, balance, totalUsed

        Example:
            info = client.balance()
            print(f"Credits remaining: {info['balance']}")
        """
        try:
            res  = self._session.get(
                f"{self.base_url}/api/balance/{self.api_key}",
                timeout=15,
            )
            data = res.json()
            if data.get("success"):
                return data.get("wallet", data)
            raise NexractError(data.get("error", "Balance check failed."))
        except requests.exceptions.ConnectionError:
            raise NexractError("Could not connect to Nexract API.")

    # ──────────────────────────────────────
    # Public: Extract Single URL
    # ──────────────────────────────────────

    def extract(
        self,
        url     : str,
        selector: str = "auto",
    ) -> ExtractionResult:
        """
        Extract clean markdown content from any URL.

        Args:
            url      : Target URL to extract
            selector : CSS selector (optional, default: auto-detect)

        Returns:
            ExtractionResult

        Raises:
            InsufficientCreditsError : No credits remaining
            AuthenticationError      : Invalid API key
            RateLimitError           : Too many requests
            NexractError             : Other API errors

        Example:
            result = client.extract("https://example.com/article")
            if result.success:
                print(result.markdown)
        """
        token = self._buy_token()
        return self._do_extract(url, token, selector)

    # ──────────────────────────────────────
    # Public: Extract Multiple URLs
    # ──────────────────────────────────────

    def extract_batch(
        self,
        urls    : List[str],
        selector: str   = "auto",
        delay   : float = 2.0,
    ) -> BatchResult:
        """
        Extract multiple URLs sequentially.

        Args:
            urls     : List of URLs to extract
            selector : CSS selector (applied to all URLs)
            delay    : Seconds to wait between requests (default: 2.0)

        Returns:
            BatchResult (iterable)

        Example:
            results = client.extract_batch([
                "https://example.com/page1",
                "https://example.com/page2",
            ])
            for r in results:
                print(r.title, r.word_count)
        """
        batch = BatchResult(total=len(urls))

        for i, url in enumerate(urls, 1):
            try:
                result = self.extract(url, selector)
            except NexractError as e:
                result = ExtractionResult(success=False, error=str(e))

            batch.results.append(result)
            if result.success:
                batch.successful += 1
            else:
                batch.failed += 1

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

        return batch

    # ──────────────────────────────────────
    # OpenAI Tool Definition (Function Calling)
    # ──────────────────────────────────────

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

        Example:
            import openai
            client = Nexract(api_key="SK-LAB-...")

            response = openai.chat.completions.create(
                model="gpt-4o",
                messages=[{"role": "user", "content": "Summarize https://example.com"}],
                tools=[Nexract.openai_tool()],
            )
        """
        return {
            "type": "function",
            "function": {
                "name"       : "extract_webpage",
                "description": (
                    "Extract clean text content from any webpage URL. "
                    "Returns the page content as markdown. "
                    "Use this whenever you need to read, summarize, or analyze a live webpage."
                ),
                "parameters": {
                    "type"      : "object",
                    "properties": {
                        "url": {
                            "type"       : "string",
                            "description": "The full URL of the webpage to extract (must include https://)",
                        }
                    },
                    "required": ["url"],
                },
            },
        }

    def handle_tool_call(self, url: str) -> str:
        """
        Handle an OpenAI tool call — extract URL and return markdown string.

        Example:
            # Inside your OpenAI tool loop:
            if tool_call.function.name == "extract_webpage":
                args   = json.loads(tool_call.function.arguments)
                result = client.handle_tool_call(args["url"])
                # Send result back as tool message
        """
        result = self.extract(url)
        if result.success:
            return result.markdown
        return f"[Extraction failed: {result.error}]"

    # ──────────────────────────────────────
    # Private: Buy Token
    # ──────────────────────────────────────

    def _buy_token(self) -> str:
        try:
            res  = self._session.post(
                f"{self.base_url}/api/buy-token",
                json={"apiKey": self.api_key},
                timeout=15,
            )
            data = res.json()

            if data.get("success"):
                return data["token"]

            error = data.get("error", "")
            if "INSUFFICIENT" in error or "balance" in error.lower():
                raise InsufficientCreditsError(
                    "No credits remaining. Top up at: https://www.nexract.ai/pricing.html"
                )
            if "invalid" in error.lower() or "key" in error.lower():
                raise AuthenticationError(f"Invalid API key: {error}")

            raise NexractError(f"Token error: {error}")

        except (InsufficientCreditsError, AuthenticationError, NexractError):
            raise
        except requests.exceptions.ConnectionError:
            raise NexractError("Could not connect to Nexract API.")

    # ──────────────────────────────────────
    # Private: Do Extraction
    # ──────────────────────────────────────

    def _do_extract(self, url: str, token: str, selector: str) -> ExtractionResult:
        try:
            res = self._session.post(
                f"{self.base_url}/api/extract-data",
                json={"token": token, "targetUrl": url, "selector": selector},
                timeout=self.timeout,
            )

            if res.status_code == 429:
                raise RateLimitError("Rate limit exceeded. Wait before retrying.")

            data = res.json()

            if data.get("success"):
                article = data.get("article_data", {})
                meta    = data.get("meta", {})
                return ExtractionResult(
                    success    = True,
                    markdown   = data.get("data", ""),
                    title      = article.get("title", ""),
                    author     = article.get("author", ""),
                    excerpt    = article.get("excerpt", ""),
                    word_count = meta.get("word_count", 0),
                    duration_ms= meta.get("duration_ms", 0),
                )

            return ExtractionResult(
                success=False,
                error=data.get("error", "Extraction failed.")
            )

        except (RateLimitError, NexractError):
            raise
        except requests.exceptions.Timeout:
            return ExtractionResult(success=False, error=f"Timeout after {self.timeout}s.")
        except Exception as e:
            return ExtractionResult(success=False, error=str(e))

    def __repr__(self):
        return f"Nexract(api_key='{self.api_key[:12]}...', base_url='{self.base_url}')"
