sakhteman/core/utils/te/engines/Mongo.py

import datetime
import uuid
from typing import Union, List
from uuid import UUID
import os
import xlsxwriter
import tempfile
import string
import copy
import jdatetime


# import dateutil.parser
from celery import shared_task
# from dateutil.relativedelta import relativedelta

from .AbstractEngine import AbstractEngineClass
from ..utils.Models import FilterItem, FilterScope, OutPutType
from ..utils.excel_value_maps import EXCEL_VALUE_MAP_REGISTRY
from app.repositories.mysql_const_repo import MySQLConstRepository


# def zero_time_data(date, is_max=False):
#     _date = dateutil.parser.isoparse(date)
#     return datetime.datetime.combine(_date, datetime.time.min if not is_max else datetime.time.max)


# def date_year_month_morph(date, is_max=False, has_month=True):
#     _date = date
#     if len(date) > 7:
#         _date = dateutil.parser.isoparse(date)
#     if not has_month and len(date) < 5:
#         _date = f'{date}-01'
#     _date = datetime.datetime.strptime(_date, '%Y-%m')
#     _date = _date + relativedelta(day=31)
#     return _date if is_max else _date + relativedelta(day=1)


class MongoEngine(AbstractEngineClass):
    rows = []
    stack = []
    filters = {}
    _const_cache: dict = None  # کش const ها از MySQL

    async def exe(self):
        _f = self.filters = self.set_filters()
        if self.unselect:
            tmp = self.query_class.find(_f, self.unselect)
        else:
            tmp = self.query_class.find(_f)
        sort_by = self.set_sort()
        if isinstance(sort_by, dict):
            tmp = tmp.sort(**sort_by)
        else:
            tmp = tmp.sort(sort_by)

        is_natural = isinstance(sort_by, dict) and sort_by.get('key_or_list') == '$natural'
        if getattr(self, 'collation', None) and not is_natural:
            tmp = tmp.collation(self.collation)

        if self.input.e and self.input.e.ot:
            res = await tmp.to_list(None)
            self.input.e.d = await self.generate_output(res)
            return {'query': _f, 'sort': self.set_sort(), 'path': self.input.e.d}

        self.rows = tmp.limit(self.input.get_limit()).skip(self.input.skip())
        if self.return_query:
            return {'query':_f, 'sort':self.set_sort(), 'limit':self.input.get_limit(), 'page':self.input.skip()}
        return [self.model(**document) async for document in self.rows], await self.count()

    def query_generator(self):
        if not self.input.has_filter():
            return []
        filters = [f for f in self.input.filters() if
                   (hasattr(f, 'items') and len(f.items)) > 0 or not hasattr(f, 'items')]
        query = self.generate_scope_filter(filters, general=True)
        return query

    def set_filters(self):
        filters = []
        query = self.query_generator()
        if query:
            filters.append(query)
        if self.input.has_text_search():
            items = []
            _search_cols = self.input.get_search_cols() if (
                    self.input.get_search_cols() and len(self.input.get_search_cols())) else self.search_cols
            if _search_cols:
                for col in _search_cols:
                    items.append({col: {'$regex': self.input.get_text_search(), '$options': 'i'}})
                filters.append({"$or": items})
        if self.input.extra_filter:
            filters.extend(self.input.extra_filter)
        if self.soft_deleted:
            filters.append({'deleted_at': {'$exists': False}})
        if not filters:
            return {}
        self.filters = {'$and': filters}
        return self.filters

    def set_sort(self):
        if self.input.o.fields:
            all_sort = []
            for item in self.input.o.fields:
                all_sort.append((item.field, 1 if item.direction == "asc" else -1))
            return all_sort
        if self.input.o.field:
            return {'key_or_list': self.input.o.field,
                    'direction': 1 if self.input.o.direction == "asc" else -1}
        elif self.default_sort_by:
            return self.default_sort_by
        else:
            return {'key_or_list': '$natural', 'direction': -1}

    @staticmethod
    def filter_item_generator(item: FilterItem) -> dict:
        if MongoEngine.is_valid_uuid(str(item.col)):
            return {
                'data': {
                    "$elemMatch": {
                        'field_id': UUID(item.col),
                        'value': {
                            **MongoEngine.get_expression_filter(item.op, MongoEngine.cast_value(item.val, item.type),
                                                                field_type=item.type)}
                    }
                }
            }
        return {item.col: {
            **MongoEngine.get_expression_filter(col=item.col, op=item.op,
                                                val=MongoEngine.cast_value(item.val, item.type), field_type=item.type)}}

    @staticmethod
    def is_valid_uuid(uuid_to_test, version=4):
        try:
            uuid_obj = UUID(str(uuid_to_test), version=version)
        except ValueError:
            return False
        return str(uuid_obj) == str(uuid_to_test)

    @staticmethod
    def cast_value(val: Union[str, List], type: str) -> any:
        if val == [] or type == 'BOOL':
            return val
        if not val:
            return None
        if type in ['NUMBER', 'STACK', 'FLOAT']:
            if isinstance(val, List):
                return [float(x) for x in val]
            return float(val)
        if type in ['DATE', 'DATETIME']:
            if isinstance(val, List):
                return [datetime.datetime.fromisoformat(str(x)) for x in val]
            return datetime.datetime.fromisoformat(val)
        if type == 'UUID':
            if isinstance(val, List):
                return [UUID(str(x)) for x in val]
            return UUID(str(val))
        return val

    @staticmethod
    def get_expression_filter(col=None, op: str = 'eq', val: str = '', field_type: str = None) -> dict:
        _op = MongoEngine.get_operator(op)

        # Date Type Operator
        if field_type in ['DATE', 'DATE_TIME', 'DATETIME', 'DATE_MONTH', 'DATE_YEAR']:
            if op == "eq":
                return {'$gte': val[0], '$lte': val[1]}
            if op == "n_eq":
                return {'$lte': val[0], '$gte': val[1]}
            if op == "lt":
                return {'$lt': val[1]}
            if op == "gt":
                return {'$gt': val[0]}
            if op == "lte":
                return {'$lt': val[1]}
            if op == "gte":
                return {'$gt': val[0]}

        if _op == "sw":
            return {'$regex': f'^{val}', '$options': 'i'}
        elif _op == "ew":
            return {'$regex': f'{val}$', '$options': 'i'}
        elif _op == "n_ew":
            return {'$regex': f'(?<!{val})$', '$options': 'i'}
        elif _op == "n_sw":
            return {'$regex': f'^(?!{val})', '$options': 'i'}
        elif _op == "n_like":
            return {'$not': {'$regex': val, '$options': 'i'}}
        elif _op == "like":
            return {'$regex': val, '$options': 'i'}
        elif _op == "bet":
            if isinstance(val, str):
                return {'$gte': val.split('-')[0], '$lte': val.split('-')[1]}
            return {'$gte': val[0], '$lte': val[1]}
        elif _op == "n_bet":
            if isinstance(val, str):
                return {'$gte': val.split('-')[0], '$lte': val.split('-')[1]}
            return {'$not': {'$gte': val[0], '$lte': val[1]}}
        elif _op == "exi":
            return {'$exists': True}
        elif _op == "n_exi":
            return {'$exists': False}
        elif _op == "in":
            return {'$in': val}
        elif _op == "not_in":
            return {'$nin': val}
        return {_op: val}

    @staticmethod
    def get_operator(op: str) -> str:
        ops = {
            "eq": "$eq",
            "n_eq": "$ne",
            "gt": "$gt",
            "ge": "$ge",
            "lt": "$lt",
            "le": "$le",
            "like": "like",
            "n_like": "n_like",
            "sw": "sw",
            "ew": "ew",
            'n_ew': 'n_ew',
            'n_sw': 'n_sw',
            'exi': 'exi',
            "n_exi": "n_exi",
            'lte': '$lte',
            'gte': '$gte',
            'in': 'in',
            'n_in': 'not_in',
            'bet': 'bet',
            'n_bet': 'n_bet',
        }
        return ops[op] if op in ops else '$eq'

    async def count(self):
        return await self.query_class.count_documents(self.filters)

    @staticmethod
    def generate_scope_filter(filters: list, general=False):
        if len(filters) == 1:
            _o = filters[0]
            if isinstance(_o, FilterItem):
                return MongoEngine.filter_item_generator(filters[0])
            elif isinstance(_o, FilterScope):
                return MongoEngine.generate_scope_filter(_o.items)
            else:
                return _o
        else:
            _a = filters.pop(0)
            _a = (MongoEngine.filter_item_generator(_a) if isinstance(_a,
                                                                      FilterItem) else _a) if not general else (
                MongoEngine.generate_scope_filter(_a.items) if isinstance(_a, FilterScope) else _a)
            _b = filters.pop(0)
            op = '$' + (_b.cr if not general else _b.op)
            _b = (MongoEngine.filter_item_generator(_b) if isinstance(_b,
                                                                      FilterItem) else _b) if not general else (
                MongoEngine.generate_scope_filter(_b.items) if isinstance(_b, FilterScope) else _b)
            filters = [{op: [
                _a, _b
            ]}] + filters
            return MongoEngine.generate_scope_filter(filters, general)

    def set_joins(self):
        pass

    async def generate_output(self, _f):
        print(self.input.e.ot)
        if self.input.e.ot == OutPutType.EXEL:
            file_name = str(uuid.uuid4().hex) + '.xlsx'
            return await self.output_generator(_f, file_name=file_name, cols=self.input.e.c)
        elif self.input.e.ot == OutPutType.JSON:
            pass
        elif self.input.e.ot == OutPutType.PDF:
            pass
        else:
            raise Exception("Output type not supported.")

    def _excel_temp_dir(self):
        """پوشه temp پروژه برای ذخیره فایل‌های اکسل (همان مسیری که برای دانلود سرو می‌شود)."""
        temp_dir = os.path.join(os.getcwd(), 'temp')
        os.makedirs(temp_dir, exist_ok=True)
        return temp_dir

    def _expand_headers_with_id_fields(self, cols: dict) -> list:
        """
        اگر design.nature در فیلدها باشد، design.nature_id هم اضافه می‌شود.
        اگر design.budget در فیلدها باشد، design.budget_id هم اضافه می‌شود.
        """
        headers = list(cols.keys())
        expanded = []
        pair_rules = [
            ('design.nature', 'design.nature_id', 'نام ماهیت'),
            ('design.budget', 'design.budget_id', 'نام دسته بودجه'),
        ]
        for h in headers:
            expanded.append(h)
            for field, id_field, default_title in pair_rules:
                if h == field and id_field not in headers and id_field not in expanded:
                    expanded.append(id_field)
                    cols[id_field] = cols.get(field, default_title) + ' (نام)'
                    break
        return expanded

    async def output_generator(self, query, file_name, cols):
        datas = query
        file_dir = os.path.join(self._excel_temp_dir(), file_name)
        head_cols = list(str(string.ascii_uppercase))
        headers = self._expand_headers_with_id_fields(cols)

        workbook = xlsxwriter.Workbook(file_dir)
        worksheet = workbook.add_worksheet()

        # نوشتن هدرها
        bold_format = workbook.add_format({'bold': True})

        # نوشتن هدرها
        for col_index, header in enumerate(headers):
            worksheet.write(0, col_index, cols[header], bold_format)

        # ساختن value_maps از vk (نام مپ) با استفاده از رجیستری
        value_maps = self._build_value_maps()
        
        # بارگذاری const ها برای فیلدهای خاص
        const_maps = await self._load_const_maps(headers)
        
        for row_index, item in enumerate(datas, start=1):
            for col_index, col in enumerate(headers):
                val = self.get_nested_value(item, col)
                val = self._apply_value_map(col, val, value_maps)
                val = await self._apply_const_map(col, val, const_maps)
                worksheet.write(row_index, col_index, val)
        workbook.close()
        return file_name

    def _build_value_maps(self):
        """اگر کاربر در e.vk برای ستونی کلید مپ داده (مثلاً design_sts)، مپ از رجیستری برمی‌گردد."""
        value_maps = {}
        vk = getattr(self.input.e, "vk", None) or {}
        if not isinstance(vk, dict):
            return value_maps
        for col_key, map_name in vk.items():
            if map_name and map_name in EXCEL_VALUE_MAP_REGISTRY:
                value_maps[col_key] = EXCEL_VALUE_MAP_REGISTRY[map_name]
        return value_maps

    def _apply_value_map(self, col_key: str, val, value_maps: dict):
        """اگر برای این ستون مپ تعریف شده، مقدار را به متن متناظر تبدیل می‌کند."""
        if not value_maps or col_key not in value_maps:
            return val
        mapping = value_maps[col_key]
        if not isinstance(mapping, dict):
            return val
        if val is None:
            return mapping.get(None, val)
        key = val if isinstance(val, (int, float)) else (int(val) if isinstance(val, str) and str(val).strip().isdigit() else val)
        if key in mapping:
            return mapping[key]
        if isinstance(key, int) and str(key) in mapping:
            return mapping[str(key)]
        return val

    async def _load_const_maps(self, headers: List[str]) -> dict:
        """
        شناسایی فیلدهای خاص (مثلاً nature_id, type_id, budget_id) و بارگذاری const ها از MySQL.
        Returns: dict به صورت {col_key: {id: title_fa}} مثلاً {"design.nature_id": {2: "طبیعت 2", ...}}
        """
        # لیست فیلدهای خاص که باید از const خوانده شوند
        const_field_names = ['nature_id', 'type_id', 'budget_id']  # می‌تونی اضافه کنی
        
        const_maps = {}
        const_keys_to_load = set()
        const_column_map = {}  # {col_key: const_key} مثلاً {"design.nature_id": "nature"}
        
        # شناسایی ستون‌هایی که باید const بخونن
        for col_key in headers:
            col_name = col_key.split('.')[-1]  # آخرین بخش (nature_id)
            if col_name in const_field_names:
                const_key = col_name.replace('_id', '')  # nature_id -> nature
                const_keys_to_load.add(const_key)
                const_column_map[col_key] = const_key
        
        # بارگذاری const ها از MySQL
        if const_keys_to_load:
            repo = MySQLConstRepository()
            for col_key, const_key in const_column_map.items():
                consts_for_key = await repo.get_by_key(const_key)
                if consts_for_key:
                    const_maps[col_key] = consts_for_key
        
        return const_maps

    async def _apply_const_map(self, col_key: str, val, const_maps: dict):
        """اگر این ستون در const_maps باشد، id را به title_fa تبدیل می‌کند."""
        if not const_maps or col_key not in const_maps:
            return val
        
        mapping = const_maps[col_key]
        if not isinstance(mapping, dict):
            return val
        
        if val is None:
            return mapping.get(None, val)
        
        # تبدیل val به int برای جستجو در mapping
        key = val if isinstance(val, (int, float)) else (int(val) if isinstance(val, str) and str(val).strip().isdigit() else val)
        
        if key in mapping:
            return mapping[key]
        if isinstance(key, int) and str(key) in mapping:
            return mapping[str(key)]
        
        return val

    def get_nested_value(self, data, path, default='-'):
        """
        دریافت مقدار از دیکشنری تو در تو با مسیر dot notation
        مثال: path = "design.title"
        """
        if not isinstance(path, str):
            return default

        keys = path.split('.')
        val = data

        for key in keys:
            if isinstance(val, dict):
                val = val.get(key, default)
            else:
                return default

            if val is None:
                return default

        # تبدیل datetime و UUID به رشته برای اکسل
        if isinstance(val, datetime.datetime):
            return self.create_date(val)
        if isinstance(val, uuid.UUID):
            return str(val)

        return val

    def create_date(self, date):
        shamsi_dt = jdatetime.datetime.fromgregorian(datetime=date)
        return shamsi_dt.strftime("%Y/%m/%d - %H:%M:%S")