# -*- coding: utf-8 -*-
"""
reset_db command

originally from https://www.djangosnippets.org/snippets/828/ by dnordberg
"""

import importlib.util
import os
import logging
import warnings

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.db import DEFAULT_DB_ALIAS

from django_extensions.settings import SQLITE_ENGINES, POSTGRESQL_ENGINES, MYSQL_ENGINES
from django_extensions.management.mysql import parse_mysql_cnf
from django_extensions.management.utils import signalcommand
from django_extensions.utils.deprecation import RemovedInNextVersionWarning


class Command(BaseCommand):
    help = "Resets the database for this project."

    def add_arguments(self, parser):
        super().add_arguments(parser)
        parser.add_argument(
            "--noinput",
            "--no-input",
            action="store_false",
            dest="interactive",
            default=True,
            help="Tells Django to NOT prompt the user for input of any kind.",
        )
        parser.add_argument(
            "--no-utf8",
            action="store_true",
            dest="no_utf8_support",
            default=False,
            help="Tells Django to not create a UTF-8 charset database",
        )
        parser.add_argument(
            "-U",
            "--user",
            action="store",
            dest="user",
            default=None,
            help="Use another user for the database than defined in settings.py",
        )
        parser.add_argument(
            "-O",
            "--owner",
            action="store",
            dest="owner",
            default=None,
            help="Use another owner for creating the database than the user defined "
            "in settings or via --user",
        )
        parser.add_argument(
            "-P",
            "--password",
            action="store",
            dest="password",
            default=None,
            help="Use another password for the database than defined in settings.py",
        )
        parser.add_argument(
            "-D",
            "--dbname",
            action="store",
            dest="dbname",
            default=None,
            help="Use another database name than defined in settings.py",
        )
        parser.add_argument(
            "-R",
            "--router",
            action="store",
            dest="router",
            default=DEFAULT_DB_ALIAS,
            help="Use this router-database other than defined in settings.py",
        )
        parser.add_argument(
            "--database",
            default=DEFAULT_DB_ALIAS,
            help='Nominates a database to run command for. Defaults to the "%s".'
            % DEFAULT_DB_ALIAS,
        )
        parser.add_argument(
            "-c",
            "--close-sessions",
            action="store_true",
            dest="close_sessions",
            default=False,
            help="Close database connections before dropping database "
            "(currently works on PostgreSQL only)",
        )

    @signalcommand
    def handle(self, *args, **options):
        """
        Reset the database for this project.

        Note: Transaction wrappers are in reverse as a work around for
        autocommit, anybody know how to do this the right way?
        """
        database = options["database"]
        if options["router"] != DEFAULT_DB_ALIAS:
            warnings.warn(
                "--router is deprecated. You should use --database.",
                RemovedInNextVersionWarning,
                stacklevel=2,
            )
            database = options["router"]

        dbinfo = settings.DATABASES.get(database)
        if dbinfo is None:
            raise CommandError("Unknown database %s" % database)

        engine = dbinfo.get("ENGINE")

        user = password = database_name = database_host = database_port = ""
        if engine == "mysql":
            (user, password, database_name, database_host, database_port) = (
                parse_mysql_cnf(dbinfo)
            )

        user = options["user"] or dbinfo.get("USER") or user
        password = options["password"] or dbinfo.get("PASSWORD") or password
        owner = options["owner"] or user

        database_name = options["dbname"] or dbinfo.get("NAME") or database_name
        if database_name == "":
            raise CommandError(
                "You need to specify DATABASE_NAME in your Django settings file."
            )

        database_host = dbinfo.get("HOST") or database_host
        database_port = dbinfo.get("PORT") or database_port

        verbosity = options["verbosity"]
        if options["interactive"]:
            confirm = input(
                """
You have requested a database reset.
This will IRREVERSIBLY DESTROY
ALL data in the database "%s".
Are you sure you want to do this?

Type 'yes' to continue, or 'no' to cancel: """
                % (database_name,)
            )
        else:
            confirm = "yes"

        if confirm != "yes":
            print("Reset cancelled.")
            return

        if engine in SQLITE_ENGINES:
            try:
                logging.info("Unlinking %s database", engine)
                os.unlink(database_name)
            except OSError:
                pass
        elif engine in MYSQL_ENGINES:
            import MySQLdb as Database

            kwargs = {
                "user": user,
                "passwd": password,
            }
            if database_host.startswith("/"):
                kwargs["unix_socket"] = database_host
            else:
                kwargs["host"] = database_host

            if database_port:
                kwargs["port"] = int(database_port)

            connection = Database.connect(**kwargs)
            drop_query = "DROP DATABASE IF EXISTS `%s`" % database_name
            utf8_support = "" if options["no_utf8_support"] else "CHARACTER SET utf8"
            create_query = "CREATE DATABASE `%s` %s" % (database_name, utf8_support)
            logging.info('Executing... "%s"', drop_query)
            connection.query(drop_query)
            logging.info('Executing... "%s"', create_query)
            connection.query(create_query.strip())
        elif engine in POSTGRESQL_ENGINES:
            has_psycopg3 = importlib.util.find_spec("psycopg")
            if has_psycopg3:
                import psycopg as Database  # NOQA
            else:
                import psycopg2 as Database  # NOQA

            conn_params = {"dbname": "template1"}
            if user:
                conn_params["user"] = user
            if password:
                conn_params["password"] = password
            if database_host:
                conn_params["host"] = database_host
            if database_port:
                conn_params["port"] = database_port

            connection = Database.connect(**conn_params)
            if has_psycopg3:
                connection.autocommit = True
            else:
                connection.set_isolation_level(0)  # autocommit false
            cursor = connection.cursor()

            if options["close_sessions"]:
                close_sessions_query = (
                    """
                    SELECT pg_terminate_backend(pg_stat_activity.pid)
                    FROM pg_stat_activity
                    WHERE pg_stat_activity.datname = '%s';
                """
                    % database_name
                )
                logging.info('Executing... "%s"', close_sessions_query.strip())
                try:
                    cursor.execute(close_sessions_query)
                except Database.ProgrammingError as e:
                    logging.exception("Error: %s", str(e))

            drop_query = 'DROP DATABASE "%s";' % database_name
            logging.info('Executing... "%s"', drop_query)
            try:
                cursor.execute(drop_query)
            except Database.ProgrammingError as e:
                logging.exception("Error: %s", str(e))

            create_query = 'CREATE DATABASE "%s"' % database_name
            if owner:
                create_query += ' WITH OWNER = "%s" ' % owner
            create_query += " ENCODING = 'UTF8'"

            if settings.DEFAULT_TABLESPACE:
                create_query += " TABLESPACE = %s;" % settings.DEFAULT_TABLESPACE
            else:
                create_query += ";"

            logging.info('Executing... "%s"', create_query)
            cursor.execute(create_query)
        else:
            raise CommandError("Unknown database engine %s" % engine)

        if verbosity >= 2 or options["interactive"]:
            print("Reset successful.")
