import os
import shutil
import hashlib
import logging
import requests
from itertools import groupby
from cStringIO import StringIO
from contextlib import contextmanager

import numpy as np
import sqlalchemy as sa
from unstdlib import to_str, to_unicode
from unstdlib.standard.contextlib_ import open_atomic
from sqlalchemy.ext.declarative import declarative_base

from akindi.utils import json

from .global_ import g
from .utils import get_file_sha1, mkdir

log = logging.getLogger(__name__)

SHA1Type = sa.String(40)

session = None
def set_engine(engine):
    global session
    session = sa.orm.sessionmaker(bind=engine, autocommit=True)()


@contextmanager
def atomic():
    session.begin(subtransactions=True)
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise


class Base(declarative_base()):
    __abstract__ = True

    @classmethod
    def create(cls, **kwargs):
        obj = cls(**kwargs)
        session.add(obj)
        return obj


class BaseRepo(object):
    def __init__(self, repo):
        self.repo = repo
        self._init()

    def _init(self):
        pass

    def refresh(self):
        if not os.path.exists(self.storage_root):
            raise ValueError("Repository storage root does not exist: %s" %(self.storage_root, ))

        added_count = 0
        paths_in_db = set(v[0] for v in self.pdfs.values(PDF.path))
        paths_on_disk = set()
        for path in self._walk_files():
            paths_on_disk.add(path)
            if path not in paths_in_db:
                log.info("Adding %s", path)
                PDF.create_from_file(self, path)
                added_count += 1
        print "Added %d PDFs" %(added_count, )

        for path in (paths_in_db - paths_on_disk):
            log.info("Removing %s", path)
            self.pdfs.filter(PDF.path == path).delete(synchronize_session=False)
        print "Removed %d PDFs not found on disk" %(len(paths_in_db - paths_on_disk), )

    def _walk_files(self, cur=None):
        if cur is None:
            cur = self.storage_root
        if os.path.basename(cur).startswith("."):
            return
        if cur.lower().endswith(".pdf") and os.path.isfile(cur):
            yield cur[len(self.storage_root):].lstrip("/")
        elif os.path.isdir(cur):
            for filename in os.listdir(cur):
                filename = os.path.join(cur, filename)
                for subfile in self._walk_files(filename):
                    yield to_unicode(subfile)

    def __getattr__(self, attr):
        try:
            return getattr(self.repo, attr)
        except AttributeError:
            raise AttributeError(attr)


class LocalRepo(BaseRepo):
    def _init(self):
        self.storage_root = self.root

    def get_absolute_path(self, path):
        return os.path.join(self.storage_root, path)

    def deleted(self):
        pass


class RemoteRepo(BaseRepo):
    # URL: https://akindi.com/uploads/_proctor/uploads

    def _init(self):
        self.blob_root = os.path.join(g.env.root, "blobs/")
        self.storage_root = os.path.join(g.env.root, "repos/%s" %(self.name, ))

    def get_absolute_path(self, path):
        return os.path.join(self.storage_root, path)

    def fetch(self):
        try:
            session.commit()
            self._fetch()
        finally:
            session.begin()

    def refresh(self):
        if not os.path.exists(self.storage_root):
            os.makedirs(self.storage_root)
        return super(RemoteRepo, self).refresh()

    def _fetch(self):
        resp = requests.get(self.root)
        resp.raise_for_status()
        data = resp.json()
        if not data["ok"]:
            raise AssertionError("Error refreshing repo: %r" %(data, ))

        mkdir(os.path.dirname(self.storage_root))
        mkdir(self.storage_root)
        mkdir(self.blob_root)
        file_list = data["data"]

        for file in file_list:
            fname = "%s-%s-t%s-c%s.pdf" %(
                self.name,
                file["id"],
                file["metadata"]["teacher_id"],
                file["metadata"]["course_id"],
            )
            fpath = os.path.join(self.storage_root, fname)
            if os.path.exists(fpath):
                continue

            log.info("importing: %s --> %s", file["url"], fname)
            with open_atomic(fpath) as outf:
                resp = requests.get(file["url"], stream=True)
                if resp.status_code > 299:
                    outf.abort()
                    log.warning("bad status code %r while fetching %r: %s",
                                resp.status_code, file["url"], resp.text)
                    continue

                sha1_hash = hashlib.sha1()
                for chunk in resp.iter_content(chunk_size=4096):
                    sha1_hash.update(chunk)
                    outf.write(chunk)

                sha1 = sha1_hash.hexdigest()
                outf.target_name = os.path.join(self.blob_root, sha1)
                if os.path.exists(outf.target_name):
                    log.info("ignoring ignoring duplicate: %s (%s)", fname, sha1)
                    continue

            with atomic():
                os.symlink(outf.target_name, fpath)
                pdf = PDF.create_from_file(self, fname)
                pdf.set_attribute(None, "source/url", file["url"])
                for name, val in file["metadata"].items():
                    pdf.set_attribute(None, "source/meta/%s" %(name, ), val)

    def delete(self):
        shutil.rmtree(self.storage_root)


class Repository(Base):
    __tablename__ = "repos"
    id = sa.Column(sa.Integer, primary_key=True)
    name = sa.Column(sa.String(16), unique=True)
    root = sa.Column(sa.String(256))
    pdfs = sa.orm.relationship("PDF", lazy="dynamic", cascade="all, delete, delete-orphan")

    def __init__(self, *args, **kwargs):
        super(Base, self).__init__(*args, **kwargs)
        self._init()

    @sa.orm.reconstructor
    def _init(self):
        is_remote = (
            self.root.startswith("http://") or
            self.root.startswith("https://")
        )
        self._impl = is_remote and RemoteRepo(self) or LocalRepo(self)

    @classmethod
    def get(cls, name):
        repo = session.query(cls).filter_by(name=name).first()
        if repo is None:
            raise ValueError("Repository does not exist: %r" %(name, ))
        return repo

    def get_pdf(self, path):
        pdf = self.pdfs.filter_by(path=path).first()
        if pdf is None:
            raise ValueError("Repository does not contain PDF: %r" %(path, ))
        return pdf

    def fetch(self):
        return self._impl.fetch()

    def refresh(self):
        self._impl.refresh()

    def get_absolute_path(self, path):
        return self._impl.get_absolute_path(path)

    def open_path(self, path):
        return open(self.get_absolute_path(path))

    def drop(self):
        session.delete(self)
        self._impl.delete()


class PDF(Base):
    __tablename__ = "pdfs"
    id = sa.Column(sa.Integer, primary_key=True)
    sha1 = sa.Column(SHA1Type, index=True)
    repo_name = sa.Column(
        sa.String(16),
        sa.ForeignKey(Repository.name, ondelete="CASCADE"),
    )
    path = sa.Column(sa.String(256), index=True)

    repo = sa.orm.relationship("Repository")
    _attrs = sa.orm.relationship("PDFAttr")

    __table_args__ = (
        sa.UniqueConstraint('repo_name', 'path'),
    )

    @classmethod
    def create_from_file(cls, repo, path):
        with repo.open_path(path) as f:
            sha1 = get_file_sha1(f)
        return cls.create(
            sha1=sha1,
            repo_name=repo.name,
            path=path,
        )

    @classmethod
    def get_unique(cls, last_id=None):
        query = (
            session
                .query(cls)
                .order_by(cls.id.asc())
        )
        if last_id:
            query = query.filter(cls.id > last_id)

        seen_sha1s = set()
        for pdf in query:
            if pdf.sha1 in seen_sha1s:
                continue
            seen_sha1s.add(pdf.sha1)
            yield pdf


    def set_attribute(self, page_num, name, value):
        value_num = (
            value if isinstance(value, (int, long, float, complex)) else
            None
        )
        session.add(PDFAttr(
            pdf_sha1=self.sha1,
            page_num=page_num,
            name=name,
            value_bytes=to_str(value),
            value_num=value_num,
        ))

    def get_absolute_path(self):
        return self.repo.get_absolute_path(self.path)


class PDFAttr(Base):
    __tablename__ = "pdf_attrs"

    id = sa.Column(sa.Integer, primary_key=True)
    pdf_sha1 = sa.Column(SHA1Type, sa.ForeignKey(PDF.sha1), index=True)
    page_num = sa.Column(sa.Integer, nullable=True)
    name = sa.Column(sa.String(128), index=True)
    value_bytes = sa.Column(sa.LargeBinary())
    value_num = sa.Column(sa.Float(), nullable=True)

    __table_args__ = (
        sa.UniqueConstraint('pdf_sha1', 'page_num', 'name'),
    )


class PDFPage(object):
    def __init__(self, pdf_sha1, page_num, pdf_attrs, attrs):
        self.pdf_sha1 = pdf_sha1
        self.page_num = page_num
        self.pdf_attrs = pdf_attrs
        self.attrs = dict(attrs)
        self.attrs.update(pdf_attrs)

    def get_pdf(self):
        return session.query(PDF).filter_by(sha1=self.pdf_sha1).first()

    @classmethod
    def load_all(cls, filter):
        attr_query = (
            session
                .query(PDFAttr)
                .join(PDF)
                .order_by(
                    PDFAttr.pdf_sha1.desc(),
                    PDFAttr.page_num.asc(),
                )
        )
        if "repo" in filter:
            attr_query = attr_query.filter(PDF.repo_name.in_(filter["repo"]))
        grouped = groupby(attr_query, key=lambda a: (a.pdf_sha1, a.page_num))
        cur_pdf = {}
        for _, attrs in grouped:
            attrs = list(attrs)
            attr_dict = dict(
                (a.name, a.value_num if a.value_num is not None else a.value_bytes)
                for a in attrs
            )
            if attrs[0].page_num is None:
                cur_pdf = attr_dict
                continue
            yield cls(attrs[0].pdf_sha1, attrs[0].page_num, cur_pdf, attr_dict)


class PDFPageOracle(Base):
    __tablename__ = "pdf_page_oracle"

    id = sa.Column(sa.Integer, primary_key=True)
    pdf_sha1 = sa.Column(SHA1Type, sa.ForeignKey(PDF.sha1), index=True)
    page_num = sa.Column(sa.Integer)
    created_on = sa.Column(sa.DateTime, default=sa.func.now())
    name = sa.Column(sa.String(128), index=True)
    data = sa.Column(sa.LargeBinary())

    def get_message_json(self):
        return {
            "type": self.__tablename__,
            "id": self.id,
            "pdf_sha1": self.pdf_sha1,
            "page_num": self.page_num,
            "name": self.name,
            "data": self.data,
        }


class ImgprocTestRun(Base):
    __tablename__ = "imgproc_test_run"

    id = sa.Column(sa.Integer, primary_key=True)
    argv = sa.Column(sa.Text())
    started_on = sa.Column(sa.DateTime, default=sa.func.now())
    finished_on = sa.Column(sa.DateTime, nullable=True)

    pages = sa.orm.relationship("ImgprocPage", lazy="dynamic",
                                cascade="all, delete, delete-orphan")

    def get_message_json(self):
        return {
            "type": self.__tablename__,
            "id": self.id,
            "argv": json.loads(self.argv),
        }


class ImgprocPDF(Base):
    __tablename__ = "imgproc_pdf"

    id = sa.Column(sa.Integer, primary_key=True)
    sha1 = sa.Column(SHA1Type, index=True)
    name = sa.Column(sa.String(128))
    path = sa.Column(sa.String(512), index=True)

    def get_message_json(self):
        return {
            "type": self.__tablename__,
            "id": self.id,
            "sha1": self.sha1,
            "name": self.name,
            "path": self.path,
        }


class ImgprocPage(Base):
    __tablename__ = "imgproc_page"

    id = sa.Column(sa.Integer, primary_key=True)
    test_run_id = sa.Column(sa.Integer, sa.ForeignKey(ImgprocTestRun.id), index=True)
    test_run = sa.orm.relationship("ImgprocTestRun")
    pdf_sha1 = sa.Column(SHA1Type, index=True)
    page_num = sa.Column(sa.Integer)
    cvpage = sa.Column(sa.LargeBinary())

    oracles = sa.orm.relationship(
        "PDFPageOracle", uselist=True, lazy="dynamic", backref="imgproc_pages",
        primaryjoin="""and_(
            PDFPageOracle.pdf_sha1 == foreign(ImgprocPage.pdf_sha1),
            PDFPageOracle.page_num == foreign(ImgprocPage.page_num),
    )""")

    def get_message_json(self):
        cvpage = json.loads(self.cvpage)
        if "localizer" in cvpage.get("cv") or {}:
            l = cvpage["cv"]["localizer"] or {}
            if "barycentric_npz" in l:
                s = StringIO(l.pop("barycentric_npz").decode("base64"))
                cvpage["cv"]["localizer"] = dict(np.load(s))
        return {
            "type": self.__tablename__,
            "id": self.id,
            "test_run": self.test_run_id,
            "pdf_sha1": self.pdf_sha1,
            "page_num": self.page_num,
            "cvpage": cvpage,
            "oracles": [
                x[0] for x in self.oracles.distinct(PDFPageOracle.name).values("name")
            ],
        }
