""" 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")