"""$URL: svn://svn/repos/trunk/sitereaper/lib/page_cache.py $
$Id: page_cache.py 20952 2003-03-05 16:51:11Z dbinger $

   Provides PageCache class.
   A PageCache caches Pages keyed by URL so that we avoid
   retrieving the same page repeatedly.
"""
import sys, os
import time, atexit, shutil, md5
from struct import unpack
from urlparse import urlparse, urlunparse
from distutils.dir_util import mkpath

from sitereaper.verbose import Verbose
from sitereaper.utils import warn
from sitereaper.cookies import CookieJar, CookieSet
from sitereaper.page import Page, Redirect

TMP_DIR = "/tmp"


class PageCache (Verbose):
    """PageCache caches pages.
    The primary method used is get_page(), which returns a Page
    from the cache, creating a new Page if none already exists
    for the given URL.

    Instance attributes:
      page_map : { ( url : string, Page ) }
        the cache of Page objects (not necessarily fetched from
        the web yet)
      cookie_jar : CookieJar
        the collection of cookies seen in our wandering across the
        web; every time we make a request, we pull out a subset
        of the cookies in the cookie jar and pass them along
      size : int
        number of pages in the cache
    """

    def __init__ (self, verbose=0):
        """PageCache(verbose : int = 0)"""
        Verbose.__init__(self, verbose)
        #self.cache_dir = os.path.join(TMP_DIR, "sitereaper.%d" % os.getpid())
        self.cache_dir = os.path.join(TMP_DIR, "sitereaper")
        mkpath(self.cache_dir)
        urlmap_filename = os.path.join(self.cache_dir, "urlmap")
        self.urlmap = open(urlmap_filename, "w")

        # Register a cleanup function.
        atexit.register(self.cleanup)

        self.page_map = {}
        self.cookie_jar = CookieJar()
        self.size = 0

    def cleanup (self):
        self.note(5, "PageCache.cleanup: nuking %s" % self.cache_dir)
        shutil.rmtree(self.cache_dir, ignore_errors=1)

    def get_size (self):
        return self.size

    def get_page (self, orig_url, retrieve=0):
        """get_page(orig_url : string) -> Page

        The fragment part of the given url is stripped, and the
        resulting url is used as a key into the page_map.  If the
        page_map has a Page with this key (or for this key with a final
        '/' removed), the Page is returned.  Otherwise, a new Page is
        inserted into the page_map for this key, and the new Page is
        returned.
        """
        self.note(6, 'get_page: original URL = ', orig_url)
        url_bits = urlparse(orig_url)[:-1]
        url = urlunparse(url_bits + ('',))
        if url != orig_url:
            self.note(6, 'get_page: de-fragged URL = ', url)

        # Check the cache for this URL.
        if self.page_map.has_key(url):
            self.note(7, 'get_page: found page in cache')
            return self.page_map[url]

        # Not in the cache -- create a new Page object.
        self.note(7, 'get_page: page not in cache, creating Page object')
        page = Page(url, self.get_verbose())
        self.size += 1
        self.page_map[url] = page

        if retrieve:
            return self.retrieve_page(page)
        else:
            return page


    def _request_page (self, page, base_filename):
        req_cookies = self.cookie_jar.get_cookies(page.get_netloc(),
                                                  page.get_path())
        if req_cookies:
            self.note(6, "will send %d cookie(s) with request: %s"
                      % (len(req_cookies), req_cookies))
            headers = {"Cookie" : req_cookies.make_request_header()}
        else:
            headers = None

        try:
            if base_filename:
                ok = page.retrieve(headers=headers,
                                   base_filename=base_filename)
            else:
                ok = page.check(headers=headers)
        except Redirect, redirect:
            # Start all over again with new URL (new cache filename,
            # new cookies, etc.)
            new_page = self.get_page(redirect.url)
            if base_filename:
                return self.retrieve_page(new_page)
            else:
                return self.check_page(new_page)

        if ok:
            cookie_hdr = page.response.getheader("Set-cookie")
            if cookie_hdr:
                self.note(7, "got a set-cookie header: %r" % cookie_hdr)
                resp_cookies = CookieSet()
                resp_cookies.parse_header(cookie_hdr,
                                          page.get_netloc(),
                                          page.get_path())
                self.note(6, "server set %d cookie(s): %s"
                          % (len(resp_cookies), resp_cookies))
                self.cookie_jar.add_cookies(resp_cookies)

        return page

    def retrieve_page (self, page):
        """retrieve_page(page : Page) -> Page

        Fetch the page content by doing an HTTP GET request.  Follows
        redirects, so the returned Page might be different from the
        requested Page (but they will be identical if there's no
        redirect).
        """
        if page.retrieved:
            # XXX what if page was a redirect?
            return page
        elif page.checked:
            warn("retrieving page %s after previously checking it" % page)

        # Compute a cache base filename for this URL by taking the MD5
        # checksum of the URL.  (The Page will add an extension to this
        # base filename once it knows the content-type.)
        md5sum = md5.new(page.url).digest()
        basename = "".join(["%08x" % i for i in unpack(">LLLL", md5sum)])
        self.urlmap.write("%s %s\n" % (basename, page.url))
        base_filename = os.path.join(self.cache_dir, basename)

        # And do the HTTP request.
        return self._request_page(page, base_filename)

    def check_page (self, page):
        """check_page(page : Page) -> Page

        Check the URL of 'page' by doing an HTTP HEAD request.  Handles
        redirects just like retrieve_page(), ie. the returned Page will
        be the passed-in object unless there is a redirect involved.
        """
        if page.checked:
            # XXX again, what if page was a redirect?
            return page
        else:
            return self._request_page(page, None)


    def has_page (self, url):
        return self.page_map.has_key(url)

    def page_checked (self, url):
        map = self.page_map
        return map.has_key(url) and map[url].checked

    def page_retrieved (self, url):
        map = self.page_map
        return map.has_key(url) and map[url].retrieved

    def get_urls (self):
        return self.page_map.keys()

    def get_urls_pages (self):
        return self.page_map.items()

    def get_pages (self):
        return self.page_map.values()


