scraper/manager.py

"""
Download manager: orchestrates discovery → downloading → writing with full retry/state logic.
"""
from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

import aiofiles

from .engine import (
    Chapter, ChapterStatus, NovelJob,
    create_session, smart_fetch, FetchMode,
    _is_cloudflare, _is_login_page, _is_locked_chapter,
)
from .discovery import discover_chapters
from .extractor import extract_chapter, is_locked_chapter

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

@dataclass
class DownloadConfig:
    mode: FetchMode = FetchMode.BALANCED
    max_concurrent: int = 3
    rate_limit: float = 1.0       # seconds between requests
    max_retries: int = 30
    retry_base_delay: float = 2.0
    timeout: int = 30
    cookies: dict = field(default_factory=dict)
    output_dir: Path = field(default_factory=lambda: Path("output"))
    save_state: bool = True

    def to_dict(self) -> dict:
        return {
            "mode": self.mode.value,
            "max_concurrent": self.max_concurrent,
            "rate_limit": self.rate_limit,
            "max_retries": self.max_retries,
            "timeout": self.timeout,
        }


# ---------------------------------------------------------------------------
# State persistence
# ---------------------------------------------------------------------------

def _state_path(job: NovelJob) -> Path:
    safe = "".join(c if c.isalnum() else "_" for c in job.url)[:80]
    return job.output_dir / f".state_{safe}.json"


def save_state(job: NovelJob) -> None:
    try:
        path = _state_path(job)
        path.parent.mkdir(parents=True, exist_ok=True)
        data = {
            "url": job.url,
            "title": job.title,
            "status": job.status,
            "chapters": [c.to_dict() for c in job.chapters],
            "start_chapter": job.start_chapter,
            "end_chapter": job.end_chapter,
            "specific_chapters": job.specific_chapters,
        }
        path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    except Exception as exc:
        logger.warning(f"State save failed: {exc}")


def load_state(url: str, output_dir: Path) -> Optional[NovelJob]:
    """Try to restore a previous job state."""
    safe = "".join(c if c.isalnum() else "_" for c in url)[:80]
    path = output_dir / f".state_{safe}.json"
    if not path.exists():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        job = NovelJob(
            url=data["url"],
            title=data.get("title", ""),
            status="queued",
            output_dir=output_dir,
            start_chapter=data.get("start_chapter", 1),
            end_chapter=data.get("end_chapter", 0),
            specific_chapters=data.get("specific_chapters", []),
        )
        for cd in data.get("chapters", []):
            ch = Chapter(
                index=cd["index"],
                title=cd["title"],
                url=cd["url"],
                status=ChapterStatus(cd.get("status", "pending")),
                retries=cd.get("retries", 0),
                error=cd.get("error", ""),
                word_count=cd.get("word_count", 0),
            )
            # Reset failed/downloading to pending for retry
            if ch.status in (ChapterStatus.DOWNLOADING, ChapterStatus.FAILED):
                if ch.retries >= 30:
                    ch.status = ChapterStatus.FAILED
                else:
                    ch.status = ChapterStatus.PENDING
            job.chapters.append(ch)
        logger.info(f"Restored state: {len(job.chapters)} chapters, {job.title}")
        return job
    except Exception as exc:
        logger.warning(f"State load failed: {exc}")
        return None


# ---------------------------------------------------------------------------
# Writer (stream chapters to disk, never hold full novel in RAM)
# ---------------------------------------------------------------------------

class NovelWriter:
    def __init__(self, path: Path):
        self.path = path
        self._file = None

    async def __aenter__(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._file = await aiofiles.open(self.path, "a", encoding="utf-8")
        return self

    async def __aexit__(self, *_):
        if self._file:
            await self._file.close()

    async def write_chapter(self, title: str, content: str) -> None:
        text = f"{title}\n\n{content}\n\n\n"
        await self._file.write(text)
        await self._file.flush()


# ---------------------------------------------------------------------------
# Download manager
# ---------------------------------------------------------------------------

class DownloadManager:
    def __init__(
        self,
        job: NovelJob,
        config: DownloadConfig,
        log_cb: Callable[[str], None] | None = None,
        progress_cb: Callable[[NovelJob], None] | None = None,
    ):
        self.job = job
        self.config = config
        self.log_cb = log_cb
        self.progress_cb = progress_cb
        self._semaphore: Optional[asyncio.Semaphore] = None
        self._rate_limiter: asyncio.Lock = asyncio.Lock()
        self._last_request_time: float = 0.0
        self._paused = False
        self._cancelled = False
        self._failed_later: list[Chapter] = []  # chapters to retry at the end

    def log(self, msg: str) -> None:
        logger.info(msg)
        if self.log_cb:
            self.log_cb(msg)

    def _notify_progress(self) -> None:
        if self.progress_cb:
            self.progress_cb(self.job)

    async def _rate_wait(self) -> None:
        async with self._rate_limiter:
            now = time.time()
            elapsed = now - self._last_request_time
            wait = self.config.rate_limit - elapsed
            if wait > 0:
                await asyncio.sleep(wait)
            self._last_request_time = time.time()

    async def _download_chapter(
        self,
        session,
        chapter: Chapter,
        writer: NovelWriter,
    ) -> None:
        """Download a single chapter with retry logic."""
        async with self._semaphore:
            while self._paused and not self._cancelled:
                await asyncio.sleep(0.5)
            if self._cancelled:
                return

            chapter.status = ChapterStatus.DOWNLOADING
            self._notify_progress()

            for attempt in range(self.config.max_retries):
                if self._cancelled:
                    return
                try:
                    await self._rate_wait()
                    html, status = await smart_fetch(
                        session,
                        chapter.url,
                        mode=self.config.mode,
                        cookies=self.config.cookies,
                        log_cb=self.log_cb,
                    )

                    if _is_cloudflare(html, status):
                        self.log(f"⚠️  CF wall on {chapter.title}, backing off...")
                        await asyncio.sleep(10 * (attempt + 1))
                        chapter.retries += 1
                        continue

                    if _is_login_page(html):
                        chapter.status = ChapterStatus.SKIPPED
                        chapter.error = "Login required"
                        self.log(f"🔒 Skipped (login required): {chapter.title}")
                        return

                    if is_locked_chapter(html) or _is_locked_chapter(html):
                        chapter.status = ChapterStatus.LOCKED
                        chapter.error = "Locked/premium chapter"
                        self.log(f"💰 Skipped (locked): {chapter.title}")
                        return

                    title, content, paragraphs = extract_chapter(html, fallback_title=chapter.title)
                    if not content or len(content) < 50:
                        raise ValueError(f"Extracted content too short ({len(content)} chars)")

                    # Update title if extraction improved it
                    if title and title != chapter.title and len(title) < 200:
                        chapter.title = title

                    chapter.word_count = len(content.split())
                    await writer.write_chapter(chapter.title, content)

                    chapter.status = ChapterStatus.DONE
                    chapter.content = ""  # Don't keep in RAM
                    self.log(f"✅ {chapter.title} ({chapter.word_count} words)")
                    self._notify_progress()
                    return

                except Exception as exc:
                    chapter.retries += 1
                    delay = self.config.retry_base_delay * (2 ** min(attempt, 6))
                    self.log(f"⚠️  Retry {attempt+1}/{self.config.max_retries} for {chapter.title}: {exc}")

                    if attempt < 3:
                        # Fast retry for first few attempts
                        await asyncio.sleep(delay)
                    else:
                        # Defer for later retry
                        self._failed_later.append(chapter)
                        self.log(f"⏳ Deferred {chapter.title} for later retry")
                        return

            chapter.status = ChapterStatus.FAILED
            chapter.error = f"Failed after {self.config.max_retries} retries"
            self.log(f"❌ Permanently failed: {chapter.title}")
            self._notify_progress()

    async def run(self) -> Path:
        """Run the full download pipeline. Returns path to output TXT."""
        self._semaphore = asyncio.Semaphore(self.config.max_concurrent)
        self.job.status = "running"
        self.job.started_at = time.time()
        self.job.output_dir.mkdir(parents=True, exist_ok=True)

        session = await create_session(self.config.cookies, self.config.timeout)

        try:
            # Discovery phase
            if not self.job.chapters:
                self.log(f"🔍 Starting chapter discovery for {self.job.url}")
                chapters = await discover_chapters(
                    session, self.job.url,
                    mode=self.config.mode,
                    cookies=self.config.cookies,
                    log_cb=self.log_cb,
                )
                if not chapters:
                    raise RuntimeError("No chapters discovered. Try Compatibility mode or check the URL.")
                self.job.chapters = chapters
                self.log(f"📚 Found {len(chapters)} chapters")

            # Detect novel title from first chapter URL if not set
            if not self.job.title:
                from urllib.parse import urlparse
                parts = urlparse(self.job.url).path.strip("/").split("/")
                self.job.title = parts[-1].replace("-", " ").replace("_", " ").title() if parts else "Novel"

            # Filter chapters by range
            chapters_to_download = [
                c for c in self.job.chapters
                if self.job.chapter_in_range(c) and c.status == ChapterStatus.PENDING
            ]

            self.log(f"⬇️  Downloading {len(chapters_to_download)} chapters...")

            safe_title = "".join(c if c.isalnum() or c in " -_" else "_" for c in self.job.title)[:60]
            out_path = self.job.output_dir / f"{safe_title}.txt"

            # Clear existing output if fresh start
            if out_path.exists() and not any(
                c.status == ChapterStatus.DONE for c in self.job.chapters
            ):
                out_path.unlink()

            async with NovelWriter(out_path) as writer:
                tasks = [
                    self._download_chapter(session, ch, writer)
                    for ch in chapters_to_download
                ]
                await asyncio.gather(*tasks, return_exceptions=True)

                # Retry deferred chapters
                if self._failed_later:
                    self.log(f"🔄 Retrying {len(self._failed_later)} deferred chapters...")
                    retry_tasks = []
                    seen_urls = set()
                    for ch in self._failed_later:
                        if ch.url not in seen_urls:
                            seen_urls.add(ch.url)
                            ch.status = ChapterStatus.PENDING
                            retry_tasks.append(self._download_chapter(session, ch, writer))
                    await asyncio.gather(*retry_tasks, return_exceptions=True)

            self.job.status = "done"
            self.job.finished_at = time.time()

            # Write failed chapters report
            failed = [c for c in self.job.chapters if c.status == ChapterStatus.FAILED]
            failed_path = self.job.output_dir / "failed_chapters.txt"
            if failed:
                lines = [f"Failed Chapters Report\n{'='*40}\n"]
                for c in failed:
                    lines.append(f"Title: {c.title}\nURL: {c.url}\nError: {c.error}\n")
                failed_path.write_text("\n".join(lines), encoding="utf-8")
                self.log(f"⚠️  {len(failed)} chapters failed — see failed_chapters.txt")
            elif failed_path.exists():
                failed_path.unlink()

            if self.config.save_state:
                save_state(self.job)

            self.log(f"🎉 Done! {self.job.done_count}/{self.job.total} chapters → {out_path}")
            self._notify_progress()
            return out_path

        finally:
            await session.close()

    def pause(self) -> None:
        self._paused = True
        self.job.status = "paused"
        self.log("⏸️  Paused")

    def resume(self) -> None:
        self._paused = False
        self.job.status = "running"
        self.log("▶️  Resumed")

    def cancel(self) -> None:
        self._cancelled = True
        self.job.status = "cancelled"
        self.log("🚫 Cancelled")