
from typing import Any, Collection, Dict, Iterable, List
import sqlite3
import shutil
from copy import deepcopy as dcp

import logging

logger = logging.getLogger(__name__)
record_cols = [
    'lib_id',
    'resident_id',
    'name',
    'gender',
    'age',
    'clinical',
    'pathol',
    'pre_test',
    'history',
    'type',
    'smpl_dt',
    'recv_dt',
    'rprt_dt',
    'q30',
    'aver_depth',
    'target_reads',
    'cov_pos',
    'cov_t1_t2',
    'dna_total',
    'lib_total',
    'err_code',
    'kdb_file',
    'kdb_version',
    'template_version',
    'plasma_volume',
    'sample_source',
    'med_history',
    'institution'
]
variant_cols = [
    'record_id',
    'gene',
    'gene_transcript',
    'exon',
    'cds',
    'aa',
    'vaf',
]
fusion_cols = [
    'record_id',
    'gene',
    'exon',
    'gene2',
    'exon2',
    'cds',
]
create_tables_queries = [
    """CREATE TABLE IF NOT EXISTS "record" (
        "id" INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
        "lib_id" TEXT NOT NULL UNIQUE,
        "resident_id" TEXT NOT NULL,
        "name" TEXT,
        "gender" TEXT,
        "age" INTEGER,
        "clinical" TEXT,
        "pathol" TEXT,
        "pre_test" TEXT,
        "history" TEXT,
        "type" TEXT,
        "smpl_dt" TEXT,
        "recv_dt" TEXT NOT NULL,
        "rprt_dt" TEXT NOT NULL,
        "q30" REAL,
        "aver_depth" REAL,
        "target_reads" REAL,
        "cov_pos" REAL,
        "cov_t1_t2" REAL,
        "dna_total" REAL,
        "lib_total" REAL,
        "err_code" INTEGER,
        "kdb_file" TEXT,
        "kdb_version" TEXT NOT NULL,
        "template_version" TEXT NOT NULL,
        "plasma_volume" REAL,
        "sample_source" TEXT,
        "med_history" TEXT,
        "institution" TEXT
    );""",
    """CREATE TABLE IF NOT EXISTS "variant" (
        "id" integer NOT NULL,
        "record_id" integer NOT NULL,
        "gene" text NOT NULL,
        "gene_transcript" text,
        "exon" text,
        "cds" text,
        "aa" text,
        "vaf" real,
        PRIMARY KEY ("id"),
        CONSTRAINT "fk_variant_record_1" FOREIGN KEY ("record_id") REFERENCES "record" ("id") ON DELETE CASCADE ON UPDATE CASCADE DEFERRABLE INITIALLY DEFERRED
    );""",
    """CREATE TABLE IF NOT EXISTS "fusion" (
        "id" INTEGER NOT NULL,
        "record_id" INTEGER,
        "gene" TEXT NOT NULL,
        "exon" TEXT NOT NULL,
        "gene2" TEXT NOT NULL,
        "exon2" TEXT NOT NULL,
        "cds" TEXT NOT NULL,
        PRIMARY KEY ("id"),
        CONSTRAINT "fk_fusion_record_1" FOREIGN KEY ("record_id") REFERENCES "record" ("id") ON DELETE CASCADE ON UPDATE CASCADE DEFERRABLE INITIALLY DEFERRED
    );"""]
upgdate_table_cols_queries = [
    """ALTER TABLE record ADD COLUMN sample_source TEXT;""",
    """ALTER TABLE record ADD COLUMN med_history TEXT;""",
    """ALTER TABLE record ADD COLUMN institution TEXT;"""]


class RecordDatabase:
    """Database module for handle YUEDONG tests analysis records."""

    def __init__(self, db_file: str = ''):
        self._db_file = ''
        if db_file:
            self.set_db_file(db_file)

    def set_db_file(self, db_file: str) -> None:
        """Set database file, and use it as for all database operations.

        Args:
            db_file(str): Database file path.
        """
        self._db_file = db_file
        self._create_tables()
        self._add_new_columns()

    def _get_conn(self) -> sqlite3.Connection:
        """ create a database connection to the SQLite database
            specified by db_file
        : return: Connection object or None
        """
        try:
            conn = sqlite3.connect(self._db_file)
            return conn
        except Exception as e:
            logger.error(e)
            raise(e)

    def _fetchall(self, query) -> List[Dict[str, Any]]:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute(query)
        res = cur.fetchall()
        data_all = [dict(zip([c[0] for c in cur.description], r)) for r in res]
        conn.close()
        return data_all

    def _fetchone(self, query) -> Dict[str, Any]:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute(query)
        res = cur.fetchone()
        data_one = {}
        try:
            data_one = dict(zip([c[0] for c in cur.description], res))
        except TypeError:
            pass
        finally:
            conn.close()
            return data_one

    def _update_db(self, sql: str, parameters: Iterable[Any] = {}) -> int:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute(sql, parameters)
        conn.commit()
        conn.close()
        return cur.lastrowid

    def _create_tables(self):
        for query in create_tables_queries:
            self._update_db(query)

    def _add_new_columns(self):
        for query in upgdate_table_cols_queries:
            try:
                self._update_db(query)
            except sqlite3.OperationalError:
                continue

    def insert(self, table_name: str, data_dict: Dict[str, Any], cols: Collection[str]) -> int:
        """Insert data in table.

        Args:
            table_name(str): name of table
            data_dict(Dict[str, Any]): data in dict form with column name as key
            cols(Collection[str]): column names

        Returns:
            int: last row id
        """
        query_template = f'INSERT INTO {table_name} (' \
            + ','.join(cols) + ') VALUES(' \
            + ','.join(['?' for i in range(len(cols))]) \
            + ');'
        data_tuple = tuple(data_dict[col] for col in cols)
        return self._update_db(query_template, data_tuple)

    def update(
            self,
            table_name: str,
            data_dict: Dict[str, Any],
            cols: Collection[str],
            where_col: str,
            where_val: Any) -> int:
        """Update data in table.

        Args:
            table_name(str): name of table
            data_dict(Dict[str, Any]): data in dict form with column name as key
            cols(Collection[str]): column names
            where_col(str): column name for where clause
            where_val(Any): value for where clause

        Returns:
            int: last row id
        """
        query_template = f'UPDATE {table_name} SET ' \
            + ','.join([f'{col}=?' for col in cols]) \
            + f' WHERE {where_col}=?;'
        data_tuple = tuple(data_dict[col] for col in cols) + (where_val,)
        return self._update_db(query_template, data_tuple)

    def save_sample(self, sample_dict2) -> int:
        """Save sample dict version into record database.

        Args:
            sample_dict(dict): sample dict
        """
        sample_dict = dcp(sample_dict2)
        sample_dict['kdb_version'] = "server"
        record_dict = {col: sample_dict[col] for col in record_cols}
        # template_version 列，如果是阅动，则设置成阅动的模板版本
        if sample_dict['report_type'] and 'yuedong' in sample_dict['report_type']:
            record_dict['template_version'] = "yuedong_" + sample_dict['yuedong_template_version']
        record_id = self.insert('record', record_dict, record_cols)

        noid_variant_cols = variant_cols.copy()
        noid_variant_cols.remove('record_id')
        for variant in sample_dict['variants']:
            variant_dict = {col: variant[col] for col in noid_variant_cols}
            variant_dict['record_id'] = record_id
            self.insert('variant', variant_dict, variant_cols)

        noid_fusion_cols = fusion_cols.copy()
        noid_fusion_cols.remove('record_id')
        for fusion in sample_dict['fusions']:
            fusion_dict = {col: fusion[col] for col in noid_fusion_cols}
            fusion_dict['record_id'] = record_id
            self.insert('fusion', fusion_dict, fusion_cols)

        return record_id

    def update_sample_rprt_dt(self, sample_dict) -> int:
        """Update sample report date in record database.

        Args:
            sample_dict(dict): sample dict
        """
        return self.update(
            'record',
            sample_dict,
            ['rprt_dt'],
            'lib_id',
            sample_dict['lib_id'])

    def get_variants(self, record_id: int) -> Collection[Dict[str, Any]]:
        query = f'SELECT * FROM variant WHERE record_id="{record_id}";'
        variant_dicts = self._fetchall(query)
        return variant_dicts

    def get_fusions(self, record_id: int) -> Collection[Dict[str, Any]]:
        query = f'SELECT * FROM fusion WHERE record_id="{record_id}";'
        fusion_dicts = self._fetchall(query)
        return fusion_dicts

    def get_sample_dict(self, lib_id: str) -> Dict[str, Any]:
        query = f'SELECT * FROM record WHERE lib_id="{lib_id}";'
        sample_dict = self._fetchone(query)
        if not sample_dict:
            return sample_dict
        sample_dict['variants'] = self.get_variants(
            sample_dict['id'])
        sample_dict['fusions'] = self.get_fusions(
            sample_dict['id'])
        return sample_dict

    def get_patient_sample_dicts(self, resident_id: str) -> Collection[Dict[str, Any]]:
        query = f'SELECT * FROM record WHERE resident_id="{resident_id}";'
        sample_dicts = self._fetchall(query)
        for i in range(len(sample_dicts)):
            sample_dicts[i]['variants'] = self.get_variants(
                sample_dicts[i]['id'])
            sample_dicts[i]['fusions'] = self.get_fusions(
                sample_dicts[i]['id'])
            # nan改为空值
            for k, v in sample_dicts[i].items():
                if v == 'nan':
                    sample_dicts[i][k] = ''
        return sample_dicts

    def make_backup(self, backup_file: str) -> None:
        shutil.copyfile(self._db_file, backup_file)
