from __future__ import annotations

import asyncio
import re
from collections import defaultdict
from typing import Optional
from urllib.parse import urlparse

import httpx
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

router = APIRouter(prefix="/dataset")

# ── Known AI crawler bot names to check in robots.txt ─────────────────────────
# Sources: ai-robots-txt/ai.robots.txt (github.com/ai-robots-txt/ai.robots.txt, 272 entries)
# and knownagents.com. Grouped by operator for maintainability.

AI_BOTS = [
    # OpenAI
    "GPTBot", "ChatGPT-User", "OAI-SearchBot",
    # Anthropic
    "ClaudeBot", "Claude-Web", "Claude-SearchBot", "anthropic-ai",
    # Google
    "Google-Extended", "Googlebot", "Googlebot-Image", "Googlebot-Video",
    "Googlebot-News", "APIs-Google", "AdsBot-Google", "Google-InspectionTool",
    "GoogleOther", "GoogleOther-Image", "GoogleOther-Video",
    # Meta
    "FacebookBot", "FacebookExternalHit", "Meta-ExternalAgent",
    "Meta-ExternalFetcher", "facebookexternalhit",
    # Apple
    "Applebot", "Applebot-Extended",
    # Microsoft / Bing
    "bingbot", "BingPreview", "msnbot", "MicrosoftPreview",
    "adidxbot", "MSNBot-Media",
    # Amazon
    "Amazonbot",
    # ByteDance / TikTok
    "Bytespider", "TikTokBot",
    # Common Crawl (used by many LLM datasets)
    "CCBot",
    # Allen Institute for AI
    "AI2Bot", "ai2-agent",
    # Cohere
    "cohere-ai",
    # Perplexity
    "PerplexityBot",
    # Diffbot
    "Diffbot",
    # LAION
    "Omgilibot", "omgili",
    # You.com
    "YouBot",
    # Brave
    "Brave",
    # Tavily
    "TavilyBot",
    # DeepSeek
    "DeepSeekBot",
    # Mistral / Le Chat
    "MistralBot",
    # xAI / Grok
    "xAI-Bot",
    # Hugging Face
    "HuggingFaceBot",
    # Scrapers / data providers used for AI training
    "Scrapy", "ia_archiver", "archive.org_bot",
    "SemrushBot", "AhrefsBot", "MJ12bot", "DotBot", "SeznamBot",
    "PetalBot", "BLEXBot", "YandexBot", "YandexImages",
    "Baiduspider", "Sogou", "360Spider",
    "AddSearchBot", "Applebot",
    # Dataset crawlers
    "DataForSeoBot", "DataBot", "Scrapy",
    "newspaper", "news-please",
    # Search AI crawlers
    "PerplexityBot", "PhindBot", "kagi-bot",
    "Qwant", "DuckAssistBot",
    # Agent frameworks
    "Crawl4AI", "firecrawl", "Apify",
    "browserless", "playwright",
    # Image / multimodal scrapers
    "ImgProxy", "CCBot",
    # Academic / research
    "Semanticscholar", "PubMedBot", "ResearchBot",
    # Generic AI training crawlers
    "img2dataset", "laion-prepro",
    "DataProvider", "TrainingDataBot",
]

# Richer metadata for bots where we have operator + purpose info (subset)
AI_BOT_META: dict[str, dict] = {
    "GPTBot":           {"operator": "OpenAI",      "purpose": "LLM training",   "respects_robots": True},
    "ChatGPT-User":     {"operator": "OpenAI",      "purpose": "AI assistant",   "respects_robots": True},
    "OAI-SearchBot":    {"operator": "OpenAI",      "purpose": "search indexing","respects_robots": True},
    "ClaudeBot":        {"operator": "Anthropic",   "purpose": "LLM training",   "respects_robots": True},
    "Claude-Web":       {"operator": "Anthropic",   "purpose": "AI assistant",   "respects_robots": True},
    "anthropic-ai":     {"operator": "Anthropic",   "purpose": "LLM training",   "respects_robots": True},
    "Google-Extended":  {"operator": "Google",      "purpose": "AI training",    "respects_robots": True},
    "CCBot":            {"operator": "Common Crawl","purpose": "LLM datasets",   "respects_robots": True},
    "FacebookBot":      {"operator": "Meta",        "purpose": "AI training",    "respects_robots": True},
    "Bytespider":       {"operator": "ByteDance",   "purpose": "LLM training",   "respects_robots": False},
    "Amazonbot":        {"operator": "Amazon",      "purpose": "AI training",    "respects_robots": True},
    "PerplexityBot":    {"operator": "Perplexity",  "purpose": "AI search",      "respects_robots": True},
    "Applebot-Extended":{"operator": "Apple",       "purpose": "AI training",    "respects_robots": True},
    "AI2Bot":           {"operator": "AI2",         "purpose": "research",       "respects_robots": True},
    "cohere-ai":        {"operator": "Cohere",      "purpose": "LLM training",   "respects_robots": True},
    "Diffbot":          {"operator": "Diffbot",     "purpose": "data extraction","respects_robots": False},
    "TavilyBot":        {"operator": "Tavily",      "purpose": "AI search",      "respects_robots": True},
    "DeepSeekBot":      {"operator": "DeepSeek",    "purpose": "LLM training",   "respects_robots": True},
    "YouBot":           {"operator": "You.com",     "purpose": "AI search",      "respects_robots": True},
}

# ── Spawning AI API ────────────────────────────────────────────────────────────

SPAWNING_URL_API   = "https://opts.api.spawning.ai/opts-api/urls"
SPAWNING_HF_API    = "https://opts.api.spawning.ai/opts-api/dataset"
SPAWNING_TIMEOUT   = 10.0
SPAWNING_BATCH_SIZE = 100


async def _check_spawning_urls(urls: list[str]) -> dict[str, bool]:
    """Batch-query Spawning AI for opt-out status. Returns {url: opted_out}."""
    results: dict[str, bool] = {}
    if not urls:
        return results

    async with httpx.AsyncClient(timeout=SPAWNING_TIMEOUT) as client:
        for i in range(0, len(urls), SPAWNING_BATCH_SIZE):
            batch = urls[i : i + SPAWNING_BATCH_SIZE]
            try:
                resp = await client.post(
                    SPAWNING_URL_API,
                    json={"urls": batch},
                    headers={"Content-Type": "application/json"},
                )
                if resp.status_code == 200:
                    data = resp.json()
                    # Response: {"urls": [{"url": "...", "optOut": true/false}, ...]}
                    for item in data.get("urls", []):
                        results[item["url"]] = bool(item.get("optOut", False))
                else:
                    for u in batch:
                        results[u] = False
            except Exception:
                for u in batch:
                    results[u] = False

    return results


async def _check_spawning_hf(dataset_id: str) -> dict:
    """Check a HuggingFace dataset ID against Spawning's dataset opt-out registry."""
    async with httpx.AsyncClient(timeout=SPAWNING_TIMEOUT) as client:
        try:
            resp = await client.post(
                SPAWNING_HF_API,
                json={"dataset": dataset_id},
                headers={"Content-Type": "application/json"},
            )
            if resp.status_code == 200:
                return resp.json()
        except Exception:
            pass
    return {}


# ── robots.txt parsing ─────────────────────────────────────────────────────────

_ROBOTS_CACHE: dict[str, dict] = {}


def _parse_robots(content: str) -> dict[str, list[str]]:
    """Parse robots.txt into {user_agent: [disallow_paths]}."""
    rules: dict[str, list[str]] = defaultdict(list)
    current_agents: list[str] = []
    for raw_line in content.splitlines():
        line = raw_line.split("#")[0].strip()
        if not line:
            current_agents = []
            continue
        if line.lower().startswith("user-agent:"):
            agent = line.split(":", 1)[1].strip()
            current_agents.append(agent)
        elif line.lower().startswith("disallow:") and current_agents:
            path = line.split(":", 1)[1].strip()
            if path:
                for agent in current_agents:
                    rules[agent].append(path)
    return dict(rules)


def _robots_blocks_ai(rules: dict[str, list[str]]) -> list[str]:
    """Return list of AI bot names that are disallowed from '/' (root)."""
    blocked = []
    all_disallows = rules.get("*", [])
    star_blocks_root = "/" in all_disallows or "/*" in all_disallows

    for bot in AI_BOTS:
        bot_disallows = rules.get(bot, [])
        if "/" in bot_disallows or "/*" in bot_disallows:
            blocked.append(bot)
        elif star_blocks_root and bot not in rules:
            blocked.append(f"* (via {bot})")
    return blocked


async def _fetch_robots(domain: str) -> dict:
    """Fetch and parse robots.txt for a domain. Returns summary dict."""
    if domain in _ROBOTS_CACHE:
        return _ROBOTS_CACHE[domain]

    result = {"domain": domain, "fetched": False, "blocked_bots": [], "error": None}
    url = f"https://{domain}/robots.txt"

    async with httpx.AsyncClient(timeout=8.0, follow_redirects=True) as client:
        try:
            resp = await client.get(url, headers={"User-Agent": "Mozilla/5.0"})
            if resp.status_code == 200:
                rules = _parse_robots(resp.text)
                blocked = _robots_blocks_ai(rules)
                result["fetched"] = True
                result["blocked_bots"] = blocked
                result["rules_found"] = len(rules)
            else:
                result["error"] = f"HTTP {resp.status_code}"
        except Exception as e:
            result["error"] = str(e)[:80]

    _ROBOTS_CACHE[domain] = result
    return result


def _extract_domains(urls: list[str]) -> list[str]:
    seen = set()
    domains = []
    for url in urls:
        try:
            parsed = urlparse(url if url.startswith("http") else "https://" + url)
            host = (parsed.hostname or "").lower().lstrip("www.")
            if host and host not in seen:
                seen.add(host)
                domains.append(host)
        except Exception:
            pass
    return domains


# ── Overall status helpers ─────────────────────────────────────────────────────

def _overall_status(
    opted_out_count: int,
    robots_blocked_count: int,
    hf_opted_out: bool,
) -> str:
    if hf_opted_out or opted_out_count > 0:
        return "opted_out"
    if robots_blocked_count > 0:
        return "restricted"
    return "clear"


# ── Request model ──────────────────────────────────────────────────────────────

class OptOutRequest(BaseModel):
    rows: list[dict]
    url_columns: list[str] = []
    hf_dataset_id: Optional[str] = None
    max_urls: int = 200


# ── Route ──────────────────────────────────────────────────────────────────────

@router.post("/optout")
async def optout_check(req: OptOutRequest):
    rows = req.rows
    if not rows and not req.hf_dataset_id:
        raise HTTPException(status_code=400, detail="Provide rows with url_columns or a hf_dataset_id")

    # ── 1. Collect URLs from row data ─────────────────────────────────────────
    seen_urls: set[str] = set()
    url_list: list[str] = []
    for row in rows[:req.max_urls]:
        for col in req.url_columns:
            val = row.get(col)
            if val and isinstance(val, str) and val.startswith("http") and val not in seen_urls:
                seen_urls.add(val)
                url_list.append(val)

    domains = _extract_domains(url_list)

    # ── 2. Run Spawning AI URL check + robots.txt checks in parallel ──────────
    spawning_task = asyncio.create_task(_check_spawning_urls(url_list))
    robots_tasks  = [asyncio.create_task(_fetch_robots(d)) for d in domains[:30]]
    hf_task       = asyncio.create_task(
        _check_spawning_hf(req.hf_dataset_id) if req.hf_dataset_id else asyncio.coroutine(lambda: {})()
    )

    spawning_results, *robots_results_list, hf_result = await asyncio.gather(
        spawning_task, *robots_tasks, hf_task
    )

    # ── 3. Tally Spawning URL results ─────────────────────────────────────────
    opted_out_urls = [u for u, v in spawning_results.items() if v]
    not_opted_urls = [u for u, v in spawning_results.items() if not v]

    # Group opted-out URLs by domain
    domain_optout: dict[str, list[str]] = defaultdict(list)
    for u in opted_out_urls:
        try:
            host = urlparse(u).hostname or ""
            domain_optout[host.lstrip("www.")].append(u)
        except Exception:
            pass

    # ── 4. Tally robots.txt results ───────────────────────────────────────────
    robots_results = robots_results_list if isinstance(robots_results_list, list) else []
    domains_with_blocks = [r for r in robots_results if r.get("blocked_bots")]
    domains_clear       = [r for r in robots_results if r.get("fetched") and not r.get("blocked_bots")]

    # ── 5. HF dataset opt-out ─────────────────────────────────────────────────
    hf_opted_out = bool(hf_result.get("optOut", False)) if hf_result else False

    # ── 6. Overall status ─────────────────────────────────────────────────────
    status = _overall_status(
        opted_out_count=len(opted_out_urls),
        robots_blocked_count=len(domains_with_blocks),
        hf_opted_out=hf_opted_out,
    )

    return {
        "status": status,
        "urls_checked": len(url_list),
        "domains_checked": len(domains),

        "spawning": {
            "opted_out_count": len(opted_out_urls),
            "clear_count": len(not_opted_urls),
            "opted_out_urls": opted_out_urls[:50],
            "domain_breakdown": [
                {"domain": d, "url_count": len(urls)}
                for d, urls in sorted(domain_optout.items(), key=lambda x: -len(x[1]))
            ],
        },

        "robots": {
            "domains_blocked": len(domains_with_blocks),
            "domains_clear": len(domains_clear),
            "blocked": [
                {
                    "domain": r["domain"],
                    "blocked_bots": r["blocked_bots"],
                }
                for r in domains_with_blocks
            ],
        },

        "hf_dataset": {
            "dataset_id": req.hf_dataset_id,
            "opted_out": hf_opted_out,
            "raw": hf_result if hf_result else None,
        } if req.hf_dataset_id else None,
    }
