import os
import re
from types import GeneratorType

from unstdlib import to_unicode
from akindi.utils import json
from akindi.imgproc.process import is_image_file_path

from .utils import get_exif, get_file_sha1, unsafe_gs_get_page_count

def flatten(gen, res=None):
    if res is None:
        res = []

    for bit in gen:
        if isinstance(bit, GeneratorType):
            flatten(bit, res)
            continue
        res.append(bit)
    return res

class ProctorTestCase(object):
    def __init__(self, pdf_testcase, page_num):
        self._pdf = pdf_testcase
        self.env = pdf_testcase.env
        self.pdf_sha1 = self._pdf.pdf_sha1
        self.pdf_name = self._pdf.pdf_name
        self.pdf_path = self._pdf.pdf_path
        self.page_num = page_num

    def save_oracle(self, name, data):
        db = self.env.db
        db.PDFPageOracle.create(
            pdf_sha1=self.pdf_sha1,
            page_num=self.page_num,
            name=name,
            data=json.dumps(data),
        )
        db.session.flush()

    def load_oracle(self, name):
        db = self.env.db
        oracle = (
            db.session.query(db.PDFPageOracle)
                .filter_by(
                    name=name,
                    pdf_sha1=self.pdf_sha1,
                    page_num=self.page_num,
                )
                .order_by(db.PDFPageOracle.id.desc())
                .first()
        )
        if oracle is None:
            return {}
        return json.loads(oracle.data)


class PDFTests(object):
    def __init__(self, env, pdf_sha1, pdf_name, pdf_path, page_nums):
        self.env = env
        self.pdf_sha1 = pdf_sha1
        self.pdf_name = pdf_name
        self.pdf_path = pdf_path
        self.page_nums = page_nums
        self.is_image = is_image_file_path(pdf_path)

    def _get_page_count(self):
        exif = get_exif(self.pdf_path)
        if "page-count" in exif:
            return int(exif["page-count"])
        return int(unsafe_gs_get_page_count(self.pdf_path))

    def iter_page_nums(self):
        if self.is_image:
            yield 1
            return
        last_num = 0
        last_tok = None
        for tok in self.page_nums:
            if tok == "-":
                last_tok = tok
                continue
            if tok == "*":
                last_tok = "-"
                tok = self._get_page_count()
            if not isinstance(tok, int):
                raise AssertionError(
                    "invalid page num token %r in %r" %(tok, self.page_nums)
                )
            if last_tok == "-":
                while last_num < tok - 1:
                    last_num += 1
                    yield last_num
            last_num = tok
            yield tok
            continue

    def iter_test_cases(self):
        for page_num in self.iter_page_nums():
            yield ProctorTestCase(self, page_num)


class ProctorTestCaseSet(object):
    def __init__(self, pdf_tests):
        self._pdf_tests = pdf_tests

    @classmethod
    def load_from_argv(cls, env, argv):
        return cls(flatten(cls._load_test_cases_from_argv(env, argv)))

    @classmethod
    def _load_test_cases_from_argv(cls, env, argv):
        for arg in argv:
            if ":" in arg or arg.endswith(".pdf") or is_image_file_path(arg):
                yield cls._load_test_cases_from_token(env, arg)
            elif os.path.isdir(arg):
                yield cls._load_test_cases_from_dir(env, arg)
            else:
                with open(arg) as f:
                    yield cls._load_test_cases_from_file(env, f)

    @classmethod
    def _load_test_cases_from_file(cls, env, f):
        for raw_line in f:
            line = to_unicode(re.sub(r"^\s*(.*?)\s*(#.*)?$", r"\1", raw_line).strip())
            if not line:
                continue
            yield cls._load_test_cases_from_token(env, line)

    @classmethod
    def _load_test_cases_from_token(cls, env, token):
        parts = token.split(":")
        if len(parts) == 1:
            repo_name, path, pages = None, parts[0], "*"
        elif len(parts) == 2:
            repo_name, path, pages = None, parts[0], parts[1]
        elif len(parts) == 3:
            repo_name, path, pages = parts
        else:
            raise AssertionError("invalid test case specification: %r" %(token, ))

        expect = "any"
        page_nums = []
        for tok in re.split("([0-9]+)", pages):
            if not tok:
                continue
            tok = tok.strip()
            if tok.isdigit() and expect in ("any", "num"):
                page_nums.append(int(tok))
                expect = "sep"
            elif tok in "-,*" and expect in ("any", "sep"):
                if tok != ",":
                    page_nums.append(tok)
                expect = "num"
            else:
                raise AssertionError("invalid page specification %r in %r" %(tok, pages, ))

        if repo_name is not None:
            repo = env.repository_get(repo_name)
            pdf = repo.get_pdf(path)
            pdf_sha1 = pdf.sha1
            pdf_name = path
            pdf_path = pdf.get_absolute_path()
        else:
            with open(path) as f:
                pdf_sha1 = get_file_sha1(f)
            pdf_name = path
            pdf_path = path
        yield PDFTests(env, pdf_sha1, pdf_name, pdf_path, page_nums)
    
    def iter_test_cases(self):
        for pdf_testcase in self._pdf_tests:
            for testcase in pdf_testcase.iter_test_cases():
                yield testcase
