#!/usr/bin/env python3

# -*- coding: utf-8 -*-

import gzip

from singlecellmultiomics.pyutils.handlelimiter import HandleLimiter

class FastqHandle:

    def __init__(

            self,

            path,

            pairedEnd=False,

            single_cell=False,

            maxHandles=500):

        self.pe = pairedEnd

        self.sc = single_cell

        self.path = path

        if not self.sc:

            if pairedEnd:

                self.handles = [

                    gzip.open(

                        path +

                        'R1.fastq.gz',

                        'wt',compresslevel=1),

                    gzip.open(

                        path +

                        'R2.fastq.gz',

                        'wt',compresslevel=1)]

            else:

                self.handles = [gzip.open(path + 'R1.fastq.gz', 'wt',compresslevel=1)]

        else:

            self.handles = HandleLimiter(

                compressionLevel=1, maxHandles=maxHandles)

    def write(self, records):

        if self.sc:

            for readIdx, record in zip(('R1', 'R2'), records):

                # Obtain cell from record:

                cell = f"{record.tags.get('bi','no_cell_id')}.{record.tags.get('MX','unk')}"

                self.handles.write(

                    f'{self.path}.{cell}.{readIdx}.fastq.gz',

                    str(record),

                    method=1)

        else:

            for handle, record in zip(self.handles, records):

                handle.write(str(record))

    def close(self):

        if self.sc:

            self.handles.close()

        else:

            for handle in self.handles:

                handle.close()