# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
# One which is durable, easy to repair, & distributed like tap water
# for machine learning intelligence.
#
# The permacomputer is community-owned infrastructure optimized around
# four values:
#
#   TRUTH      First principles, math & science, open source code freely distributed
#   FREEDOM    Voluntary partnerships, freedom from tyranny & corporate control
#   HARMONY    Minimal waste, self-renewing systems with diverse thriving connections
#   LOVE       Be yourself without hurting others, cooperation through natural law
#
# This software contributes to that vision by archiving the web, preserving digital knowledge before it disappears.
# Code is seeds to sprout on any abandoned technology.

"""
Screenshot backfill operations for neopig.

Re-captures screenshots as JPEG to replace old PNGs.
"""

import asyncio
import json
import logging
import multiprocessing
import os
import time
from datetime import datetime, timezone
from pathlib import Path

from miniuri import Uri
from tqdm import tqdm

from ..state import get_state_file_path

logger = logging.getLogger(__name__)


async def backfill_screenshots(
    db_path: str,
    vault_path: str = "vault",
    domain_filter: str = None,
    delete_old: bool = True,
    fast_mode: bool = False,
    quiet: bool = False,
):
    """Re-capture screenshots as JPEG to replace old PNGs.

    Args:
        db_path: Path to SQLite database
        vault_path: Path to vault directory
        domain_filter: Only process pages matching this domain
        delete_old: Delete old PNG files after successful JPEG capture
        fast_mode: Skip delay between captures (for sites without robots.txt)
        quiet: Disable progress bar
    """
    from screenshot import ScreenshotCapture, ScreenshotConfig
    from filevault import AsyncVault, hash_to_path
    from database import Database, Media, MediaSource, Page, SCORE_SCREENSHOT
    from sqlalchemy import select, update

    # Load completed URIs from unified state file
    completed_uris = set()
    state_path = get_state_file_path(domain_filter) if domain_filter else None
    if state_path and state_path.exists():
        try:
            full_state = json.loads(state_path.read_text())
            completed_uris = set(full_state.get('backfill_screenshots', []))
            if completed_uris:
                logger.info(f"Resuming: {len(completed_uris)} already backfilled")
        except Exception as e:
            logger.warning(f"Could not load state: {e}")

    config = ScreenshotConfig(enabled=True, full_page=True)
    capture = ScreenshotCapture(config)
    vault = AsyncVault(vault_path)
    await vault.init()

    db = Database(db_path)
    await db.init()

    async with db.session() as session:
        # Find all screenshot records using ORM, excluding already-completed URIs
        if domain_filter:
            pattern = f"%://{domain_filter}%" if '://' not in domain_filter else f"%{domain_filter}%"
            stmt = (
                select(MediaSource.page_uri, Media.md5_hash)
                .join(Media, MediaSource.md5_hash == Media.md5_hash)
                .where(Media.media_type == 'screenshot')
                .where(MediaSource.page_uri.like(pattern))
                .where(MediaSource.page_uri.notin_(completed_uris) if completed_uris else True)
                .distinct()
            )
            logger.info(f"Backfilling screenshots for domain: {domain_filter}")
        else:
            stmt = (
                select(MediaSource.page_uri, Media.md5_hash)
                .join(Media, MediaSource.md5_hash == Media.md5_hash)
                .where(Media.media_type == 'screenshot')
                .where(MediaSource.page_uri.notin_(completed_uris) if completed_uris else True)
                .distinct()
            )
            logger.info("Backfilling screenshots for ALL pages")

        result = await session.execute(stmt)
        rows = result.fetchall()
        to_process = len(rows)
        skipped_count = len(completed_uris)
        total = to_process + skipped_count
        logger.info(f"Found {to_process} to process ({skipped_count} already completed)")

        domain_last_fetched = {}  # Track last fetch time per domain
        crawl_delay = 2.0  # Default crawl delay in seconds

        # Concurrent workers in fast mode
        max_workers = multiprocessing.cpu_count() * 2 if fast_mode else 1
        semaphore = asyncio.Semaphore(max_workers)
        if fast_mode and max_workers > 1:
            logger.info(f"Fast mode: {max_workers} concurrent workers")

        pbar = tqdm(total=total, initial=skipped_count, desc="Screenshots", unit="pages", disable=quiet, smoothing=0.1)

        # Thread-safe counters
        stats_lock = asyncio.Lock()
        stats = {'captured': 0, 'failed': 0, 'bytes_saved': 0}

        async def capture_one(page_uri: str, old_hash: str):
            """Capture screenshot for one page."""
            async with semaphore:
                try:
                    async with db.session() as local_session:
                        from sqlalchemy import func as sqlfunc
                        result = await local_session.execute(
                            select(sqlfunc.length(Page.raw_html)).where(Page.uri == page_uri)
                        )
                        content_length = result.scalar() or 0

                    result = await capture.capture(page_uri, content_length=content_length)
                    if not result:
                        return False, None
                    return True, (page_uri, old_hash, result)
                except Exception as e:
                    logger.debug(f"Capture error {page_uri}: {e}")
                    return False, None

        async def process_result(success, data):
            """Process a captured screenshot - store and update DB."""
            nonlocal completed_uris
            if not success or not data:
                async with stats_lock:
                    stats['failed'] += 1
                pbar.update(1)
                return

            page_uri, old_hash, result = data
            domain = Uri(page_uri).hostname.lower()

            try:
                chunks = result if isinstance(result, list) else [result]
                old_path = Path(vault_path) / hash_to_path(old_hash, ext='png')
                old_size = old_path.stat().st_size if old_path.exists() else 0

                # Store chunks
                total_new_size = 0
                for i, chunk in enumerate(chunks):
                    new_hash = chunk['md5_hash']
                    new_data = chunk['data']
                    new_ext = chunk.get('format', 'jpg')
                    total_new_size += len(new_data)

                    if not await vault.exists(new_hash):
                        await vault.store(new_hash, new_data, new_ext)

                    # Symlink
                    parsed = Uri(page_uri)
                    url_path = (parsed.path or '').lstrip('/') or 'index'
                    suffix = f'_{i}' if len(chunks) > 1 else ''
                    new_symlink = Path(vault_path) / 'linkpeek_vault' / domain / (url_path.replace('/', '_') + f'{suffix}.{new_ext}')
                    new_symlink.parent.mkdir(parents=True, exist_ok=True)
                    hash_vault_path = Path(vault_path) / hash_to_path(new_hash, ext=new_ext)
                    rel_path = os.path.relpath(hash_vault_path, new_symlink.parent)
                    if not new_symlink.exists():
                        new_symlink.symlink_to(rel_path)

                # Delete old
                if delete_old and old_path.exists():
                    old_path.unlink()
                    async with stats_lock:
                        stats['bytes_saved'] += old_size - total_new_size

                # Remove old symlink
                parsed = Uri(page_uri)
                url_path = (parsed.path or '').lstrip('/') or 'index'
                old_symlink = Path(vault_path) / 'linkpeek_vault' / domain / (url_path.replace('/', '_') + '.png')
                if old_symlink.exists():
                    old_symlink.unlink()

                # Update database
                new_hash = chunks[0]['md5_hash']
                new_ext = chunks[0].get('format', 'jpg')
                new_mime = chunks[0].get('mime_type', 'image/jpeg')
                new_size = total_new_size

                async with db.session() as local_session:
                    from sqlalchemy import delete as sqldelete
                    old_media_result = await local_session.execute(
                        select(Media).where(Media.md5_hash == old_hash)
                    )
                    old_media = old_media_result.scalar_one_or_none()

                    if old_media and old_hash != new_hash:
                        existing = await local_session.execute(
                            select(Media).where(Media.md5_hash == new_hash)
                        )
                        if not existing.scalar_one_or_none():
                            new_media = Media(
                                md5_hash=new_hash,
                                media_type='screenshot',
                                mime_type=new_mime,
                                file_size=new_size,
                                keywords=old_media.keywords,
                                alt_text=old_media.alt_text,
                                title=old_media.title,
                                first_seen_at=datetime.now(timezone.utc).isoformat(),
                                analysis_status=old_media.analysis_status,
                                analysis_result=old_media.analysis_result,
                                score=SCORE_SCREENSHOT,
                            )
                            local_session.add(new_media)
                            await local_session.flush()
                        await local_session.execute(
                            update(MediaSource).where(MediaSource.md5_hash == old_hash).values(md5_hash=new_hash)
                        )
                        await local_session.execute(sqldelete(Media).where(Media.md5_hash == old_hash))
                    elif old_media:
                        old_media.mime_type = new_mime
                        old_media.file_size = new_size
                        old_media.first_seen_at = datetime.now(timezone.utc).isoformat()
                    await local_session.commit()

                async with stats_lock:
                    stats['captured'] += 1
                    completed_uris.add(page_uri)

                # Save state
                if state_path:
                    async with stats_lock:
                        try:
                            full_state = json.loads(state_path.read_text()) if state_path.exists() else {}
                        except Exception:
                            full_state = {}
                        full_state['backfill_screenshots'] = list(completed_uris)
                        state_path.write_text(json.dumps(full_state, indent=2))

                pbar.update(1)

            except Exception as e:
                err_msg = str(e).split('\n')[0][:60]
                logger.warning(f"Failed: {page_uri} - {err_msg}")
                async with stats_lock:
                    stats['failed'] += 1
                pbar.update(1)

        # Process
        if fast_mode and max_workers > 1:
            async def capture_and_process(row):
                success, data = await capture_one(row.page_uri, row.md5_hash)
                await process_result(success, data)

            tasks = [asyncio.create_task(capture_and_process(row)) for row in rows]
            await asyncio.gather(*tasks)
        else:
            for row in rows:
                if not fast_mode:
                    domain = Uri(row.page_uri).hostname.lower()
                    last_fetched = domain_last_fetched.get(domain, 0)
                    elapsed = time.time() - last_fetched
                    if elapsed < crawl_delay:
                        await asyncio.sleep(crawl_delay - elapsed)
                    domain_last_fetched[domain] = time.time()

                success, data = await capture_one(row.page_uri, row.md5_hash)
                await process_result(success, data)

        pbar.close()
        logger.info(f"Backfill complete: {stats['captured']} captured, {skipped_count} skipped, {stats['failed']} failed")
        if stats['bytes_saved'] > 0:
            logger.info(f"Space saved: {stats['bytes_saved'] / 1024 / 1024:.1f} MB")
