from django.contrib.postgres import lookups
from django.contrib.postgres.fields.array import ArrayContainedBy, ArrayContains, \
    ArrayField
from django.db import models


class PgLtreeField(models.TextField):
    description = 'Ltree'
    
    def db_type(self, connection):
        return 'ltree'
    
    def get_internal_type(self):
        return "CharField"
    
    def get_prep_value(self, value):
        if isinstance(value, list) or isinstance(value, tuple):
            return [self.to_python(val) for val in value]
        return self.to_python(value)


class LtreeLookup(lookups.PostgresSimpleLookup):
    def as_sql(self, qn, connection):
        sql, params = super().as_sql(qn, connection)
        sql = '%s::%s' % (sql, self.lhs.output_field.db_type(connection))
        if isinstance(self.rhs, list) or isinstance(self.rhs, tuple):
            # Cast to array
            sql += '[]'
        return sql, params


@PgLtreeField.register_lookup
class AncestorOf(LtreeLookup):
    lookup_name = 'ancestor_of'
    operator = '@>'


@PgLtreeField.register_lookup
class DescendantOf(LtreeLookup):
    lookup_name = 'descendant_of'
    operator = '<@'


@PgLtreeField.register_lookup
class Match(LtreeLookup):
    lookup_name = 'match'
    operator = '~'


class LtreeArrayField(ArrayField):
    def __init__(self, **kwargs):
        kwargs['base_field'] = PgLtreeField()
        super().__init__(**kwargs)
    
    def get_db_prep_value(self, value, connection, prepared=False):
        if value:
            return '{%s}' % ','.join(value)
        return value
    
    def from_db_value(self, value, expression, connection, context):
        if value is None:
            return value
        # value example: '{1.2,1.3,1.4}'
        # remove brackets {}, then split by comma
        # return ['1.2', '1.3', '1.4']
        return [] if value == '{}' else value[1:-1].split(',')


@LtreeArrayField.register_lookup
class ArrayAncestorOf(ArrayContainedBy):
    lookup_name = 'ancestor_of'


@LtreeArrayField.register_lookup
class ArrayDescendantOf(ArrayContains):
    lookup_name = 'descendant_of'


@LtreeArrayField.register_lookup
class ArrayMatch(lookups.PostgresSimpleLookup):
    lookup_name = 'match'
    operator = '~'
    
    def as_sql(self, qn, connection):
        sql, params = super().as_sql(qn, connection)
        sql = '%s::%s' % (sql, self.lhs.output_field.db_type(connection))
        return sql, params

@@ -649,37 +649,41 @@
 ction):%0A        
-sql, 
+lhs, lhs_
 params = super()
@@ -672,37 +672,39 @@
 s_params = s
-uper().as_sql
+elf.process_lhs
 (qn, connect
@@ -720,61 +720,47 @@
     
-sql = '%25s::%25s' %25 (sql, self.lhs.output_field.db_type(
+rhs, rhs_params = self.process_rhs(qn, 
 conn
@@ -758,33 +758,32 @@
 (qn, connection)
-)
 %0A        if isin
@@ -881,40 +881,251 @@
     
-sql += '%5B%5D'%0A        return sql, 
+# With any clause GIST index will work%0A            rhs = 'ANY(%25s::%25s%5B%5D)' %25 (%0A                rhs,%0A                self.lhs.output_field.db_type(connection)%0A            )%0A%0A        return '%25s %25s %25s' %25 (lhs, self.operator, rhs), lhs_params + rhs_
 para

use ANY clause with <@ ltree[] queries