from django.db import connection, migrations, models
from django.db.migrations.state import ProjectState
from django.db.migrations.writer import MigrationWriter
from django.test import TestCase
from django.test.utils import isolate_apps

from tsvector_field import SearchVectorField, WeightedColumn
from tsvector_field.receivers import inject_trigger_operations
from tsvector_field.schema import DatabaseSchemaEditor, DatabaseTriggerEditor
from tsvector_field.operations import IndexSearchVector


@isolate_apps('tsvector_test_app')
class MigrationWriterTests(TestCase):

    def test_deconstruct_with_no_arguments(self):
        svf = SearchVectorField()
        self.assertEqual(
            ("tsvector_field.SearchVectorField()",
             {'import tsvector_field'}),
            MigrationWriter.serialize(svf)
        )

    def test_deconstruct_default_arguments(self):

        svf = SearchVectorField([
            WeightedColumn('name', 'A'),
            WeightedColumn('description', 'D'),
        ], language=None, language_column=None, force_update=False)

        definition, path = MigrationWriter.serialize(svf)

        self.assertEqual(
            "tsvector_field.SearchVectorField("
            "columns=["
            "tsvector_field.WeightedColumn('name', 'A'), "
            "tsvector_field.WeightedColumn('description', 'D')]"
            ")",
            definition
        )

        self.assertSetEqual(
            {'import tsvector_field'},
            path
        )

    def test_deconstruct_all_arguments(self):

        class TextDocument(models.Model):
            svf = SearchVectorField([
                WeightedColumn('name', 'A'),
                WeightedColumn('description', 'D'),
            ], language='english', language_column='lang', force_update=True)

        name, path, args, kwargs = TextDocument._meta.get_field('svf').deconstruct()

        self.assertEqual(name, "svf")
        self.assertEqual(path, "tsvector_field.SearchVectorField")
        self.assertFalse(args)
        self.assertSetEqual(set(kwargs.keys()), {
            'columns', 'language', 'language_column', 'force_update'
        })


@isolate_apps('tsvector_test_app')
class SchemaEditorTests(TestCase):

    @property
    def trigger_editor(self):
        with DatabaseSchemaEditor(connection) as schema_editor:
            return DatabaseTriggerEditor(schema_editor)

    def test_sql_setweight(self):

        def check_sql(model, sql):
            field = model._meta.get_field('search')
            self.assertEqual(
                sql, self.trigger_editor._to_tsvector_weights(field)
            )

        class WithLanguageTwoColumn(models.Model):
            search = SearchVectorField([
                WeightedColumn('title', 'A'),
                WeightedColumn('body', 'D'),
            ], language='ukrainian')

        class WithLanguage(models.Model):
            search = SearchVectorField([
                WeightedColumn('body', 'D'),
            ], language='ukrainian')

        class WithLanguageColumn(models.Model):
            search = SearchVectorField([
                WeightedColumn('body', 'D'),
            ], language_column='lang')

        class WithLanguageAndLanguageColumn(models.Model):
            search = SearchVectorField([
                WeightedColumn('body', 'D'),
            ], language='ukrainian', language_column='lang')

        check_sql(
            WithLanguageTwoColumn, [
                """ setweight(to_tsvector('ukrainian', COALESCE(NEW."title", '')), 'A') ||""",
                """ setweight(to_tsvector('ukrainian', COALESCE(NEW."body", '')), 'D');"""
            ]
        )

        check_sql(
            WithLanguage, [
                """ setweight(to_tsvector('ukrainian', COALESCE(NEW."body", '')), 'D');"""
            ]
        )

        check_sql(
            WithLanguageColumn, [
                """ setweight(to_tsvector(NEW."lang"::regconfig, COALESCE(NEW."body", '')), 'D');"""
            ]
        )

        check_sql(
            WithLanguageAndLanguageColumn, [
                """ setweight(to_tsvector(COALESCE(NEW."lang"::regconfig, 'ukrainian'),"""
                """ COALESCE(NEW."body", '')), 'D');"""
            ]
        )

    def test_sql_update_column_checks(self):

        def check_sql(model, sql):
            field = model._meta.get_field('search')
            self.assertEqual(
                sql, list(self.trigger_editor._to_tsvector_preconditions(field))
            )

        class OneColumn(models.Model):
            search = SearchVectorField([
                WeightedColumn('name', 'A'),
            ])

        class ThreeColumns(models.Model):
            search = SearchVectorField([
                WeightedColumn('name', 'A'),
                WeightedColumn('title', 'B'),
                WeightedColumn('body', 'C'),
            ])

        check_sql(
            OneColumn, [
                "IF (TG_OP = 'INSERT') THEN do_update = true;",
                "ELSIF (TG_OP = 'UPDATE') THEN",
                ' IF (NEW."search" IS NULL) THEN do_update = true;',
                ' ELSIF (NEW."name" IS DISTINCT FROM OLD."name") THEN do_update = true;',
                ' END IF;',
                'END IF;',
            ]
        )

        check_sql(
            ThreeColumns, [
                "IF (TG_OP = 'INSERT') THEN do_update = true;",
                "ELSIF (TG_OP = 'UPDATE') THEN",
                ' IF (NEW."search" IS NULL) THEN do_update = true;',
                ' ELSIF (NEW."name" IS DISTINCT FROM OLD."name") THEN do_update = true;',
                ' ELSIF (NEW."title" IS DISTINCT FROM OLD."title") THEN do_update = true;',
                ' ELSIF (NEW."body" IS DISTINCT FROM OLD."body") THEN do_update = true;',
                ' END IF;',
                'END IF;',
            ]
        )

    def test_sql_update_function(self):

        def check_sql(model, sql):
            field = model._meta.get_field('search')
            self.assertEqual(
                sql, self.trigger_editor._create_function('thefunction()', field)
            )

        class TextDocument(models.Model):
            search = SearchVectorField([
                WeightedColumn('title', 'A'),
                WeightedColumn('body', 'D'),
            ], 'english')

        self.maxDiff = None
        check_sql(
            TextDocument,
            "CREATE FUNCTION thefunction() RETURNS trigger AS $$\n"
            "DECLARE\n"
            " do_update bool default false;\n"
            "BEGIN\n"
            " IF (TG_OP = 'INSERT') THEN do_update = true;\n"
            " ELSIF (TG_OP = 'UPDATE') THEN\n"
            '  IF (NEW."search" IS NULL) THEN do_update = true;\n'
            '  ELSIF (NEW."title" IS DISTINCT FROM OLD."title") THEN do_update = true;\n'
            '  ELSIF (NEW."body" IS DISTINCT FROM OLD."body") THEN do_update = true;\n'
            "  END IF;\n"
            " END IF;\n"
            " IF do_update THEN\n"
            '  NEW."search" :=\n'
            "   setweight(to_tsvector('english', COALESCE(NEW.\"title\", '')), 'A') ||\n"
            "   setweight(to_tsvector('english', COALESCE(NEW.\"body\", '')), 'D');\n"
            " END IF;\n"
            " RETURN NEW;\n"
            "END\n"
            "$$ LANGUAGE plpgsql"
        )

        class TextDocumentForceUpdate(models.Model):
            search = SearchVectorField([
                WeightedColumn('body', 'D'),
            ], 'english', force_update=True)

        check_sql(
            TextDocumentForceUpdate,
            "CREATE FUNCTION thefunction() RETURNS trigger AS $$\n"
            "DECLARE\n"
            " do_update bool default false;\n"
            "BEGIN\n"
            " do_update = true;\n"
            " IF do_update THEN\n"
            '  NEW."search" :=\n'
            "   setweight(to_tsvector('english', COALESCE(NEW.\"body\", '')), 'D');\n"
            " END IF;\n"
            " RETURN NEW;\n"
            "END\n"
            "$$ LANGUAGE plpgsql"
        )

    def test_create_model_no_function(self):

        class NoWeightedColumns(models.Model):
            search = SearchVectorField()

        with DatabaseSchemaEditor(connection) as schema_editor:
            schema_editor.create_model(NoWeightedColumns)
            self.assertEqual(len(schema_editor.deferred_sql), 1)
            self.assertIn('CREATE INDEX', schema_editor.deferred_sql[0])
            self.assertIn('GIN', schema_editor.deferred_sql[0])

    def test_create_model(self):

        class TextDocument(models.Model):
            title = models.CharField(max_length=128)
            search = SearchVectorField([
                WeightedColumn('title', 'A'),
            ], 'english')

        with DatabaseSchemaEditor(connection) as schema_editor:
            schema_editor.create_model(TextDocument)
            self.assertEqual(len(schema_editor.deferred_sql), 3)
            self.assertIn('CREATE INDEX', schema_editor.deferred_sql[0])
            self.assertIn('GIN', schema_editor.deferred_sql[0])
            self.assertIn('CREATE FUNCTION', schema_editor.deferred_sql[1])
            self.assertIn('CREATE TRIGGER', schema_editor.deferred_sql[2])


class MigrationTestCase(TestCase):

    def migrate(self, ops, state=None):
        class Migration(migrations.Migration):