# -*- coding: utf-8 -*-
"""
    module
    ~~~~~~~~~~~~~~~~

    @Copyright: (c) 2018-04 by Lingxi Chen (chanlingxi@gmail.com).
    @License: LICENSE_NAME, see LICENSE for more details.
"""
import os
from itertools import tee


class Module(object):

    def __init__(self, in_data=None, in_fn=None,
                 reader_cls=None, file_type=None,
                 record_cls=None,
                 write=False, write_stats=True, write_headers=False,
                 write_fields=True, dynamic_fields=False,
                 writer_cls=None, write_mode='whole', out_dir='',
                 prefix='', suffix='', chunk_size=1000,
                 sample=None,
                 **extra):
        self.sample = sample

        self.in_fn = in_fn
        self.in_data = in_data
        self.reader_cls = reader_cls
        self.file_type = file_type
        self.record_cls = record_cls

        if in_data:
            if file_type:
                self.in_iter = self.reader_cls(in_data=in_data, sample=sample, file_type=file_type)
            else:
                self.in_iter = self.reader_cls(in_data=in_data, sample=sample)
        elif in_fn:
            if file_type:
                self.in_iter = self.reader_cls(in_fn=in_fn, sample=sample, file_type=file_type)
            else:
                self.in_iter = self.reader_cls(in_fn=in_fn, sample=sample)
        else:
            self.in_iter = in_data

        self.write = write
        self.write_stats = write_stats
        self.write_headers = write_headers
        self.write_fields = write_fields
        self.dynamic_fields = dynamic_fields
        self.writer_cls = writer_cls
        self.out_fn = None
        self.out_fns = []

        self.write_mode = write_mode
        self.out_dir = out_dir
        self.prefix = prefix
        self.suffix = suffix
        self.chunk_size = chunk_size

        self.meta_dict = {}
        self.writer_dict = {}
        self.meta_count = {}

    def evaluate(self):
        gen1, gen2 = tee(self)

        if not os.path.exists(self.out_dir):
            os.makedirs(self.out_dir)
        for meta, record in gen1:
            if self.write:
                self._write_line(meta, record)

        if self.write:
            self._write_after_iter()

        return gen2

    def __iter__(self):
        pass

    def _write_line(self, meta, line):
        # update count
        count = self.meta_count.get(meta, 0)
        self.meta_count[meta] = count + 1

        # write
        if self.write_mode == 'whole':
            meta = 'ALL'
        chunk = self.meta_dict.get(meta, [])

        if len(chunk) >= self.chunk_size:
            self._write_chunk(meta, chunk)
            self.meta_dict[meta] = [line]
        else:
            chunk.append(line)
            self.meta_dict[meta] = chunk

    def _write_chunk(self, meta, chunk):
        dynamic_fields = False
        if meta not in self.writer_dict:
            if meta == 'ALL':
                fns = [self.prefix, self.suffix]
            else:
                fns = [self.prefix, meta, self.suffix]
            fns = [x for x in fns if x]
            self.out_fn = os.path.join(self.out_dir, '.'.join(fns))
            if self.out_fn not in self.out_fns:
                self.out_fns.append(self.out_fn)
            writer = self.writer_cls(fn=self.out_fn)
            if self.write_headers:
                writer.write_headers()
            if self.write_fields:
                if not self.dynamic_fields:
                    writer.write_fields()
                else:
                    dynamic_fields = True
            self.writer_dict[meta] = writer
        else:
            writer = self.writer_dict[meta]

        writer.write_chunk(chunk, dynamic_fields=dynamic_fields)

    def _write_after_iter(self):
        self._write_last_chunk()
        if self.write_stats:
            self._write_stat()

    def _write_last_chunk(self):
        for meta, chunk in self.meta_dict.items():
            self._write_chunk(meta, chunk)

    def _write_stat(self):
        stat_fn = os.path.join(
            self.out_dir, self.prefix + '.' + self.suffix + '.stat')
        with open(stat_fn, 'w') as f:
            f.write('\t'.join(map(str, self.meta_count.keys())) + '\n')
            f.write('\t'.join(map(str, self.meta_count.values())) + '\n')
