#!/usr/bin/env python3
"""
archive.py - Site Archiver for Sunset Sites

Thin wrapper around neopig that packages crawl results into a distributable
tar.gz archive. Uses neopig for the actual crawling.

Output structure:
    {domain}-{date}/
        html/                   # Original HTML pages
        markdown/               # Converted markdown (optional)
        vault/                  # Media vault (9-layer hash paths, neopig-compatible)
        neopig.db               # Full neopig database
        neopig/                 # Embedded neopig server
        metadata.json           # Crawl metadata

Usage:
    python archive.py https://discourse-urho3d.github.io/
    python archive.py https://example.com --no-screenshot
"""

import argparse
import asyncio
import json
import logging
import os
import re
import shutil
import signal
import subprocess
import sys
import tarfile
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Any, List, Optional
from miniuri import Uri

import aiofiles
import aiofiles.os
from bs4 import BeautifulSoup
from tqdm import tqdm

from neopig import NeoPig, setup_logging, rotate_state_file, get_state_file_path
from async_web_fetcher import CrawlMode
from screenshot import ScreenshotConfig
from filevault import hash_to_path

# Optional markdown conversion
try:
    import html2text
    HAS_HTML2TEXT = True
except ImportError:
    HAS_HTML2TEXT = False

logger = logging.getLogger(__name__)


def sanitize_filename(name: str) -> str:
    """Sanitize a string for use as filename. Uses - as separator."""
    name = re.sub(r'[<>:"/\\|?*.]', '-', name)
    name = re.sub(r'-+', '-', name)  # collapse multiple dashes
    name = name.strip('- ')
    return name[:200] if name else 'unnamed'


def url_to_path(url: str) -> str:
    """Convert URL to filesystem path."""
    parsed = Uri(url)
    path = parsed.path.strip('/')
    if not path:
        return 'index.html'
    if path.endswith('.html') or path.endswith('.htm'):
        return path
    if '.' in path.split('/')[-1]:
        return path
    return f"{path}/index.html"


def html_to_markdown(html: str, base_url: str = '', trim_wrapper: bool = False) -> str:
    """Convert HTML to markdown."""
    if not HAS_HTML2TEXT:
        return html
    # Optionally strip nav/header/footer/logo before conversion
    if trim_wrapper:
        from neopig import trim_html_wrapper
        html = trim_html_wrapper(html)
    h = html2text.HTML2Text()
    h.ignore_links = False
    h.ignore_images = False
    h.body_width = 0
    h.unicode_snob = True
    if base_url:
        h.baseurl = base_url
    return h.handle(html)


class SiteArchiver:
    """
    Packages neopig crawl results into a distributable tar.gz archive.

    Uses neopig for crawling, then reads from its vaults to build the archive.
    """

    def __init__(
        self,
        output_dir: str = '.',
        include_screenshots: bool = True,
        include_markdown: bool = True,
        screenshot_config: ScreenshotConfig = None,
        fast_mode: bool = False,
        trim_wrapper: bool = False,
        show_progress: bool = True,
        fresh_start: bool = False,
    ):
        self.output_dir = Path(output_dir)
        self.include_screenshots = include_screenshots
        self.include_markdown = include_markdown and HAS_HTML2TEXT
        self.screenshot_config = screenshot_config or ScreenshotConfig(enabled=include_screenshots)
        self.fast_mode = fast_mode
        self.trim_wrapper = trim_wrapper
        self.show_progress = show_progress
        self.fresh_start = fresh_start

    async def archive(
        self,
        target_url: str,
        depth: int = -1,
        max_pages: int = -1,
        db_path: str = None,
        vault_path: str = None,
        package_only: bool = False,
    ) -> Path:
        """
        Archive a site using neopig and package into tar.gz.
        """
        parsed = Uri(target_url)
        domain = parsed.hostname.lower()
        date_str = datetime.now().strftime('%Y%m%d')
        archive_name = f"{sanitize_filename(domain)}-{date_str}"

        # Use standard neopig data paths - one vault, multiple domains
        if not db_path:
            db_path = "data/neopig.db"
        if not vault_path:
            vault_path = "data/vault"

        # Ensure data directory exists
        Path("data").mkdir(exist_ok=True)

        logger.info(f"Starting archive of {target_url}")
        logger.info(f"Archive name: {archive_name}")

        # Create neopig instance and crawl
        pig = NeoPig(
            db_path=db_path,
            vault_path=vault_path,
            screenshot_config=self.screenshot_config,
            fast_mode=self.fast_mode,
            trim_wrapper=self.trim_wrapper,
        )
        await pig.init()

        # Handle fresh start: rotate state files
        if self.fresh_start:
            pig._clear_state(target_url)
            logger.info("Fresh start: state file rotated")
        else:
            # Load previously crawled pages/media/screenshots from DB (source of truth for resume)
            crawled_pages = await pig.db.get_crawled_page_uris()
            crawled_media = await pig.db.get_crawled_media_uris()
            crawled_screenshots = await pig.db.get_crawled_screenshot_uris()
            logger.info(f"Resume state: {len(crawled_pages)} pages, {len(crawled_media)} media URIs, {len(crawled_screenshots)} screenshotted pages in DB")
            if crawled_pages:
                pig.seen_pages = crawled_pages
            if crawled_media:
                pig.seen_media = crawled_media
            if crawled_screenshots:
                pig.seen_screenshots = crawled_screenshots
        logger.info(f"Screenshots enabled: {pig.screenshot_config.enabled}")

        # Run the crawl (unless package_only)
        stats = {}
        if package_only:
            logger.info("Package-only mode: skipping crawl, packaging existing data")
        else:
            stats = await pig.crawl(
                target_uri=target_url,
                mode=CrawlMode.ALL,
                depth=depth,
                max_pages=max_pages,
                download_media=True,
            )

        # Now package the results
        logger.info("Packaging archive (streaming mode)...")

        tar_path = self.output_dir / f"{archive_name}.tar.gz"

        # Get URL-to-hash mapping for rewriting external URLs to local copies
        logger.info("Loading media URL mappings...")
        url_to_hash = await pig.db.get_all_media_uri_mappings()
        logger.info(f"Loaded {len(url_to_hash)} URL mappings for rewriting")

        def rewrite_urls(html_content: str) -> str:
            """Rewrite external image/media URLs to local archive paths."""
            import re
            def replace_url(match):
                url = match.group(1)
                if url in url_to_hash:
                    md5 = url_to_hash[url]
                    # Get extension from original URL
                    ext = Path(url.split('?')[0]).suffix or '.bin'
                    return match.group(0).replace(url, f'../media/{md5}{ext}')
                return match.group(0)

            # Replace src="url" and href="url" patterns
            html_content = re.sub(r'src=["\']([^"\']+)["\']', replace_url, html_content)
            html_content = re.sub(r'href=["\']([^"\']+\.(jpg|jpeg|png|gif|webp|mp4|webm|svg))["\']',
                                  replace_url, html_content, flags=re.IGNORECASE)
            return html_content

        # Step 1: Load pages from database (has both raw_html and markdown)
        logger.info("Loading pages from database...")
        sitemap = []
        html_contents = {}
        markdown_contents = {}

        pages = await pig.db.get_pages_by_domain(domain)
        for page in pages:
            if not page.get('raw_html'):
                continue
            # Convert URI to path
            uri_path = page.get('path') or ''
            if not uri_path or uri_path == '/':
                rel_path = 'index.html'
            elif uri_path.endswith('.html'):
                rel_path = uri_path.lstrip('/')
            else:
                rel_path = uri_path.strip('/') + '/index.html'

            try:
                content = page['raw_html']
                # Rewrite external URLs to local copies
                content = rewrite_urls(content)
                title = page.get('title') or self._extract_title(content) or rel_path
                sitemap.append({'path': f'html/{rel_path}', 'title': title})
                html_contents[rel_path] = content
                # Store markdown if available
                if page.get('markdown'):
                    markdown_contents[rel_path] = page['markdown']
            except Exception:
                pass

        # Step 2: Create generated files in small temp dir
        local_tmpdir = self.output_dir / '.tmp'
        local_tmpdir.mkdir(parents=True, exist_ok=True)

        with tempfile.TemporaryDirectory(dir=local_tmpdir) as tmpdir:
            tmpdir_path = Path(tmpdir)

            # Write metadata
            metadata = {
                'domain': domain,
                'target_url': target_url,
                'created': datetime.now(timezone.utc).isoformat(),
                'stats': stats,
                'include_screenshots': self.include_screenshots,
                'include_markdown': self.include_markdown,
            }
            (tmpdir_path / 'metadata.json').write_text(json.dumps(metadata, indent=2))

            # Copy state file
            state_file = get_state_file_path(domain)
            if state_file.exists():
                shutil.copy(state_file, tmpdir_path / 'crawl_state.json')
                logger.info(f"Included crawl state for future delta crawls")

            # Write index.html
            self._write_index_html(tmpdir_path, sitemap, domain)

            # Close database BEFORE copying to ensure WAL is checkpointed
            # (SQLite WAL mode keeps data in -wal file until close)
            await pig.db.close()
            logger.info("Database closed, WAL checkpointed")

            # Step 3: Stream everything to tar.gz in one pass
            logger.info("Streaming to archive...")

            # Capture for closure
            show_progress = self.show_progress
            include_markdown = self.include_markdown
            trim_wrapper = self.trim_wrapper
            include_screenshots = self.include_screenshots

            def stream_to_tar():
                with tarfile.open(tar_path, 'w:gz') as tar:
                    # Add generated files first (from temp)
                    for f in tmpdir_path.iterdir():
                        tar.add(f, arcname=f"{archive_name}/{f.name}")

                    # Stream HTML files (with rewritten URLs)
                    html_iter = tqdm(html_contents.items(), desc="HTML", unit="pages", disable=not show_progress)
                    for rel_path_str, content in html_iter:
                        try:
                            html_bytes = content.encode('utf-8')
                            arcname = f"{archive_name}/html/{rel_path_str}"
                            info = tarfile.TarInfo(name=arcname)
                            info.size = len(html_bytes)
                            tar.addfile(info, io.BytesIO(html_bytes))

                            # Use stored markdown (or generate if not available)
                            if include_markdown:
                                md_content = markdown_contents.get(rel_path_str)
                                if not md_content:
                                    md_content = html_to_markdown(content, trim_wrapper=trim_wrapper)
                                md_bytes = md_content.encode('utf-8')
                                md_rel = rel_path_str.replace('.html', '.md')
                                md_arcname = f"{archive_name}/markdown/{md_rel}"
                                md_info = tarfile.TarInfo(name=md_arcname)
                                md_info.size = len(md_bytes)
                                tar.addfile(md_info, io.BytesIO(md_bytes))
                        except Exception as e:
                            logger.debug(f"Error adding HTML {rel_path_str}: {e}")

                    # Stream media files by hash (matching rewritten URLs)
                    hash_vault = Path(vault_path)
                    added_hashes = set()
                    media_iter = tqdm(url_to_hash.items(), desc="Media", unit="files", disable=not show_progress)
                    for url, md5 in media_iter:
                        if md5 in added_hashes:
                            continue
                        # Skip screenshots if not included
                        if not include_screenshots and url.startswith('screenshot:'):
                            continue
                        # Find file in hash vault using 9-layer deep path: vault/ab/cd/.../hash.ext
                        hash_path = hash_vault / hash_to_path(md5)
                        hash_dir = hash_path.parent
                        if hash_dir.exists():
                            for f in hash_dir.iterdir():
                                if f.stem == md5:
                                    try:
                                        # Preserve vault structure for neopig compatibility
                                        rel_path = f.relative_to(hash_vault)
                                        arcname = f"{archive_name}/vault/{rel_path}"
                                        tar.add(f, arcname=arcname)
                                        added_hashes.add(md5)
                                    except Exception as e:
                                        logger.debug(f"Error adding media {f}: {e}")
                                    break

                    # Bundle neopig source files for self-contained crawling
                    neopig_src = Path(__file__).parent
                    neopig_files = [
                        'neopig.py', 'database.py', 'async_web_fetcher.py',
                        'storage.py', 'domain_vault.py', 'screenshot.py',
                        'html2md.py', 'serp.py', 'filevault.py', 'async_filevault.py',
                    ]
                    for pyfile in neopig_files:
                        src_path = neopig_src / pyfile
                        if src_path.exists():
                            tar.add(src_path, arcname=f"{archive_name}/neopig/{pyfile}")

                    # Add static vendor assets (highlight.js for markdown code blocks)
                    vendor_dir = neopig_src / 'static' / 'vendor'
                    if vendor_dir.exists():
                        for asset in vendor_dir.iterdir():
                            if asset.is_file():
                                tar.add(asset, arcname=f"{archive_name}/neopig/static/vendor/{asset.name}")

                    # Add requirements.txt for neopig dependencies
                    req_path = neopig_src / 'requirements.txt'
                    if req_path.exists():
                        tar.add(req_path, arcname=f"{archive_name}/requirements.txt")

                    # Add full neopig database for self-containment
                    db_file = Path(db_path)
                    if db_file.exists():
                        tar.add(db_file, arcname=f"{archive_name}/neopig.db")

            await asyncio.to_thread(stream_to_tar)

        # Clear html_contents to free memory
        html_contents.clear()

        final_size = tar_path.stat().st_size
        logger.info(f"Archive complete: {tar_path} ({final_size / 1024 / 1024:.1f} MB)")

        return tar_path

    def _write_index_html(self, tmpdir: Path, sitemap: list, domain: str):
        """Write index.html with sitemap."""
        html = f'''<!DOCTYPE html>
<html>
<head>
    <title>{domain} Archive</title>
    <style>
        body {{ font-family: sans-serif; max-width: 800px; margin: 40px auto; padding: 0 20px; }}
        h1 {{ color: #333; }}
        ul {{ list-style: none; padding: 0; }}
        li {{ padding: 8px 0; border-bottom: 1px solid #eee; }}
        a {{ color: #0066cc; text-decoration: none; }}
        a:hover {{ text-decoration: underline; }}
    </style>
</head>
<body>
    <h1>{domain} Archive</h1>
    <p>{len(sitemap)} pages archived</p>
    <ul>
'''
        for item in sitemap[:1000]:  # Limit to 1000 in index
            html += f'        <li><a href="{item["path"]}">{item["title"]}</a></li>\n'
        if len(sitemap) > 1000:
            html += f'        <li>... and {len(sitemap) - 1000} more pages</li>\n'
        html += '''    </ul>
</body>
</html>'''
        (tmpdir / 'index.html').write_text(html)

    def _extract_title(self, html: str) -> Optional[str]:
        """Extract title from HTML."""
        try:
            soup = BeautifulSoup(html, 'html.parser')
            if soup.title and soup.title.string:
                return soup.title.string.strip()
            h1 = soup.find('h1')
            if h1:
                return h1.get_text().strip()
        except Exception:
            pass
        return None

def upgrade_neopig_in_archive(archive_path: Path, output_dir: Path = None) -> Path:
    """Replace neopig/ directory in existing archive with current source.

    Uses system tar + pigz for speed (10x faster than Python tarfile).
    Requires temp disk space for extraction (~5-10x compressed size).

    Args:
        archive_path: Path to the tar.gz archive
        output_dir: Directory for temp files (default: cwd)
    """
    import shutil
    import subprocess

    if not archive_path.exists():
        raise FileNotFoundError(f"Archive not found: {archive_path}")

    # Estimate required temp space (compressed * 5 is conservative)
    archive_size = archive_path.stat().st_size
    required_space = archive_size * 5
    logger.info(f"Archive: {archive_size / 1024 / 1024:.0f} MB (need ~{required_space / 1024 / 1024 / 1024:.1f} GB temp space)")

    neopig_src = Path(__file__).parent

    # Check for pigz (parallel gzip) - much faster
    has_pigz = shutil.which('pigz') is not None
    if has_pigz:
        logger.info("Using pigz for parallel compression")

    # Create temp directory for extraction (use output_dir or cwd for space)
    work_dir = Path(output_dir) if output_dir else Path.cwd()
    temp_base = work_dir / '.upgrade_tmp'
    temp_base.mkdir(exist_ok=True)
    with tempfile.TemporaryDirectory(dir=temp_base) as tmpdir:
        tmpdir = Path(tmpdir)

        # Extract archive using system tar (much faster than Python)
        logger.info(f"Extracting archive: {archive_path}")
        subprocess.run(['tar', '-xzf', str(archive_path), '-C', str(tmpdir)], check=True)

        # Find archive root directory
        contents = list(tmpdir.iterdir())
        if len(contents) != 1 or not contents[0].is_dir():
            raise ValueError("Expected single directory in archive")
        archive_root = contents[0]
        archive_name = archive_root.name
        logger.info(f"Archive root: {archive_name}")

        # Remove old neopig directory
        old_neopig = archive_root / 'neopig'
        if old_neopig.exists():
            shutil.rmtree(old_neopig)
            logger.info("Removed old neopig/")

        # Copy new neopig files
        new_neopig = archive_root / 'neopig'
        new_neopig.mkdir()
        neopig_files = [
            'neopig.py', 'database.py', 'async_web_fetcher.py',
            'storage.py', 'domain_vault.py', 'screenshot.py',
            'html2md.py', 'serp.py', 'filevault.py', 'async_filevault.py',
        ]
        for pyfile in neopig_files:
            src = neopig_src / pyfile
            if src.exists():
                shutil.copy2(src, new_neopig / pyfile)
                logger.info(f"  Added: neopig/{pyfile}")

        # Copy static vendor assets
        vendor_src = neopig_src / 'static' / 'vendor'
        if vendor_src.exists():
            vendor_dst = new_neopig / 'static' / 'vendor'
            vendor_dst.mkdir(parents=True, exist_ok=True)
            for asset in vendor_src.iterdir():
                if asset.is_file():
                    shutil.copy2(asset, vendor_dst / asset.name)
                    logger.info(f"  Added: neopig/static/vendor/{asset.name}")

        # Copy requirements.txt
        req_src = neopig_src / 'requirements.txt'
        if req_src.exists():
            shutil.copy2(req_src, archive_root / 'requirements.txt')
            logger.info("  Added: requirements.txt")

        # Repack using system tar (with pigz if available)
        logger.info("Repacking archive...")
        output_path = archive_path.with_suffix('.new.tar.gz')

        if has_pigz:
            # tar + pigz for parallel compression
            with open(output_path, 'wb') as out:
                tar_proc = subprocess.Popen(
                    ['tar', '-cf', '-', '-C', str(tmpdir), archive_name],
                    stdout=subprocess.PIPE
                )
                pigz_proc = subprocess.Popen(
                    ['pigz', '-c'],
                    stdin=tar_proc.stdout,
                    stdout=out
                )
                tar_proc.stdout.close()
                pigz_proc.wait()
                tar_proc.wait()
        else:
            # Standard tar with gzip
            subprocess.run(
                ['tar', '-czf', str(output_path), '-C', str(tmpdir), archive_name],
                check=True
            )

        # Replace original
        logger.info("Replacing original archive...")
        shutil.move(output_path, archive_path)

    # Cleanup temp base
    try:
        temp_base.rmdir()
    except OSError:
        pass

    logger.info(f"Done! Upgraded: {archive_path}")
    return archive_path


async def main():
    parser = argparse.ArgumentParser(
        description="Archive a website for preservation (uses neopig for crawling)",
        epilog="Example: python archive.py https://discourse-urho3d.github.io/"
    )

    parser.add_argument("url", nargs='?', help="URL of the site to archive")
    parser.add_argument("-o", "--output", default=".", help="Output directory for tar.gz")
    parser.add_argument("-d", "--depth", type=int, default=-1, help="Crawl depth: 0=single page, 1=page+links, -1=unlimited")
    parser.add_argument("-p", "--max-pages", type=int, default=-1, help="Max pages (-1 = unlimited)")
    parser.add_argument("--no-screenshot", "--no-screenshots", action="store_true", help="Disable screenshots")
    parser.add_argument("--screenshot-width", type=int, default=1280, help="Screenshot width")
    parser.add_argument("--screenshot-height", type=int, default=1024, help="Screenshot height")
    parser.add_argument("--screenshot-engine", type=str, default=None, help="Screenshot engine")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")
    parser.add_argument("--fresh", action="store_true", help="Start fresh (rotates old state files instead of resuming)")
    parser.add_argument("--fast", action="store_true", help="Fast mode: no crawl delay (for sites without robots.txt)")
    parser.add_argument("--package-only", action="store_true", help="Skip crawling, just package existing data from vault")
    parser.add_argument("--serve", action="store_true", help="Start SERP server to watch crawl live")
    parser.add_argument("--port", type=int, default=31337, help="Port for SERP server (default: 31337)")
    parser.add_argument("--backfill-markdown", metavar="DOMAIN", help="Re-process stored HTML for DOMAIN to regenerate markdown with absolute URLs")
    parser.add_argument("--trim-wrapper", action="store_true", help="With --backfill-markdown: strip nav/header/footer/logo before conversion")
    parser.add_argument("--db", default="data/neopig.db", help="Database path (for --backfill-markdown)")
    parser.add_argument("--upgrade-neopig", metavar="TARBALL", help="Upgrade neopig inside an existing archive")
    parser.add_argument("--no-progress", action="store_true", help="Disable progress bars")

    args = parser.parse_args()

    # Validate depth - max 18 before suggesting unlimited
    if args.depth > 18:
        print(f"Error: depth {args.depth} is too high. Use --depth -1 for unlimited crawling.")
        return

    setup_logging(level=logging.DEBUG if args.verbose else logging.INFO)

    # Handle --upgrade-neopig
    if args.upgrade_neopig:
        archive_path = Path(args.upgrade_neopig)
        output_dir = Path(args.output) if args.output != "." else None
        logger.info(f"Upgrading neopig in {archive_path}...")
        upgrade_neopig_in_archive(archive_path, output_dir=output_dir)
        logger.info(f"Done! Archive updated: {archive_path}")
        return

    # Handle --backfill-markdown
    if args.backfill_markdown:
        from neopig import backfill_markdown
        await backfill_markdown(args.db, domain_filter=args.backfill_markdown, trim_wrapper=args.trim_wrapper)
        return

    # URL required for archiving
    if not args.url:
        parser.error("URL required (use --upgrade-neopig or --backfill-markdown for other operations)")

    if not HAS_HTML2TEXT:
        logger.warning("html2text not installed - markdown disabled")

    screenshot_config = ScreenshotConfig(
        enabled=not args.no_screenshot,
        width=args.screenshot_width,
        height=args.screenshot_height,
        engine=args.screenshot_engine,
        full_page=True,  # Archive captures full page by default
    )

    archiver = SiteArchiver(
        output_dir=args.output,
        include_screenshots=not args.no_screenshot,
        include_markdown=True,
        screenshot_config=screenshot_config,
        fast_mode=args.fast,
        trim_wrapper=args.trim_wrapper,
        show_progress=not args.no_progress,
        fresh_start=args.fresh,
    )

    # Start SERP server if requested
    serp_process = None
    if args.serve:
        serp_script = Path(__file__).parent / 'serp.py'
        if serp_script.exists():
            serp_cmd = [
                sys.executable, str(serp_script),
                '--port', str(args.port),
                '--db', 'data/neopig.db',
                '--vault', 'data/vault',
            ]
            serp_process = subprocess.Popen(
                serp_cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            logger.info(f"SERP server started at http://localhost:{args.port}")
            logger.info("Watch the crawl live - pages appear as they're indexed!")
        else:
            logger.warning("serp.py not found - --serve disabled")

    try:
        archive_path = await archiver.archive(
            target_url=args.url,
            depth=args.depth,
            max_pages=args.max_pages,
            package_only=args.package_only,
        )

        print(f"\nArchive created: {archive_path}")
        print(f"Extract with: tar -xzf {archive_path.name}")
    finally:
        # Cleanup SERP server
        if serp_process:
            logger.info("Stopping SERP server...")
            serp_process.terminate()
            try:
                serp_process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                serp_process.kill()


if __name__ == "__main__":
    asyncio.run(main())