from .sql import create_table_statements, create_view_statements
from .sql import table_triggers
from .util import addr_to_int, AttrCheck
import storm.locals as rain

class Execution:
    def __init__(self, db):
        self.db = db
        self.execution = None

    def __enter__(self):
        from .util import right_now
        self.execution = self.db.start_execution(right_now())
        return self.execution

    def __exit__(self, exc_type, exc_value, traceback):
        from .util import right_now
        completed = exc_type is None
        self.db.finish_execution(self.execution, right_now(), completed)

class TException(rain.Storm, AttrCheck):
    __storm_table__ = "Exceptions"
    exception_id = rain.Int("ExceptionId", primary=True)
    name = rain.Unicode("Name")
    fail = rain.Bool("Fail")

class TExecution(rain.Storm, AttrCheck):
    __storm_table__ = "Executions"
    execution_id = rain.Int("ExecutionId", primary=True)
    start_date = rain.DateTime("StartDate")
    finish_date = rain.DateTime("FinishDate")
    completed = rain.Bool("Completed")

class TAddress(rain.Storm, AttrCheck):
    __storm_table__ = "Ips"
    address_id = rain.Int("IpId", primary=True)
    str = rain.Unicode("AsStr")
    ip = rain.Int("AsInt")
    china = rain.Bool("China")
    block_target = rain.Bool("BlockTarget")
    server = rain.Bool("Server")
    redirect_target = rain.Bool("RedirectTarget")
    gfw_target = rain.Bool("GfwTarget")

class TKind(rain.Storm, AttrCheck):
    __storm_table__ = "Kinds"
    kind_id = rain.Int("KindId", primary=True)
    name = rain.Unicode("Name")
    xxid = rain.Int("ExpectExceptionId")
    exception = rain.Reference(xxid, "TException.exception_id")

class TDomain(rain.Storm, AttrCheck):
    __storm_table__ = "Domains"
    domain_id = rain.Int("DomainId", primary=True)
    name = rain.Unicode("Name")
    kind_id = rain.Int("KindId")
    kind = rain.Reference(kind_id, "TKind.kind_id")

class TRecord(rain.Storm, AttrCheck):
    __storm_table__ = "Records"
    row_id = rain.Int("rowid", primary=True)
    record_id = rain.Int("RecordId")
    address_id = rain.Int("IpId")
    address = rain.Reference(address_id, "TAddress.address_id")

class TMessage(rain.Storm, AttrCheck):
    __storm_table__ = "Messages"
    message_id = rain.Int("MessageId", primary=True)
    execution_id = rain.Int("ExecutionId")
    server_id = rain.Int("ServerId")
    domain_id = rain.Int("DomainId")
    record_id = rain.Int("RecordId")
    exception_id = rain.Int("ExceptionId")
    execution = rain.Reference(execution_id, "TExecution.execution_id")
    server = rain.Reference(server_id, "TAddress.address_id")
    domain = rain.Reference(domain_id, "TDomain.domain_id")
    #record = rain.Reference(record_id, "TRecord.record_id")
    exception = rain.Reference(exception_id, "TException.exception_id")

def apply_properties(obj, d):
    from storm.properties import PropertyColumn
    for k, v in d.items():
        ref = getattr(obj.__class__, k)
        assert ref is not None, (type(obj), k)
        assert isinstance(ref, PropertyColumn) or isinstance(ref, rain.Reference), \
            (type(obj), k)
        setattr(obj, k, v)
    return obj

class RespoDB:
    def __init__(self, uri, setup=False, create=False):
        self.uri = uri
        db_exists = self._db_exists(self.uri)
        self.db = rain.create_database(self.uri)
        self._conn = None

        if setup or (create and not db_exists):
            with self:
                self.setup_executions()
                self.setup_exceptions()
                self.setup_ips()
                self.setup_kinds()
                self.setup_domains()
                self.setup_records()
                self.setup_messages()

                for q in create_view_statements:
                    self._conn.execute(q, noresult=True)
        assert setup or create or db_exists, "database was never setup"

        self.execution = Execution(self)

    @staticmethod
    def _db_exists(uri):
        from os.path import exists
        _, _, fp = uri.partition(":")
        if fp.startswith("//"):
            _, _, fp = fp[2:].partition("/")
        return fp and exists(fp)

    def __enter__(self):
        self._conn = rain.Store(self.db)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.commit()
        self._conn.close()
        self._conn = None

    def find_one(self, cls_spec, *args, **kwargs):
        assert self._conn is not None
        return self._conn.find(cls_spec, *args, **kwargs).one()

    def flush(self):
        assert self._conn is not None
        self._conn.flush()

    def commit(self):
        assert self._conn is not None
        self._conn.commit()

    def new_exception(self, **kwargs):
        assert self._conn is not None
        return self._conn.add(apply_properties(TException(), kwargs))

    def new_kind(self, **kwargs):
        assert self._conn is not None
        return self._conn.add(apply_properties(TKind(), kwargs))

    def new_domain(self, **kwargs):
        assert self._conn is not None
        return self._conn.add(apply_properties(TDomain(), kwargs))

    def new_address(self, **kwargs):
        assert self._conn is not None
        return self._conn.add(apply_properties(TAddress(), kwargs))

    def new_record(self, **kwargs):
        assert self._conn is not None
        return self._conn.add(apply_properties(TRecord(), kwargs))

    def new_message(self, **kwargs):
        assert self._conn is not None
        return self._conn.add(apply_properties(TMessage(), kwargs))

    def setup_executions(self):
        self._conn.execute(create_table_statements["executions"], noresult=True)

    def setup_exceptions(self):
        # careful not to call them "errors" since NXDOMAIN is not an error.
        self._conn.execute(create_table_statements["exceptions"], noresult=True)

        # TODO: upsert?

        self.new_exception(name="NXDOMAIN", fail=False)
        self.new_exception(name="NoAnswer", fail=True)
        self.new_exception(name="NoNameservers", fail=True)
        self.new_exception(name="Timeout", fail=True)

    def setup_ips(self):
        from .ips import china, blocks

        self._conn.execute(create_table_statements["ips"], noresult=True)

        # TODO: upsert?

        self.new_address(ip=addr_to_int("0.0.0.0"), block_target=True)
        self.new_address(ip=addr_to_int("127.0.0.1"), block_target=True)
        for ip in china:
            self.new_address(ip=addr_to_int(ip), china=True)
        for ip in blocks:
            self.new_address(ip=addr_to_int(ip), block_target=True)

    def setup_kinds(self):
        self._conn.execute(create_table_statements["kinds"], noresult=True)

        # TODO: upsert?

        #NXDOMAIN = self.find_one(TException, TException.name == "NXDOMAIN")
        #self.new_kind(name="bad", exception=NXDOMAIN)
        #self.new_kind(name="badsub", exception=NXDOMAIN)

    def setup_domains(self):
        self._conn.execute(create_table_statements["domains"], noresult=True)

    def setup_records(self):
        self._conn.execute(create_table_statements["records"], noresult=True)

    def setup_messages(self):
        self._conn.execute(create_table_statements["messages"], noresult=True)
        for trig in table_triggers["messages"]:
            self._conn.execute(trig)

    def start_execution(self, dt):
        execution = TExecution()
        execution.start_date = dt
        self.flush()
        return execution

    def finish_execution(self, execution, dt, completed):
        # TODO: fail if ExecutionId is missing?
        execution.finish_date = dt
        execution.completed = completed
        self.flush()

    def next_record_id(self):
        from storm.expr import Add, Max, Coalesce
        expr = Add(Coalesce(Max(TRecord.record_id), 0), 1)
        return self.find_one(expr)

    def find_record_id(self, addresses):
        address_ids = list(address.address_id for address in addresses)
        record_ids = list(self._conn.find(TRecord, TRecord.address_id.is_in(address_ids)).values(TRecord.record_id))
        if not record_ids:
            return None
        unique_ids = sorted(set(record_ids))
        for needle in unique_ids:
            if sum(1 for id in record_ids if id == needle) == len(addresses):
                found = True
                return needle
        return None

    def push_entry(self, entry):
        kind = self.find_one(TKind, TKind.name == entry.kind)
        if not kind:
            kind = self.new_kind(name=entry.kind)
        if entry.kind.startswith("bad"):
            exception = self.find_one(TException, TException.name == "NXDOMAIN")
            assert exception is not None
            kind.exception = exception

        domain = self.find_one(TDomain, TDomain.name == entry.domain)
        if not domain:
            domain = self.new_domain(name=entry.domain)
        domain.kind = kind

        addresses = []
        as_ints = sorted(set(map(addr_to_int, entry.addrs)))
        for numeric in as_ints:
            address = self.find_one(TAddress, TAddress.ip == numeric)
            if not address:
                address = self.new_address(ip=numeric)
            addresses.append(address)

        for address in addresses:
            if entry.reason == "block":
                address.block_target = True
            elif entry.reason == "redirect":
                address.redirect_target = True
            elif entry.reason == "gfw":
                address.gfw_target = True

        if addresses:
            record_id = self.find_record_id(addresses)
            if record_id is None:
                record_id = self.next_record_id()
                for address in addresses:
                    self.new_record(record_id=record_id, address=address)
        else:
            record_id = None

        numeric = addr_to_int(entry.server)
        server = self.find_one(TAddress, TAddress.ip == numeric)
        if not server:
            server = self.new_address(ip=numeric)
            self.flush()
        server.server = True

        if entry.exception:
            exception = self.find_one(TException, TException.name == entry.exception)
            assert exception is not None
        else:
            exception = None

        message = self.new_message(
            execution=entry.execution,
            server=server, domain=domain,
            record_id=record_id, exception=exception)
        self.flush()