import json
from idm.core.management.base import IdmBaseCommand
from django.db import connections

SQL_BLOCKING = """
SELECT queries.blocking_query AS blocking_query,
       queries.blocking_pid AS blocking_pid,
       queries.blocking_mode AS blocking_mode,
       max(queries.blocking_time) AS blocking_time,
       count(queries.blocked_query) AS query_count,
       max(queries.blocked_time) AS max_blocked_time
  FROM (
       SELECT 
              blockeda.query AS blocked_query,
              blockinga.pid AS blocking_pid,
              blockinga.query AS blocking_query,
              blockingl.mode AS blocking_mode,
              now()-blockeda.query_start AS blocked_time,
              now()-blockinga.query_start AS blocking_time
         FROM pg_catalog.pg_locks blockedl
         JOIN pg_stat_activity blockeda ON blockedl.pid = blockeda.pid
         JOIN pg_catalog.pg_locks blockingl
              ON (
                     (
                         (blockingl.transactionid=blockedl.transactionid)
                         OR (blockingl.relation=blockedl.relation AND blockingl.locktype=blockedl.locktype)
                     )
                     AND blockedl.pid != blockingl.pid
              )
         JOIN pg_stat_activity blockinga ON blockingl.pid = blockinga.pid
        WHERE NOT blockedl.granted
  ) AS queries
GROUP BY queries.blocking_query, queries.blocking_pid, queries.blocking_mode
ORDER BY {order} desc
LIMIT {limit}
;
"""


SQL_BLOCKED = """
SELECT COALESCE(blockingl.relation::regclass::text,blockingl.locktype) as locked_item,
       blockeda.pid AS blocked_pid,
       blockeda.query AS blocked_query,
       blockedl.mode AS blocked_mode,
       blockinga.pid AS blocking_pid,
       blockinga.query AS blocking_query,
       blockingl.mode AS blocking_mode,
       now()-blockeda.query_start AS blocked_time,
       now()-blockinga.query_start AS blocking_time
  FROM pg_catalog.pg_locks blockedl
  JOIN pg_stat_activity blockeda ON blockedl.pid = blockeda.pid
  JOIN pg_catalog.pg_locks blockingl
       ON (
              (
                  (blockingl.transactionid=blockedl.transactionid)
                  OR (blockingl.relation=blockedl.relation AND blockingl.locktype=blockedl.locktype)
              )
              AND blockedl.pid != blockingl.pid
       )
  JOIN pg_stat_activity blockinga ON blockingl.pid = blockinga.pid
 WHERE NOT blockedl.granted {pid}
 ORDER BY blocked_time desc
 LIMIT {limit}
;
"""

class Command(IdmBaseCommand):
    ORDER = ('count', 'time',)
    ORDER_SUB = {'count': 'query_count', 'time': 'blocking_time'}
    QUERIES = ('blocking', 'blocked')
    BLOCKING_FIELDS = ['blocking_query', 'blocking_pid', 'blocking_mode',
                       'blocking_time', 'blocking_count', 'max_blocked_time']
    BLOCKED_FIELDS = ['locked_item', 'blocked_pid', 'blocked_query', 'blocked_mode',
                      'blocking_pid', 'blocking_query', 'blocking_mode', 'blocked_time', 'blocking_time']

    def add_arguments(self, parser):
        super(Command, self).add_arguments(parser)
        parser.add_argument(
            '--query-type',
            action='store',
            dest='query_type',
            choices=self.QUERIES,
            required=False,
            default='blocking',
            help='Показать блокирующиеся или заблокированные запросы'
        )
        parser.add_argument(
            '--order',
            action='store',
            dest='order',
            choices=self.ORDER,
            required=False,
            default='count',
            help='Сортировка блокирующих зарпосов'
        )
        parser.add_argument(
            '--blocking-pid',
            type='int',
            dest='pid',
            required=False,
            help='Фильтрация заблокированных запросов по блокирующему pid',
        )
        parser.add_argument(
            '--limit',
            type='int',
            dest='limit',
            required=False,
            default=10,
            help='Ограничение размера результата'
        )
        parser.add_argument(
            '--dbuser',
            dest='dbuser',
            required=False,
            default='default',
            help='Имя пользователя базы'
        )

    def get_sql(self, **options: dict) -> str:
        if options.get('query_type') == 'blocking':
            return SQL_BLOCKING.format(order=self.ORDER_SUB[options.get('order')], limit=options.get('limit'))
        else:
            if options.get('pid'):
                return SQL_BLOCKED.format(pid='AND blockinga.pid={}'.format(options['pid']), limit=options.get('limit'))
            else:
                return SQL_BLOCKED.format(pid='', limit=options.get('limit'))

    def format_record(self, line: list) -> list:
        return [str(element) for element in line]

    def get_result(self, **options: dict) -> list:
        sql = self.get_sql(**options)
        dbuser = options.get('dbuser')
        with connections[dbuser].cursor() as cursor:
            cursor.execute(sql)
            result = cursor.fetchall()
        if options.get('query_type') == 'blocking':
            return [dict(zip(self.BLOCKING_FIELDS, self.format_record(line))) for line in result]
        else:
            return [dict(zip(self.BLOCKED_FIELDS, self.format_record(line))) for line in result]

    def idm_handle(self, *args, **options):
        print(json.dumps(self.get_result(**options), indent=4))
