"""
LangChain tools for Olostep web scraping API.

Comprehensive tools matching all Olostep API capabilities.
"""

import os
from typing import Literal, List, Dict, Any, Optional, Union
from langchain_core.tools import tool
from langchain_core.exceptions import LangChainException
import requests
import json


class OlostepClient:
    """Client for Olostep API with full feature support."""
    
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv("OLOSTEP_API_KEY")
        if not self.api_key:
            raise ValueError("OLOSTEP_API_KEY environment variable is required")
        self.base_url = "https://api.olostep.com/v1"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json"
        }
    
    def scrape(
        self, 
        url: str,
        formats: Optional[List[str]] = None,
        country: Optional[str] = None,
        wait_before_scraping: int = 0,
        parser_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Scrape a single URL using Olostep API.
        
        Args:
            url: URL to scrape
            formats: Output formats (markdown, html, json, text)
            country: Country code for location-specific scraping
            wait_before_scraping: Wait time in milliseconds before scraping
            parser_id: Optional parser ID for specialized extraction
            
        Returns:
            Scraped content and metadata
        """
        payload = {
            "url_to_scrape": url,
            "formats": formats or ["markdown", "html", "json", "text"]
        }
        
        if country:
            payload["country"] = country
        if wait_before_scraping:
            payload["wait_before_scraping"] = wait_before_scraping
        if parser_id:
            payload["parser"] = {"id": parser_id}
        
        try:
            response = requests.post(
                f"{self.base_url}/scrapes",
                headers=self.headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise LangChainException(f"Olostep scrape failed: {str(e)}")
    
    def batch_scrape(
        self,
        urls: List[Dict[str, str]],
        formats: Optional[List[str]] = None,
        country: Optional[str] = None,
        wait_before_scraping: int = 0,
        parser_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Scrape multiple URLs in batch.
        
        Args:
            urls: List of dicts with 'url' and optional 'custom_id'
            formats: Output formats
            country: Country code
            wait_before_scraping: Wait time in milliseconds
            parser_id: Optional parser ID
            
        Returns:
            Batch information
        """
        payload = {
            "items": urls
        }
        
        if formats:
            payload["formats"] = formats
        if country:
            payload["country"] = country
        if wait_before_scraping:
            payload["wait_before_scraping"] = wait_before_scraping
        if parser_id:
            payload["parser"] = {"id": parser_id}
        
        try:
            response = requests.post(
                f"{self.base_url}/batches",
                headers=self.headers,
                json=payload,
                timeout=120
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise LangChainException(f"Olostep batch scrape failed: {str(e)}")
    
    def answer(
        self,
        url: str,
        question: str,
        formats: Optional[List[str]] = None,
        country: Optional[str] = None,
        wait_before_scraping: int = 0
    ) -> Dict[str, Any]:
        """
        Ask a question about a website and get an AI-powered answer.
        
        Args:
            url: URL to analyze
            question: Question to answer about the content
            formats: Output formats
            country: Country code
            wait_before_scraping: Wait time in milliseconds
            
        Returns:
            Answer and scraped content
        """
        payload = {
            "url_to_scrape": url,
            "question": question,
            "formats": formats or ["markdown"]
        }
        
        if country:
            payload["country"] = country
        if wait_before_scraping:
            payload["wait_before_scraping"] = wait_before_scraping
        
        try:
            response = requests.post(
                f"{self.base_url}/answers",
                headers=self.headers,
                json=payload,
                timeout=90
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise LangChainException(f"Olostep answer failed: {str(e)}")
    
    def map_extract(
        self,
        url: str,
        fields: List[str],
        formats: Optional[List[str]] = None,
        country: Optional[str] = None,
        wait_before_scraping: int = 0
    ) -> Dict[str, Any]:
        """
        Extract specific fields from a website using AI mapping.
        
        Args:
            url: URL to extract from
            fields: List of field names to extract
            formats: Output formats
            country: Country code
            wait_before_scraping: Wait time in milliseconds
            
        Returns:
            Extracted fields and metadata
        """
        payload = {
            "url_to_scrape": url,
            "fields": fields,
            "formats": formats or ["json"]
        }
        
        if country:
            payload["country"] = country
        if wait_before_scraping:
            payload["wait_before_scraping"] = wait_before_scraping
        
        try:
            response = requests.post(
                f"{self.base_url}/maps",
                headers=self.headers,
                json=payload,
                timeout=90
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise LangChainException(f"Olostep map extraction failed: {str(e)}")


# LangChain Tools

@tool
def scrape_website(
    url: str,
    format: Literal["markdown", "html", "json", "text"] = "markdown",
    country: Optional[str] = None,
    wait_before_scraping: int = 0,
    parser: Optional[str] = None,
    api_key: Optional[str] = None
) -> str:
    """
    Scrape content from any website using Olostep's powerful web scraping API.
    
    Perfect for extracting content from dynamic websites, handling JavaScript rendering,
    and bypassing anti-scraping measures. Supports multiple output formats and
    specialized parsers for specific websites.
    
    Args:
        url: Website URL to scrape (must include http:// or https://)
        format: Output format - markdown, html, json, or text
        country: Country code for location-specific content (e.g., 'US', 'GB', 'CA')
        wait_before_scraping: Wait time in milliseconds for JavaScript rendering (0-10000)
        parser: Optional parser ID for specialized extraction (e.g., '@olostep/amazon-product')
        api_key: Olostep API key (uses OLOSTEP_API_KEY env var if not provided)
    
    Returns:
        Scraped content in requested format with metadata
    
    Examples:
        >>> # Basic scraping
        >>> content = scrape_website("https://example.com")
        
        >>> # With JavaScript wait time
        >>> content = scrape_website("https://example.com", wait_before_scraping=2000)
        
        >>> # With specialized parser
        >>> content = scrape_website(
        ...     "https://amazon.com/product/123",
        ...     parser="@olostep/amazon-product"
        ... )
    """
    try:
        client = OlostepClient(api_key)
        result = client.scrape(
            url=url,
            formats=[format],
            country=country,
            wait_before_scraping=wait_before_scraping,
            parser_id=parser
        )
        
        # Extract content from result
        content_result = result.get("result", {})
        
        # Get the requested format content
        format_key = f"{format}_content"
        content = content_result.get(format_key, "")
        
        # If content is a dict (for json format), convert to string
        if isinstance(content, dict):
            content = json.dumps(content, indent=2)
        
        # Build response with metadata
        response = {
            "content": content,
            "url": url,
            "format": format,
            "scrape_id": result.get("retrieve_id", ""),
            "metadata": content_result.get("page_metadata", {})
        }
        
        # Add hosted URLs if available
        if f"{format}_hosted_url" in content_result:
            response[f"{format}_url"] = content_result[f"{format}_hosted_url"]
        
        return json.dumps(response, indent=2)
        
    except Exception as e:
        raise LangChainException(f"Failed to scrape {url}: {str(e)}")


@tool
def scrape_batch(
    urls: List[str],
    format: Literal["markdown", "html", "json", "text"] = "markdown",
    country: Optional[str] = None,
    wait_before_scraping: int = 0,
    parser: Optional[str] = None,
    api_key: Optional[str] = None
) -> str:
    """
    Scrape multiple websites in parallel using Olostep's batch processing API.
    
    Perfect for large-scale data extraction, competitor analysis, or building
    comprehensive datasets. Supports up to 100,000 URLs in a single batch.
    
    Args:
        urls: List of website URLs to scrape
        format: Output format for all URLs - markdown, html, json, or text
        country: Country code for location-specific content
        wait_before_scraping: Wait time in milliseconds for JavaScript rendering
        parser: Optional parser ID for specialized extraction
        api_key: Olostep API key (uses OLOSTEP_API_KEY env var if not provided)
    
    Returns:
        Batch job information with status and batch ID
    
    Examples:
        >>> # Scrape multiple websites
        >>> urls = [
        ...     "https://example1.com",
        ...     "https://example2.com",
        ...     "https://example3.com"
        ... ]
        >>> result = scrape_batch(urls)
        
        >>> # With custom options
        >>> result = scrape_batch(
        ...     urls,
        ...     format="json",
        ...     country="US",
        ...     wait_before_scraping=2000
        ... )
    """
    try:
        client = OlostepClient(api_key)
        
        # Convert URLs to batch format
        batch_items = [
            {"url": url, "custom_id": f"url_{i}"}
            for i, url in enumerate(urls)
        ]
        
        result = client.batch_scrape(
            urls=batch_items,
            formats=[format],
            country=country,
            wait_before_scraping=wait_before_scraping,
            parser_id=parser
        )
        
        response = {
            "batch_id": result.get("batch_id", result.get("id", "")),
            "status": result.get("status", "processing"),
            "total_urls": len(urls),
            "format": format,
            "urls": urls
        }
        
        if country:
            response["country"] = country
        if parser:
            response["parser"] = parser
        
        return json.dumps(response, indent=2)
        
    except Exception as e:
        raise LangChainException(f"Failed to scrape batch: {str(e)}")


@tool
def scrape_with_answer(
    url: str,
    question: str,
    format: Literal["markdown", "html", "json", "text"] = "markdown",
    country: Optional[str] = None,
    wait_before_scraping: int = 0,
    api_key: Optional[str] = None
) -> str:
    """
    Scrape a website and get an AI-powered answer to a specific question about its content.
    
    This powerful feature combines web scraping with AI analysis to extract
    specific information from websites. Perfect for research, data extraction,
    and competitive intelligence.
    
    Args:
        url: Website URL to analyze
        question: Question to answer about the website content
        format: Output format for scraped content
        country: Country code for location-specific content
        wait_before_scraping: Wait time in milliseconds for JavaScript rendering
        api_key: Olostep API key (uses OLOSTEP_API_KEY env var if not provided)
    
    Returns:
        AI-generated answer with scraped content and confidence score
    
    Examples:
        >>> # Extract specific information
        >>> result = scrape_with_answer(
        ...     "https://company.com",
        ...     "What is the company's main product?"
        ... )
        
        >>> # Get pricing information
        >>> result = scrape_with_answer(
        ...     "https://competitor.com/pricing",
        ...     "What are the pricing tiers and their features?"
        ... )
        
        >>> # Extract contact details
        >>> result = scrape_with_answer(
        ...     "https://business.com/contact",
        ...     "What are the email addresses and phone numbers?"
        ... )
    """
    try:
        client = OlostepClient(api_key)
        result = client.answer(
            url=url,
            question=question,
            formats=[format],
            country=country,
            wait_before_scraping=wait_before_scraping
        )
        
        response = {
            "answer": result.get("answer", ""),
            "question": question,
            "url": url,
            "confidence": result.get("confidence", ""),
            "content": result.get("result", {}).get(f"{format}_content", ""),
            "metadata": result.get("result", {}).get("page_metadata", {})
        }
        
        return json.dumps(response, indent=2)
        
    except Exception as e:
        raise LangChainException(f"Failed to get answer from {url}: {str(e)}")


@tool
def scrape_with_map(
    url: str,
    fields: List[str],
    format: Literal["json", "markdown"] = "json",
    country: Optional[str] = None,
    wait_before_scraping: int = 0,
    api_key: Optional[str] = None
) -> str:
    """
    Extract specific structured data fields from a website using AI-powered mapping.
    
    Perfect for extracting product details, contact information, pricing data,
    or any structured information from websites. The AI automatically identifies
    and extracts the requested fields.
    
    Args:
        url: Website URL to extract from
        fields: List of field names to extract (e.g., ['price', 'title', 'description'])
        format: Output format - json or markdown
        country: Country code for location-specific content
        wait_before_scraping: Wait time in milliseconds for JavaScript rendering
        api_key: Olostep API key (uses OLOSTEP_API_KEY env var if not provided)
    
    Returns:
        Extracted fields as structured data
    
    Examples:
        >>> # Extract product information
        >>> result = scrape_with_map(
        ...     "https://store.com/product/123",
        ...     fields=["product_name", "price", "description", "rating"]
        ... )
        
        >>> # Extract contact information
        >>> result = scrape_with_map(
        ...     "https://company.com/about",
        ...     fields=["company_name", "email", "phone", "address"]
        ... )
        
        >>> # Extract article data
        >>> result = scrape_with_map(
        ...     "https://blog.com/article",
        ...     fields=["title", "author", "publish_date", "content"]
        ... )
    """
    try:
        client = OlostepClient(api_key)
        result = client.map_extract(
            url=url,
            fields=fields,
            formats=[format],
            country=country,
            wait_before_scraping=wait_before_scraping
        )
        
        response = {
            "extracted_data": result.get("result", {}).get("extracted_fields", {}),
            "url": url,
            "fields": fields,
            "metadata": result.get("result", {}).get("page_metadata", {})
        }
        
        return json.dumps(response, indent=2)
        
    except Exception as e:
        raise LangChainException(f"Failed to extract fields from {url}: {str(e)}")
