#!/usr/bin/env python3
"""
Tavily Search v1.0 - AI-powered web search platform with 5 modes.
Author: Leo 🦁
Created: 2026-02-07

Commands:
- search: General web search with optional LLM answer
- news: News-optimized search (topic=news)
- finance: Finance-optimized search (topic=finance)
- extract: Extract content from URLs
- crawl: Crawl a website
- map: Discover sitemap URLs
- research: Deep AI research with citations

Environment Variables:
- TAVILY_API_KEY: Required. Tavily API key.
"""

import argparse
import json
import os
import sys
import urllib.request
import urllib.error
from typing import Any, Optional

# ─── Configuration ───────────────────────────────────────────────────────────

API_KEY: str = os.environ.get("TAVILY_API_KEY", "")
BASE_URL: str = "https://api.tavily.com"
REQUEST_TIMEOUT: int = 30
RESEARCH_TIMEOUT: int = 120  # Research can take longer


# ─── HTTP Helper ─────────────────────────────────────────────────────────────

def _api_request(
    endpoint: str,
    payload: dict[str, Any],
    timeout: int = REQUEST_TIMEOUT,
) -> dict[str, Any]:
    """
    Make a POST request to the Tavily API.

    Args:
        endpoint: API endpoint path (e.g., '/search').
        payload: JSON request body.
        timeout: Request timeout in seconds.

    Returns:
        Parsed JSON response.

    Raises:
        SystemExit: On API errors with descriptive messages.
    """
    url = f"{BASE_URL}{endpoint}"
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")[:500]
        if e.code == 401:
            print("Error: Invalid TAVILY_API_KEY. Check your key at https://app.tavily.com", file=sys.stderr)
        elif e.code == 429:
            print("Error: Rate limit exceeded. Try again later.", file=sys.stderr)
        elif e.code == 400:
            # Try to extract error message from JSON response
            try:
                err_data = json.loads(body)
                msg = err_data.get("detail", err_data.get("message", body))
                print(f"Error: Bad request - {msg}", file=sys.stderr)
            except (json.JSONDecodeError, KeyError):
                print(f"Error: Bad request - {body}", file=sys.stderr)
        else:
            print(f"Error: Tavily API returned {e.code}: {body}", file=sys.stderr)
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"Error: Network error - {e.reason}", file=sys.stderr)
        sys.exit(1)
    except TimeoutError:
        print(f"Error: Request timed out after {timeout}s", file=sys.stderr)
        sys.exit(1)


# ─── Output Formatting ──────────────────────────────────────────────────────

def _format_search_results(data: dict[str, Any], as_json: bool = False) -> str:
    """Format search/news/finance results for display."""
    if as_json:
        return json.dumps(data, ensure_ascii=False, indent=2)

    lines: list[str] = []

    # LLM answer
    answer = data.get("answer")
    if answer:
        lines.append(f"Answer: {answer}")
        lines.append("")

    # Images
    images = data.get("images")
    if images:
        lines.append("Images:")
        for img in images:
            if isinstance(img, dict):
                lines.append(f"  - {img.get('url', img)}")
            else:
                lines.append(f"  - {img}")
        lines.append("")

    # Results
    results = data.get("results", [])
    if results:
        lines.append("Results:")
        for i, r in enumerate(results, 1):
            title = r.get("title", "Untitled")
            url = r.get("url", "")
            content = r.get("content", "")
            score = r.get("score")
            published = r.get("published_date", "")

            lines.append(f"  {i}. {title}")
            lines.append(f"     {url}")
            if published:
                lines.append(f"     Published: {published}")
            if score is not None:
                lines.append(f"     Score: {score:.4f}")
            if content:
                # Truncate long content to keep output readable
                snippet = content[:500].strip()
                if len(content) > 500:
                    snippet += "..."
                lines.append(f"     {snippet}")

            # Raw content (if requested)
            raw = r.get("raw_content")
            if raw:
                lines.append(f"     --- Raw Content ---")
                raw_snippet = raw[:1000].strip()
                if len(raw) > 1000:
                    raw_snippet += f"... [{len(raw)} chars total]"
                lines.append(f"     {raw_snippet}")

            lines.append("")
    elif not answer:
        lines.append("No results found.")

    return "\n".join(lines).rstrip()


def _format_extract_results(data: dict[str, Any], as_json: bool = False) -> str:
    """Format extract results for display."""
    if as_json:
        return json.dumps(data, ensure_ascii=False, indent=2)

    lines: list[str] = []
    results = data.get("results", [])

    if not results:
        return "No content extracted."

    for r in results:
        url = r.get("url", "Unknown URL")
        content = r.get("raw_content", "")
        lines.append(f"URL: {url}")
        lines.append("─" * 50)
        if content:
            lines.append(content.strip())
        else:
            lines.append("(No content extracted)")
        lines.append("")

    # Failed URLs
    failed = data.get("failed_results", [])
    if failed:
        lines.append("Failed URLs:")
        for f in failed:
            url = f.get("url", "Unknown")
            error = f.get("error", "Unknown error")
            lines.append(f"  ✗ {url}: {error}")

    return "\n".join(lines).rstrip()


def _format_crawl_results(data: dict[str, Any], as_json: bool = False) -> str:
    """Format crawl results for display."""
    if as_json:
        return json.dumps(data, ensure_ascii=False, indent=2)

    lines: list[str] = []
    results = data.get("results", [])
    base_url = data.get("base_url", "")

    lines.append(f"Crawled {len(results)} pages from {base_url}")
    lines.append("")

    for i, r in enumerate(results, 1):
        url = r.get("url", "Unknown URL")
        content = r.get("raw_content", "")
        lines.append(f"Page {i}: {url}")
        lines.append("─" * 50)
        if content:
            # Truncate very long pages
            snippet = content[:2000].strip()
            if len(content) > 2000:
                snippet += f"\n... [{len(content)} chars total]"
            lines.append(snippet)
        else:
            lines.append("(No content)")
        lines.append("")

    # Failed
    failed = data.get("failed_results", [])
    if failed:
        lines.append("Failed URLs:")
        for f in failed:
            url = f.get("url", "Unknown")
            error = f.get("error", "Unknown error")
            lines.append(f"  ✗ {url}: {error}")

    return "\n".join(lines).rstrip()


def _format_map_results(data: dict[str, Any], url: str, as_json: bool = False) -> str:
    """Format map/sitemap results for display."""
    if as_json:
        return json.dumps(data, ensure_ascii=False, indent=2)

    urls = data.get("results", [])
    lines: list[str] = []
    lines.append(f"Sitemap for {url} ({len(urls)} URLs found):")
    lines.append("")

    for i, u in enumerate(urls, 1):
        if isinstance(u, dict):
            lines.append(f"  {i}. {u.get('url', u)}")
        else:
            lines.append(f"  {i}. {u}")

    if not urls:
        lines.append("  No URLs discovered.")

    return "\n".join(lines).rstrip()


def _format_research_results(data: dict[str, Any], as_json: bool = False) -> str:
    """Format research results for display."""
    if as_json:
        return json.dumps(data, ensure_ascii=False, indent=2)

    lines: list[str] = []

    # Topic
    topic = data.get("topic") or data.get("query", "")
    if topic:
        lines.append(f"Research: {topic}")
        lines.append("")

    # Main content
    content = data.get("content") or data.get("output") or data.get("report", "")
    if content:
        lines.append(content.strip())
    else:
        lines.append("No research output returned.")

    # Sources
    sources = data.get("sources", [])
    if sources:
        lines.append("")
        lines.append("Sources:")
        for i, src in enumerate(sources, 1):
            if isinstance(src, dict):
                url = src.get("url", src.get("link", str(src)))
                title = src.get("title", "")
                if title:
                    lines.append(f"  [{i}] {title}")
                    lines.append(f"      {url}")
                else:
                    lines.append(f"  [{i}] {url}")
            else:
                lines.append(f"  [{i}] {src}")

    return "\n".join(lines).rstrip()


# ─── Commands ────────────────────────────────────────────────────────────────

def cmd_search(args: argparse.Namespace) -> str:
    """Execute search/news/finance command."""
    topic_map = {
        "search": "general",
        "news": "news",
        "finance": "finance",
    }

    payload: dict[str, Any] = {
        "query": args.query,
        "topic": topic_map.get(args.command, "general"),
        "search_depth": args.depth,
        "max_results": args.n,
    }

    if args.answer:
        payload["include_answer"] = True
    if args.raw:
        payload["include_raw_content"] = "markdown"
    if args.images:
        payload["include_images"] = True
    if args.time:
        payload["time_range"] = args.time
    if args.include_domains:
        payload["include_domains"] = [d.strip() for d in args.include_domains.split(",")]
    if args.exclude_domains:
        payload["exclude_domains"] = [d.strip() for d in args.exclude_domains.split(",")]
    if args.country:
        payload["country"] = args.country

    data = _api_request("/search", payload)
    return _format_search_results(data, as_json=args.as_json)


def cmd_extract(args: argparse.Namespace) -> str:
    """Execute extract command."""
    urls = args.urls
    if not urls:
        print("Error: At least one URL is required for extract.", file=sys.stderr)
        sys.exit(1)

    payload: dict[str, Any] = {
        "urls": urls if len(urls) > 1 else urls[0],
    }

    if args.depth and args.depth != "basic":
        payload["extract_depth"] = args.depth
    if hasattr(args, "format_type") and args.format_type:
        payload["format"] = args.format_type
    if hasattr(args, "query") and args.query:
        payload["query"] = args.query

    data = _api_request("/extract", payload)
    return _format_extract_results(data, as_json=args.as_json)


def cmd_crawl(args: argparse.Namespace) -> str:
    """Execute crawl command."""
    payload: dict[str, Any] = {
        "url": args.url,
    }

    if args.max_depth is not None:
        payload["max_depth"] = args.max_depth
    if args.max_breadth is not None:
        payload["max_breadth"] = args.max_breadth
    if args.limit is not None:
        payload["limit"] = args.limit
    if hasattr(args, "instructions") and args.instructions:
        payload["instructions"] = args.instructions
    if hasattr(args, "select_paths") and args.select_paths:
        payload["select_paths"] = [p.strip() for p in args.select_paths.split(",")]
    if hasattr(args, "exclude_paths") and args.exclude_paths:
        payload["exclude_paths"] = [p.strip() for p in args.exclude_paths.split(",")]
    if hasattr(args, "format_type") and args.format_type:
        payload["format"] = args.format_type

    data = _api_request("/crawl", payload, timeout=60)
    return _format_crawl_results(data, as_json=args.as_json)


def cmd_map(args: argparse.Namespace) -> str:
    """Execute map/sitemap command."""
    payload: dict[str, Any] = {
        "url": args.url,
    }

    if args.max_depth is not None:
        payload["max_depth"] = args.max_depth
    if args.max_breadth is not None:
        payload["max_breadth"] = args.max_breadth
    if args.limit is not None:
        payload["limit"] = args.limit

    data = _api_request("/map", payload)
    return _format_map_results(data, args.url, as_json=args.as_json)


def cmd_research(args: argparse.Namespace) -> str:
    """Execute research command."""
    payload: dict[str, Any] = {
        "input": args.query,
    }

    if args.model:
        payload["model"] = args.model

    data = _api_request("/research", payload, timeout=RESEARCH_TIMEOUT)
    return _format_research_results(data, as_json=args.as_json)


# ─── CLI ─────────────────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    """Build the argument parser with all subcommands."""
    parser = argparse.ArgumentParser(
        description="Tavily Search v1.0 - AI-powered web search platform",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s search "latest AI news" --answer
  %(prog)s news "tech industry" --time week
  %(prog)s finance "NVIDIA stock" --depth advanced
  %(prog)s extract "https://example.com/article"
  %(prog)s crawl "https://docs.example.com" --limit 20
  %(prog)s map "https://example.com"
  %(prog)s research "Impact of AI on healthcare"
        """,
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # ── Common search options ──
    def add_search_options(sub: argparse.ArgumentParser) -> None:
        sub.add_argument("query", help="Search query")
        sub.add_argument("--depth", choices=["basic", "advanced"], default="basic",
                         help="Search depth (default: basic; advanced = 2 credits)")
        sub.add_argument("--time", choices=["day", "week", "month", "year", "d", "w", "m", "y"],
                         default=None, help="Time range filter")
        sub.add_argument("-n", type=int, default=5, help="Max results 0-20 (default: 5)")
        sub.add_argument("--answer", action="store_true", help="Include LLM-synthesized answer")
        sub.add_argument("--raw", action="store_true", help="Include raw page content")
        sub.add_argument("--images", action="store_true", help="Include image URLs")
        sub.add_argument("--include-domains", default=None,
                         help="Comma-separated domains to include")
        sub.add_argument("--exclude-domains", default=None,
                         help="Comma-separated domains to exclude")
        sub.add_argument("--country", default=None, help="Country code to boost (e.g., US, IL)")
        sub.add_argument("--json", action="store_true", dest="as_json", help="JSON output")

    # search
    p_search = subparsers.add_parser("search", help="Web search (general)")
    add_search_options(p_search)

    # news
    p_news = subparsers.add_parser("news", help="News search")
    add_search_options(p_news)

    # finance
    p_finance = subparsers.add_parser("finance", help="Finance search")
    add_search_options(p_finance)

    # extract
    p_extract = subparsers.add_parser("extract", help="Extract content from URLs")
    p_extract.add_argument("urls", nargs="+", help="URLs to extract content from")
    p_extract.add_argument("--depth", choices=["basic", "advanced"], default="basic",
                           help="Extraction depth")
    p_extract.add_argument("--format", dest="format_type", choices=["markdown", "text"],
                           default=None, help="Output format (default: markdown)")
    p_extract.add_argument("--query", default=None,
                           help="Query for relevance reranking of chunks")
    p_extract.add_argument("--json", action="store_true", dest="as_json", help="JSON output")

    # crawl
    p_crawl = subparsers.add_parser("crawl", help="Crawl a website")
    p_crawl.add_argument("url", help="Root URL to crawl")
    p_crawl.add_argument("--depth", choices=["basic", "advanced"], default=None,
                         help="Crawl depth")
    p_crawl.add_argument("--max-depth", type=int, default=None, help="Max link depth (default: 2)")
    p_crawl.add_argument("--max-breadth", type=int, default=None,
                         help="Max pages per level (default: 10)")
    p_crawl.add_argument("--limit", type=int, default=None, help="Max total pages (default: 10)")
    p_crawl.add_argument("--instructions", default=None,
                         help="Natural language crawl instructions")
    p_crawl.add_argument("--select-paths", default=None,
                         help="Comma-separated path patterns to include")
    p_crawl.add_argument("--exclude-paths", default=None,
                         help="Comma-separated path patterns to exclude")
    p_crawl.add_argument("--format", dest="format_type", choices=["markdown", "text"],
                         default=None, help="Output format")
    p_crawl.add_argument("--json", action="store_true", dest="as_json", help="JSON output")

    # map
    p_map = subparsers.add_parser("map", help="Discover sitemap URLs")
    p_map.add_argument("url", help="Root URL to map")
    p_map.add_argument("--max-depth", type=int, default=None, help="Max depth (default: 2)")
    p_map.add_argument("--max-breadth", type=int, default=None,
                       help="Max breadth per level (default: 20)")
    p_map.add_argument("--limit", type=int, default=None, help="Max URLs (default: 50)")
    p_map.add_argument("--json", action="store_true", dest="as_json", help="JSON output")

    # research
    p_research = subparsers.add_parser("research", help="Deep AI research")
    p_research.add_argument("query", help="Research question")
    p_research.add_argument("--model", choices=["mini", "pro", "auto"], default=None,
                            help="Research model (default: auto)")
    p_research.add_argument("--json", action="store_true", dest="as_json", help="JSON output")

    return parser


def main() -> None:
    """CLI entry point."""
    parser = build_parser()
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    if not API_KEY:
        print("Error: TAVILY_API_KEY environment variable not set.", file=sys.stderr)
        print("Set it in OpenClaw config or export TAVILY_API_KEY=your_key", file=sys.stderr)
        sys.exit(1)

    try:
        if args.command in ("search", "news", "finance"):
            result = cmd_search(args)
        elif args.command == "extract":
            result = cmd_extract(args)
        elif args.command == "crawl":
            result = cmd_crawl(args)
        elif args.command == "map":
            result = cmd_map(args)
        elif args.command == "research":
            result = cmd_research(args)
        else:
            parser.print_help()
            sys.exit(1)

        print(result)
    except SystemExit:
        raise
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()