scraper/discovery.py

"""
Chapter link discovery: handles paginated lists, AJAX, JSON-in-scripts, infinite scroll, etc.
"""
from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass
from typing import Callable, Optional
from urllib.parse import urljoin, urlparse, urlencode, parse_qs, urlunparse

from bs4 import BeautifulSoup

from .engine import Chapter, smart_fetch, FetchMode, _is_cloudflare
import aiohttp

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Chapter sorting helpers
# ---------------------------------------------------------------------------

_CH_NUM_RE = re.compile(
    r"(?:chapter|ch\.?|ep\.?|episode|part|vol\.?|volume)?\s*"
    r"(\d+(?:\.\d+)?)",
    re.IGNORECASE,
)
_SIDE_STORY_RE = re.compile(r"side\s*story|extra|bonus|omake|special|ss\s*\d", re.IGNORECASE)


def _chapter_sort_key(ch: "Chapter") -> tuple[float, float, str]:
    """Return (is_side_story, number, title) for natural sorting."""
    is_side = 1.0 if _SIDE_STORY_RE.search(ch.title) else 0.0
    m = _CH_NUM_RE.search(ch.title)
    num = float(m.group(1)) if m else 0.0
    return (is_side, num, ch.title.lower())


def sort_chapters(chapters: list[Chapter]) -> list[Chapter]:
    sorted_chs = sorted(chapters, key=_chapter_sort_key)
    for i, ch in enumerate(sorted_chs, start=1):
        ch.index = i
    return sorted_chs


# ---------------------------------------------------------------------------
# Link scoring helpers
# ---------------------------------------------------------------------------

_CHAPTER_URL_PATTERNS = [
    re.compile(p, re.IGNORECASE) for p in [
        r"/chapter[-_]?\d",
        r"/ch[-_]?\d",
        r"/episode[-_]?\d",
        r"/ep[-_]?\d",
        r"/chapter/",
        r"/chapters/\d",
        r"/\d+/?$",
        r"chapter=\d",
        r"ep=\d",
    ]
]

_NAV_CLASSES = re.compile(
    r"nav|navigation|menu|header|footer|sidebar|widget|ad|comment|related|recommend|share|social",
    re.IGNORECASE,
)

_CHAPTER_CLASSES = re.compile(
    r"chapter|episode|listing|list|toc|content|volume|chapters|chap-list|ep-list",
    re.IGNORECASE,
)


def _score_link(a_tag, base_url: str) -> float:
    """Score an <a> tag on likelihood of being a chapter link."""
    href = a_tag.get("href", "")
    text = a_tag.get_text(strip=True)
    if not href or not text:
        return 0.0

    score = 0.0

    # URL pattern match
    for pat in _CHAPTER_URL_PATTERNS:
        if pat.search(href):
            score += 3.0
            break

    # Text contains number
    if re.search(r"\d+", text):
        score += 2.0

    # Text contains "chapter"/"episode"/etc
    if re.search(r"chapter|episode|ep\.|ch\.", text, re.IGNORECASE):
        score += 2.0

    # Parent container looks like a chapter list
    for parent in a_tag.parents:
        cls = " ".join(parent.get("class", []))
        id_ = parent.get("id", "")
        if _CHAPTER_CLASSES.search(cls + " " + id_):
            score += 2.0
            break
        if _NAV_CLASSES.search(cls + " " + id_):
            score -= 2.0
            break

    # Same domain
    parsed_base = urlparse(base_url)
    parsed_href = urlparse(urljoin(base_url, href))
    if parsed_href.netloc and parsed_href.netloc != parsed_base.netloc:
        score -= 5.0

    return score


# ---------------------------------------------------------------------------
# JSON-in-script extraction
# ---------------------------------------------------------------------------

def _extract_urls_from_scripts(soup: BeautifulSoup, base_url: str) -> list[str]:
    """Pull chapter URLs from inline JSON blobs in <script> tags."""
    found: list[str] = []
    base_domain = urlparse(base_url).netloc

    for script in soup.find_all("script"):
        src = script.string or ""
        # Look for arrays or objects containing URLs
        urls = re.findall(r'["\'](https?://[^"\']+?/[^"\']*chapter[^"\']*)["\']', src, re.IGNORECASE)
        urls += re.findall(r'["\'](/[^"\']*chapter[^"\']*)["\']', src, re.IGNORECASE)
        for u in urls:
            full = urljoin(base_url, u)
            if urlparse(full).netloc == base_domain:
                found.append(full)

        # Look for JSON arrays with href/url keys
        try:
            json_blobs = re.findall(r'\{[^{}]*"(?:href|url|link|chapterUrl)"[^{}]*\}', src)
            for blob in json_blobs:
                try:
                    data = json.loads(blob)
                    for key in ("href", "url", "link", "chapterUrl"):
                        if key in data:
                            full = urljoin(base_url, data[key])
                            if urlparse(full).netloc == base_domain:
                                found.append(full)
                except Exception:
                    pass
        except Exception:
            pass

    return list(dict.fromkeys(found))  # deduplicate preserving order


# ---------------------------------------------------------------------------
# Pagination helpers
# ---------------------------------------------------------------------------

def _build_page_urls(base_url: str, max_pages: int = 30) -> list[str]:
    """Generate candidate paginated URLs."""
    candidates: list[str] = []
    parsed = urlparse(base_url)

    for i in range(2, max_pages + 1):
        # ?page=N
        qs = parse_qs(parsed.query)
        qs["page"] = [str(i)]
        params = "&".join(f"{k}={v[0]}" for k, v in qs.items())
        cand = urlunparse(parsed._replace(query=params))
        candidates.append(cand)

        # /page/N
        path_paged = parsed.path.rstrip("/") + f"/page/{i}"
        candidates.append(urlunparse(parsed._replace(path=path_paged, query="")))

    return candidates


_TOC_SUFFIXES = [
    "",
    "/chapter",
    "/chapters",
    "/chapters/",
    "/contents",
    "/episodes",
    "/all-chapters",
    "/table-of-contents",
    "/toc",
]


def _toc_candidates(url: str) -> list[str]:
    parsed = urlparse(url)
    base = f"{parsed.scheme}://{parsed.netloc}{parsed.path.rstrip('/')}"
    return [base + s for s in _TOC_SUFFIXES]


# ---------------------------------------------------------------------------
# Main discoverer
# ---------------------------------------------------------------------------

async def discover_chapters(
    session: aiohttp.ClientSession,
    novel_url: str,
    mode: FetchMode = FetchMode.BALANCED,
    cookies: dict | None = None,
    log_cb: Callable[[str], None] | None = None,
) -> list[Chapter]:
    """Discover all chapter links for a novel. Returns sorted Chapter list."""

    def log(msg: str):
        if log_cb:
            log_cb(msg)
        logger.info(msg)

    all_links: dict[str, str] = {}  # url -> title

    async def _scrape_page(url: str) -> tuple[list[tuple[str, str]], bool]:
        """Return [(url, title), ...] and whether any chapters were found."""
        try:
            html, status = await smart_fetch(session, url, mode=mode, cookies=cookies, log_cb=log_cb)
        except Exception as exc:
            log(f"⚠️  Cannot fetch {url}: {exc}")
            return [], False

        if _is_cloudflare(html, status):
            log(f"🛑 Cloudflare wall on {url}")
            return [], False

        soup = BeautifulSoup(html, "lxml")

        # Remove obvious noise
        for tag in soup.find_all(["script", "style", "noscript", "iframe"]):
            tag.decompose()

        results: list[tuple[str, str]] = []

        # Score all <a> tags
        scored: list[tuple[float, str, str]] = []
        for a in soup.find_all("a", href=True):
            score = _score_link(a, url)
            if score > 1.0:
                href = urljoin(url, a["href"])
                title = a.get_text(strip=True)
                scored.append((score, href, title))

        # Also pull from scripts
        script_urls = _extract_urls_from_scripts(soup, url)
        for su in script_urls:
            scored.append((2.0, su, ""))

        # Deduplicate by URL, keep highest score title
        seen: dict[str, tuple[float, str]] = {}
        for score, href, title in scored:
            if href not in seen or score > seen[href][0]:
                seen[href] = (score, title)

        results = [(href, title) for href, (_, title) in seen.items()]
        return results, bool(results)

    log(f"🔍 Discovering chapters for: {novel_url}")

    # Try TOC variants
    for toc_url in _toc_candidates(novel_url):
        log(f"  Trying TOC URL: {toc_url}")
        links, found = await _scrape_page(toc_url)
        if found:
            log(f"  ✅ Found {len(links)} links at {toc_url}")
            for href, title in links:
                if href not in all_links:
                    all_links[href] = title
            break

    # If we have few links, try pagination
    if len(all_links) < 5:
        log("📄 Trying paginated chapter lists...")
        for page_url in _build_page_urls(novel_url):
            links, found = await _scrape_page(page_url)
            if not found:
                break
            before = len(all_links)
            for href, title in links:
                if href not in all_links:
                    all_links[href] = title
            if len(all_links) == before:
                break  # No new links, stop paginating
            log(f"  Page result: {len(all_links)} total links")

    # Final dedup and filter
    base_domain = urlparse(novel_url).netloc
    clean: dict[str, str] = {}
    for href, title in all_links.items():
        parsed = urlparse(href)
        if parsed.netloc == base_domain or not parsed.netloc:
            # Skip obvious non-chapter URLs
            if any(bad in href.lower() for bad in [
                "/tag/", "/category/", "/author/", "/search", "/login", "/register",
                "/contact", "/about", "/privacy", "/sitemap", "/feed", "#"
            ]):
                continue
            clean[href] = title

    log(f"✅ Discovered {len(clean)} chapter links total")

    chapters: list[Chapter] = []
    for i, (href, title) in enumerate(clean.items(), start=1):
        if not title:
            title = f"Chapter {i}"
        chapters.append(Chapter(index=i, title=title, url=href))

    return sort_chapters(chapters)