Django/gvsdsdk/fluent.py

from __future__ import annotations

import threading
import functools
from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union, overload

from django.db import connection as default_db_connection, models, transaction
from django.db.models import Q, F, Count, Sum, Avg, Min, Max, Subquery, OuterRef, Case, When, Value


_local = threading.local()


def _get_field_expression_cls():
    """延迟导入，避免 fluent ↔ model_base 循环依赖"""
    from gvsdsdk.model_base import FieldExpression
    return FieldExpression


class _AggregateExpr:
    def __init__(self, model, aggregate):
        self.model = model
        self.aggregate = aggregate


class _ExtractExpr:
    def __init__(self, field_name, lookup_type):
        self.field_name = field_name
        self.lookup_type = lookup_type

    def __eq__(self, other):
        return Q(**{f'{self.field_name}__{self.lookup_type}': other})

    def __and__(self, other):
        if isinstance(other, Q):
            return Q(**{f'{self.field_name}__{self.lookup_type}': None}) & other
        return self.__eq__(other)


class _CompatFunc:
    def sum(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            agg = Sum(field_expr._django_name)
        elif hasattr(field_expr, 'field_name'):
            agg = Sum(field_expr.field_name)
        else:
            agg = Sum(str(field_expr))
        model = getattr(field_expr, 'model_class', None) if hasattr(field_expr, 'model_class') else None
        return _AggregateExpr(model, agg) if model else agg

    def count(self, field_expr=None):
        if field_expr is None:
            return Count('*')
        if hasattr(field_expr, '_django_name'):
            agg = Count(field_expr._django_name)
        else:
            agg = Count(str(field_expr))
        model = getattr(field_expr, 'model_class', None) if hasattr(field_expr, 'model_class') else None
        return _AggregateExpr(model, agg) if model else agg

    def max(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            agg = Max(field_expr._django_name)
        else:
            agg = Max(str(field_expr))
        model = getattr(field_expr, 'model_class', None) if hasattr(field_expr, 'model_class') else None
        return _AggregateExpr(model, agg) if model else agg

    def min(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            agg = Min(field_expr._django_name)
        else:
            agg = Min(str(field_expr))
        model = getattr(field_expr, 'model_class', None) if hasattr(field_expr, 'model_class') else None
        return _AggregateExpr(model, agg) if model else agg

    def avg(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            agg = Avg(field_expr._django_name)
        else:
            agg = Avg(str(field_expr))
        model = getattr(field_expr, 'model_class', None) if hasattr(field_expr, 'model_class') else None
        return _AggregateExpr(model, agg) if model else agg

    def year(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            return _ExtractExpr(field_expr._django_name, 'year')
        return _ExtractExpr(str(field_expr), 'year')

    def month(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            return _ExtractExpr(field_expr._django_name, 'month')
        return _ExtractExpr(str(field_expr), 'month')

    def day(self, field_expr):
        if hasattr(field_expr, '_django_name'):
            return _ExtractExpr(field_expr._django_name, 'day')
        return _ExtractExpr(str(field_expr), 'day')


func = _CompatFunc()


class _DeleteStmt:
    def __init__(self, model_class, conditions=None):
        self.model_class = model_class
        self.conditions = conditions or Q()

    def where(self, *args, **kwargs):
        new_conditions = Q()
        for a in args:
            if isinstance(a, Q):
                new_conditions &= a
        self.conditions = new_conditions
        return self

    def execute(self):
        if self.model_class and self.conditions:
            self.model_class.objects.filter(self.conditions).delete()
        elif self.model_class:
            self.model_class.objects.all().delete()


def _get_session():
    if not hasattr(_local, 'db_session'):
        _local.db_session = Session()
    return _local.db_session


def _get_cache():
    if not hasattr(_local, 'query_cache'):
        _local.query_cache = {}
    return _local.query_cache


def _invalidate_cache():
    if hasattr(_local, 'query_cache'):
        _local.query_cache.clear()


class RelationAccessor:
    def __init__(self, model_class, prefix=''):
        self._model = model_class
        self._prefix = prefix

    def __getattr__(self, name):
        if name.startswith('_'):
            return object.__getattribute__(self, name)
        try:
            self._model._meta.get_field(name)
            full_name = f'{self._prefix}__{name}' if self._prefix else name
            return _get_field_expression_cls()(self._model, full_name)
        except Exception:
            for rel in self._model._meta.get_fields(include_hidden=True):
                if hasattr(rel, 'related_model') and rel.related_model:
                    rel_name = ''
                    if hasattr(rel, 'get_accessor_name'):
                        rel_name = rel.get_accessor_name()
                    elif hasattr(rel, 'name'):
                        rel_name = rel.name
                    if rel_name and rel_name == name:
                        full_prefix = f'{self._prefix}__{rel_name}' if self._prefix else rel_name
                        return RelationAccessor(rel.related_model, full_prefix)
            return _get_field_expression_cls()(self._model, name)


class _ValuesResult:
    """
    包装 Django ValuesQuerySet，支持链式调用。
    底层使用 Django .values() 返回字典，保持与 Django 行为一致。
    额外提供 .all() 返回字典列表，.first() 返回字典或 None，
    .scalar() 返回单字段场景的标量值。
    """
    def __init__(self, qs: models.QuerySet, fields: List[str], original_fields: Tuple) -> None:
        self._qs: models.QuerySet = qs          # Django ValuesQuerySet（返回字典）
        self._fields: List[str] = fields  # 转换后的字段名列表（传给 Django .values() 的）
        self._original_fields: Tuple = original_fields  # 原始传入的字段（用于判断单字段场景）

    @property
    def query(self) -> Any:
        """暴露底层 Django QuerySet 的 query 属性，支持 Subquery() 等场景。"""
        return self._qs.query

    def filter(self, *args: Any, **kwargs: Any) -> _ValuesResult:
        self._qs = self._qs.filter(*args, **kwargs)
        return self

    def exclude(self, *args: Any, **kwargs: Any) -> _ValuesResult:
        self._qs = self._qs.exclude(*args, **kwargs)
        return self

    def order_by(self, *fields: str) -> _ValuesResult:
        self._qs = self._qs.order_by(*fields)
        return self

    def distinct(self) -> _ValuesResult:
        self._qs = self._qs.distinct()
        return self

    def annotate(self, **kwargs: Any) -> _ValuesResult:
        self._qs = self._qs.annotate(**kwargs)
        return self

    def first(self) -> Optional[Dict[str, Any]]:
        return self._qs.first()

    def all(self) -> List[Dict[str, Any]]:
        return list(self._qs)

    def scalar(self) -> Any:
        row = self._qs.first()
        if row is None:
            return None
        if self._fields:
            return row.get(self._fields[0])
        return row

    def count(self) -> int:
        return self._qs.count()

    def exists(self) -> bool:
        return self._qs.exists()

    def __getitem__(self, key: Union[int, slice]) -> Any:
        if isinstance(key, slice):
            self._qs = self._qs[key]
            return self
        else:
            return self._qs[key]

    def __iter__(self) -> Iterator[Dict[str, Any]]:
        return iter(self._qs)

    def __len__(self) -> int:
        return self.count()

    def __bool__(self) -> bool:
        return self._qs.exists()


class FluentQuery:
    """QModel 兼容层查询构建器，提供 SQLAlchemy 风格的链式 OOP 查询接口。

    用法示例::

        # 基础查询
        Dingdan.query.filter(zhuangtai=8).all()

        # 链式过滤
        Dingdan.query.filter(jine__gt=100).order_by('-create_time').limit(10)

        # 字段表达式（通过 Model.field 自动获取 FieldExpression）
        Dingdan.query.filter(Dingdan.jine > 100, Dingdan.zhuangtai.in_([1,2,3]))

        # 聚合
        Dingdan.query.aggregate(total=Sum('jine'))

        # values / annotate
        Dingdan.query.values('zhuangtai').annotate(cnt=Count('id'))

        # 分页
        Dingdan.query.paginate(page=1, per_page=20)
    """

    def __init__(self, model_class: Type[models.Model]) -> None:
        self._model: Type[models.Model] = model_class
        self._qs: models.QuerySet = model_class.objects.all()
        self._joined_models: List[Type[models.Model]] = []
        self._cache_enabled: bool = True

    @property
    def query(self) -> Any:
        """暴露底层 Django QuerySet 的 query 属性，支持 Subquery() 等场景。"""
        return self._qs.query

    def _clone(self) -> FluentQuery:
        q = FluentQuery(self._model)
        q._qs = self._qs
        q._joined_models = list(self._joined_models)
        q._cache_enabled = self._cache_enabled
        return q

    def _cache_key(self, method: str) -> Tuple[Any, str, Any]:
        try:
            query_str = str(self._qs.query)
        except Exception:
            query_str = id(self._qs)
        return (self._model, method, query_str)

    def _remap_q(self, q: Q) -> Q:
        new_children: List[Any] = []
        for child in q.children:
            if isinstance(child, Q):
                new_children.append(self._remap_q(child))
            elif isinstance(child, tuple):
                key, val = child
                new_key = self._remap_key(key)
                new_children.append((new_key, val))
            else:
                new_children.append(child)
        if new_children:
            new_q = Q(*new_children)
            new_q.connector = q.connector
            new_q.negated = q.negated
            return new_q
        return q

    def _remap_key(self, key: str) -> str:
        parts = key.split('__')
        field_name = parts[0]
        try:
            self._model._meta.get_field(field_name)
            return key
        except Exception:
            pass
        for joined in self._joined_models:
            try:
                joined._meta.get_field(field_name)
                for rel in self._model._meta.get_fields(include_hidden=True):
                    if hasattr(rel, 'related_model') and rel.related_model == joined:
                        acc = ''
                        if hasattr(rel, 'get_accessor_name'):
                            acc = rel.get_accessor_name()
                        elif hasattr(rel, 'name'):
                            acc = rel.name
                        if acc and acc != '+':
                            return acc + '__' + key
                return joined._meta.model_name + '__' + key
            except Exception:
                pass
        return key

    def filter(self, *args: Any, **kwargs: Any) -> FluentQuery:
        """过滤查询，支持 Q 对象和关键字参数，返回 self 以支持链式调用。"""
        converted_args: List[Any] = []
        for arg in args:
            if isinstance(arg, Q) and self._joined_models:
                converted_args.append(self._remap_q(arg))
            else:
                converted_args.append(arg)
        self._qs = self._qs.filter(*converted_args, **kwargs)
        return self

    def exclude(self, *args: Any, **kwargs: Any) -> FluentQuery:
        """排除查询，返回 self 以支持链式调用。"""
        converted_args: List[Any] = []
        for arg in args:
            if isinstance(arg, Q) and self._joined_models:
                converted_args.append(self._remap_q(arg))
            else:
                converted_args.append(arg)
        self._qs = self._qs.exclude(*converted_args, **kwargs)
        return self

    def order_by(self, *args: Any) -> FluentQuery:
        """排序，支持字符串和 FieldExpression，返回 self 以支持链式调用。"""
        converted: List[str] = []
        for a in args:
            if isinstance(a, str):
                converted.append(a)
            elif isinstance(a, _get_field_expression_cls()):
                converted.append(str(a))
            else:
                converted.append(str(a))
        self._qs = self._qs.order_by(*converted)
        return self

    def first(self) -> Optional[models.Model]:
        """返回第一条记录或 None。"""
        if self._cache_enabled:
            cache = _get_cache()
            key = self._cache_key('first')
            if key in cache:
                return cache[key]
        obj = self._qs.first()
        if obj is not None:
            _get_session().track(obj)
        if self._cache_enabled:
            _get_cache()[self._cache_key('first')] = obj
        return obj

    def all(self) -> FluentQuery:
        """返回 self 以支持链式调用（与 Django QuerySet.all() 一致）。
        如需获取列表，使用 .to_list() 或 list()。"""
        self._qs = self._qs.all()
        return self

    def to_list(self) -> List[models.Model]:
        """执行查询并返回模型实例列表。"""
        if self._cache_enabled:
            cache = _get_cache()
            key = self._cache_key('to_list')
            if key in cache:
                return cache[key]
        results = list(self._qs)
        session = _get_session()
        for obj in results:
            session.track(obj)
        if self._cache_enabled:
            _get_cache()[self._cache_key('to_list')] = results
        return results

    def count(self) -> int:
        """返回匹配记录数。"""
        if self._cache_enabled:
            cache = _get_cache()
            key = self._cache_key('count')
            if key in cache:
                return cache[key]
        result = self._qs.count()
        if self._cache_enabled:
            _get_cache()[self._cache_key('count')] = result
        return result

    def get(self, *args: Any, **kwargs: Any) -> models.Model:
        """获取唯一匹配记录，不存在或多个时抛异常。"""
        if args:
            obj = self._qs.filter(*args).get(**kwargs)
        else:
            obj = self._qs.get(**kwargs)
        _get_session().track(obj)
        return obj

    def create(self, **kwargs: Any) -> models.Model:
        """创建并保存新记录。"""
        obj = self._model.objects.create(**kwargs)
        _invalidate_cache()
        return obj

    def get_or_create(self, defaults: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Tuple[models.Model, bool]:
        """查询或创建记录，返回 (obj, created) 元组。"""
        obj, created = self._model.objects.get_or_create(defaults=defaults, **kwargs)
        if created:
            _invalidate_cache()
        _get_session().track(obj)
        return obj, created

    def update_or_create(self, defaults: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Tuple[models.Model, bool]:
        """更新或创建记录，返回 (obj, created) 元组。"""
        obj, created = self._model.objects.update_or_create(defaults=defaults, **kwargs)
        _invalidate_cache()
        _get_session().track(obj)
        return obj, created

    def none(self) -> FluentQuery:
        """返回空结果集，返回 self 以支持链式调用。"""
        self._qs = self._qs.none()
        return self

    def last(self) -> Optional[models.Model]:
        """返回最后一条记录或 None。"""
        return self._qs.last()

    def limit(self, n: int) -> FluentQuery:
        """限制返回记录数，返回 self 以支持链式调用。"""
        self._qs = self._qs[:n]
        return self

    def offset(self, n: int) -> FluentQuery:
        """跳过前 n 条记录，返回 self 以支持链式调用。"""
        self._qs = self._qs[n:]
        return self

    def join(self, *args: Type[models.Model]) -> FluentQuery:
        """关联模型，用于跨表查询自动映射字段，返回 self 以支持链式调用。"""
        for arg in args:
            if isinstance(arg, type) and issubclass(arg, models.Model):
                if arg not in self._joined_models:
                    self._joined_models.append(arg)
        return self

    def select_related(self, *fields: str) -> FluentQuery:
        """外键关联查询优化，返回 self 以支持链式调用。"""
        self._qs = self._qs.select_related(*fields)
        return self

    def prefetch_related(self, *fields: str) -> FluentQuery:
        """多对多/反向关联查询优化，返回 self 以支持链式调用。"""
        self._qs = self._qs.prefetch_related(*fields)
        return self

    def only(self, *fields: str) -> FluentQuery:
        """只加载指定字段，返回 self 以支持链式调用。"""
        self._qs = self._qs.only(*fields)
        return self

    def defer(self, *fields: str) -> FluentQuery:
        """延迟加载指定字段，返回 self 以支持链式调用。"""
        self._qs = self._qs.defer(*fields)
        return self

    def select_for_update(self, nowait: bool = False, skip_locked: bool = False,
                          of: Tuple = (), no_key: bool = False) -> FluentQuery:
        """行级锁，返回 self 以支持链式调用。"""
        self._qs = self._qs.select_for_update(
            nowait=nowait, skip_locked=skip_locked, of=of, no_key=no_key
        )
        return self

    def bulk_create(self, objs: List[models.Model], batch_size: Optional[int] = None,
                    ignore_conflicts: bool = False) -> List[models.Model]:
        """批量创建记录。"""
        result = self._model.objects.bulk_create(
            objs, batch_size=batch_size, ignore_conflicts=ignore_conflicts
        )
        _invalidate_cache()
        return result

    def annotate(self, **kwargs: Any) -> FluentQuery:
        """注解查询，添加聚合/计算字段，返回 self 以支持链式调用。"""
        self._qs = self._qs.annotate(**kwargs)
        return self

    def aggregate(self, **kwargs: Any) -> Dict[str, Any]:
        """聚合查询，返回聚合结果字典。"""
        return self._qs.aggregate(**kwargs)

    def values(self, *fields: Union[str, _AggregateExpr, models.Expression]) -> _ValuesResult:
        """返回指定字段的字典结果，支持链式调用。"""
        if not fields:
            # 无参数：返回 Django .values()（所有字段字典）
            return _ValuesResult(self._qs.values(), [], fields)

        annotations: Dict[str, models.Aggregate] = {}
        converted: List[str] = []
        for i, f in enumerate(fields):
            if isinstance(f, _AggregateExpr):
                annotations[f'val{i}'] = f.aggregate
                converted.append(f'val{i}')
            elif isinstance(f, (models.Aggregate, models.Expression)):
                annotations[f'val{i}'] = f
                converted.append(f'val{i}')
            elif isinstance(f, str):
                converted.append(f)
            else:
                converted.append(str(f))

        qs = self._qs
        if annotations:
            qs = qs.annotate(**annotations)
        # 使用 Django 原生 .values() 返回字典
        qs = qs.values(*converted)
        return _ValuesResult(qs, converted, fields)

    def with_entities(self, *fields: Union[str, _AggregateExpr, models.Expression]) -> _ValuesResult:
        """values() 的别名，SQLAlchemy 风格。"""
        return self.values(*fields)

    def params(self, **kwargs: Any) -> FluentQuery:
        """参数占位（SQLAlchemy 兼容），返回 self 以支持链式调用。"""
        return self

    def values_list(self, *fields: str, flat: bool = False) -> Any:
        """返回指定字段的元组列表。"""
        return self._qs.values_list(*fields, flat=flat)

    def update(self, **kwargs: Any) -> int:
        """批量更新匹配记录，返回受影响行数。"""
        result = self._qs.update(**kwargs)
        _invalidate_cache()
        return result

    def delete(self) -> Tuple[int, Dict[str, int]]:
        """批量删除匹配记录，返回 (删除数, {表名: 删除数})。"""
        result = self._qs.delete()
        _invalidate_cache()
        return result

    def exists(self) -> bool:
        """判断是否存在匹配记录。"""
        return self._qs.exists()

    def distinct(self) -> FluentQuery:
        """去重查询，返回 self 以支持链式调用。"""
        self._qs = self._qs.distinct()
        return self

    def paginate(self, page: int = 1, per_page: int = 20, error_out: bool = False) -> Any:
        """分页查询，返回 Page 对象（含 items, total, pages 等属性）。"""
        from django.core.paginator import Paginator
        if not self._qs.ordered:
            meta = getattr(self._model, '_meta', None)
            if meta and meta.pk:
                self._qs = self._qs.order_by(meta.pk.name)
        paginator = Paginator(self._qs, per_page)
        try:
            page_obj = paginator.page(page)
        except Exception:
            page_obj = paginator.page(1)
        page_obj.total = paginator.count
        page_obj.items = page_obj.object_list
        page_obj.pages = paginator.num_pages
        page_obj.page = page
        page_obj.per_page = per_page
        page_obj.has_next = page_obj.has_next()
        page_obj.has_prev = page_obj.has_previous()
        page_obj.next_num = page_obj.next_page_number() if page_obj.has_next else None
        page_obj.prev_num = page_obj.previous_page_number() if page_obj.has_prev else None
        return page_obj

    def no_cache(self) -> FluentQuery:
        """禁用查询缓存，返回 self 以支持链式调用。"""
        self._cache_enabled = False
        return self

    def using(self, alias: str) -> FluentQuery:
        """指定数据库别名，返回 self 以支持链式调用。"""
        self._qs = self._qs.using(alias)
        return self

    def __getitem__(self, key: Union[int, slice]) -> Any:
        """支持切片 [start:end] 和索引 [n]。切片返回 self，索引返回模型实例。"""
        if isinstance(key, slice):
            self._qs = self._qs[key]
            return self
        else:
            # 单个索引，返回模型实例
            return self._qs[key]

    def __iter__(self) -> Iterator[models.Model]:
        return iter(self._qs)

    def __len__(self) -> int:
        return self.count()

    def __bool__(self) -> bool:
        return self.exists()

    def __repr__(self) -> str:
        return f'FluentQuery({self._model.__name__})'


class Session:
    def __init__(self):
        self._pending_adds = []
        self._pending_deletes = []
        self._tracked = set()

    def add(self, obj):
        self._pending_adds.append(obj)

    def delete(self, obj):
        self._pending_deletes.append(obj)
        self._tracked.discard(obj)

    def track(self, obj):
        if isinstance(obj, models.Model):
            self._tracked.add(obj)

    def commit(self):
        delete_pks = {(obj.__class__, obj.pk) for obj in self._pending_deletes if obj.pk is not None}
        with transaction.atomic():
            for obj in self._pending_deletes:
                obj.delete()
            for obj in self._pending_adds:
                obj.save()
            for obj in self._tracked:
                if obj.pk is not None and (obj.__class__, obj.pk) in delete_pks:
                    continue
                try:
                    obj.save()
                except Exception:
                    pass
        self._pending_adds.clear()
        self._pending_deletes.clear()
        self._tracked.clear()
        _invalidate_cache()

    def rollback(self):
        self._pending_adds.clear()
        self._pending_deletes.clear()
        self._tracked.clear()
        _invalidate_cache()

    def query(self, *args):
        model = None
        aggregates = []
        for a in args:
            if isinstance(a, type) and issubclass(a, models.Model):
                model = a
            elif isinstance(a, _AggregateExpr):
                if a.model:
                    model = a.model
                aggregates.append(a.aggregate)
            elif isinstance(a, (models.Aggregate, models.Expression)):
                aggregates.append(a)
        if not aggregates and model:
            return FluentQuery(model)
        if aggregates and model:
            fq = FluentQuery(model)
            return fq.values(*args)
        return None

    def flush(self):
        self.commit()

    def execute(self, stmt):
        if isinstance(stmt, _DeleteStmt):
            stmt.execute()
        elif isinstance(stmt, str):
            connection = default_db_connection
            with connection.cursor() as cursor:
                cursor.execute(stmt)


class DB:
    @property
    def session(self):
        return _get_session()

    @staticmethod
    def or_(*args):
        result = None
        for q in args:
            if isinstance(q, Q):
                if result is None:
                    result = q
                else:
                    result |= q
        return result if result is not None else Q()

    @staticmethod
    def and_(*args):
        result = Q()
        for q in args:
            if isinstance(q, Q):
                result &= q
        return result


db = DB()


def FQ(model_class):
    return FluentQuery(model_class)