import gzip
import json
import os
from typing import List

from sqlalchemy import text
from sqlalchemy.engine import ResultProxy
from sqlalchemy.exc import DatabaseError
from sqlalchemy.orm import Session
from structlog import get_logger

from delphi_es_utils.constants import INDEX_PG_COL_WHITELIST, INDEX_PG_SCHEMA, INDEX_PG_TABLE
from delphi_es_utils.repository.pg_client import get_session_from_secret_key
from delphi_es_utils.settings import DB_CREDS_SECRET_NAME, TEMP_DIR

LOG = get_logger(__name__)


class DataManager:
    def __init__(self, index_name: str, table=None, schema=None):
        self.index_name = index_name
        self.table = table or INDEX_PG_TABLE.get(index_name)
        self.schema = schema or INDEX_PG_SCHEMA.get(index_name)
        self.temp_fp = os.path.join(TEMP_DIR, f'{self.schema}.{self.table}.json')
        _session_maker = get_session_from_secret_key(DB_CREDS_SECRET_NAME)
        self.session: Session = _session_maker()

    def get_data(self) -> List[dict]:
        """Main method to call"""
        try:
            items = self.export_data()
        except (DatabaseError, FileNotFoundError) as e:
            LOG.exception('Data export failed. error=%s' % e)
            raise
        if not items:
            LOG.warning('No items found in data export')
            return []
        return self.filter_cols(items)

    def export_data(self) -> List[dict]:
        table_ref = f'{self.schema}.{self.table}'
        stmt = text(f'SELECT json_agg(t) FROM {table_ref} AS t;')
        result: ResultProxy = self.session.execute(stmt)
        return result.fetchone()[0]

    def filter_cols(self, data: List[dict]):
        whitelist = INDEX_PG_COL_WHITELIST.get(self.index_name, set([]))
        return [{k: v for k, v in el.items() if k in whitelist} for el in data]

    def verify_data_export(self):
        if not os.path.isfile(self.temp_fp):
            raise FileNotFoundError(f'No data export exists at {self.temp_fp}. Aborting.')

    @classmethod
    def load_data(cls, file_path) -> List[dict]:
        """Reads data from a JSON file"""
        with open(file_path, 'rt') as file_obj:
            return json.load(file_obj)

    @classmethod
    def load_gz_data(cls, file_path) -> List[dict]:
        """Reads data from a GZIP encoded JSON file"""
        with gzip.open(file_path, 'rt') as file_obj:
            return json.load(file_obj)