# Copyright (c) 2023, JINGROW
# For license information, please see license.txt

import json
from collections import Counter, defaultdict

import jingrow
from jingrow import _
from jingrow.core.pagetype.recorder.db_optimizer import DBOptimizer, DBTable
from jingrow.custom.pagetype.property_setter.property_setter import make_property_setter
from jingrow.model.page import Page
from jingrow.recorder import RECORDER_REQUEST_HASH
from jingrow.recorder import get as get_recorder_data
from jingrow.utils import cint, cstr, evaluate_filters, get_table_name
from jingrow.utils.caching import redis_cache


class Recorder(Page):
	# begin: auto-generated types
	# This code is auto-generated. Do not modify anything in this block.

	from typing import TYPE_CHECKING

	if TYPE_CHECKING:
		from jingrow.core.pagetype.recorder_query.recorder_query import RecorderQuery
		from jingrow.core.pagetype.recorder_suggested_index.recorder_suggested_index import (
			RecorderSuggestedIndex,
		)
		from jingrow.types import DF

		cmd: DF.Data | None
		duration: DF.Float
		event_type: DF.Data | None
		form_dict: DF.Code | None
		method: DF.Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]
		number_of_queries: DF.Int
		path: DF.Data | None
		profile: DF.Code | None
		request_headers: DF.Code | None
		sql_queries: DF.Table[RecorderQuery]
		suggested_indexes: DF.Table[RecorderSuggestedIndex]
		time: DF.Datetime | None
		time_in_queries: DF.Float
	# end: auto-generated types

	def load_from_db(self):
		request_data = get_recorder_data(self.name)
		if not request_data:
			raise jingrow.DoesNotExistError
		request = serialize_request(request_data)
		super(Page, self).__init__(request)

	@staticmethod
	def get_list(args):
		start = cint(args.get("start")) or 0
		page_length = cint(args.get("page_length")) or 20
		requests = Recorder.get_filtered_requests(args)[start : start + page_length]

		if order_by_statment := args.get("order_by"):
			if "." in order_by_statment:
				order_by_statment = order_by_statment.split(".")[1]

			if " " in order_by_statment:
				sort_key, sort_order = order_by_statment.split(" ", 1)
			else:
				sort_key = order_by_statment
				sort_order = "desc"

			sort_key = sort_key.replace("`", "")
			return sorted(requests, key=lambda r: r.get(sort_key) or 0, reverse=bool(sort_order == "desc"))

		return sorted(requests, key=lambda r: r.duration, reverse=1)

	@staticmethod
	def get_count(args):
		return len(Recorder.get_filtered_requests(args))

	@staticmethod
	def get_filtered_requests(args):
		filters = args.get("filters")
		requests = [serialize_request(request) for request in get_recorder_data()]
		return [req for req in requests if evaluate_filters(req, filters)]

	@staticmethod
	def get_stats(args):
		pass

	@staticmethod
	def delete(self):
		pass

	def db_insert(self, *args, **kwargs):
		pass

	def db_update(self):
		pass


def serialize_request(request):
	request = jingrow._dict(request)
	if request.get("calls"):
		for i in request.calls:
			i["stack"] = jingrow.as_json(i["stack"])
			i["explain_result"] = jingrow.as_json(i["explain_result"])
	request.update(
		name=request.get("uuid"),
		number_of_queries=request.get("queries"),
		time_in_queries=request.get("time_queries"),
		request_headers=jingrow.as_json(request.get("headers"), indent=4),
		form_dict=jingrow.as_json(request.get("form_dict"), indent=4),
		sql_queries=request.get("calls"),
		suggested_indexes=request.get("suggested_indexes"),
		modified=request.get("time"),
		creation=request.get("time"),
	)

	return request


@jingrow.whitelist()
def add_indexes(indexes):
	jingrow.only_for("Administrator")
	indexes = json.loads(indexes)

	for index in indexes:
		jingrow.enqueue(_add_index, table=index["table"], column=index["column"])
	jingrow.msgprint(_("Enqueued creation of indexes"), alert=True)


def _add_index(table, column):
	pagetype = get_pagetype_name(table)
	jingrow.db.add_index(pagetype, [column])
	jingrow.msgprint(
		_("Index created successfully on column {0} of pagetype {1}").format(column, pagetype),
		alert=True,
		realtime=True,
	)


@jingrow.whitelist()
def optimize(recorder_id: str):
	jingrow.only_for("Administrator")
	jingrow.enqueue(_optimize, recorder_id=recorder_id, queue="long")


def _optimize(recorder_id):
	record: Recorder = jingrow.get_pg("Recorder", recorder_id)
	total_duration = record.time_in_queries

	# Any index with query time less than 5% of total time is not suggested
	PERCENT_DURATION_THRESHOLD_OVERALL = 0.05
	# Any query with duration less than 0.5% of total duration is not analyzed
	PERCENT_DURATION_THRESHOLD_QUERY = 0.005

	# Index suggestion -> Query duration
	index_suggestions = Counter()
	for idx, captured_query in enumerate(record.sql_queries, start=1):
		query = cstr(captured_query.query)
		jingrow.publish_progress(
			idx / len(record.sql_queries) * 100,
			title="Analyzing Queries",
			pagetype=record.pagetype,
			pagename=record.name,
			description=f"Analyzing query: {query[:140]}",
		)
		if captured_query.duration < total_duration * PERCENT_DURATION_THRESHOLD_QUERY:
			continue
		if not query.lower().strip().startswith(("select", "update", "delete")):
			continue
		if index := _optimize_query(query):
			index_suggestions[(index.table, index.column)] += captured_query.duration

	suggested_indexes = index_suggestions.most_common(3)
	suggested_indexes = [
		idx for idx in suggested_indexes if idx[1] > total_duration * PERCENT_DURATION_THRESHOLD_OVERALL
	]

	if not suggested_indexes:
		jingrow.msgprint(
			_("No automatic optimization suggestions available."),
			title=_("No Suggestions"),
			realtime=True,
		)
		return

	data = jingrow.cache.hget(RECORDER_REQUEST_HASH, record.name)
	data["suggested_indexes"] = [{"table": idx[0][0], "column": idx[0][1]} for idx in suggested_indexes]
	jingrow.cache.hset(RECORDER_REQUEST_HASH, record.name, data)
	jingrow.publish_realtime("recorder-analysis-complete", user=jingrow.session.user)
	jingrow.msgprint(_("Query analysis complete. Check suggested indexes."), realtime=True, alert=True)


def _optimize_query(query):
	optimizer = DBOptimizer(query=query)
	tables = optimizer.tables_examined()

	# Note: Two passes are required here because we first need basic data to understand which
	# columns need to be analyzed to get accurate cardinality.
	for table in tables:
		pagetype = get_pagetype_name(table)
		stats = _fetch_table_stats(pagetype, columns=[])
		if not stats:
			return
		db_table = DBTable.from_jingrow_ouput(stats)
		optimizer.update_table_data(db_table)

	potential_indexes = optimizer.potential_indexes()
	tablewise_columns = defaultdict(list)
	for idx in potential_indexes:
		tablewise_columns[idx.table].append(idx.column)

	for table in tables:
		pagetype = get_pagetype_name(table)
		stats = _fetch_table_stats(pagetype, columns=tablewise_columns[table])
		if not stats:
			return
		db_table = DBTable.from_jingrow_ouput(stats)
		optimizer.update_table_data(db_table)

	return optimizer.suggest_index()


def _fetch_table_stats(pagetype: str, columns: list[str]) -> dict | None:
	def sql_bool(val):
		return cstr(val).lower() in ("yes", "1", "true")

	if not jingrow.db.table_exists(pagetype):
		return

	table = get_table_name(pagetype, wrap_in_backticks=True)

	schema = []
	for field in jingrow.db.sql(f"describe {table}", as_dict=True):
		schema.append(
			{
				"column": field["Field"],
				"type": field["Type"],
				"is_nullable": sql_bool(field["Null"]),
				"default": field["Default"],
			}
		)

	def update_cardinality(column, value):
		for col in schema:
			if col["column"] == column:
				col["cardinality"] = value
				break

	indexes = []
	for idx in jingrow.db.sql(f"show index from {table}", as_dict=True):
		indexes.append(
			{
				"unique": not sql_bool(idx["Non_unique"]),
				"cardinality": idx["Cardinality"],
				"name": idx["Key_name"],
				"sequence": idx["Seq_in_index"],
				"nullable": sql_bool(idx["Null"]),
				"column": idx["Column_name"],
				"type": idx["Index_type"],
			}
		)
		if idx["Seq_in_index"] == 1:
			update_cardinality(idx["Column_name"], idx["Cardinality"])

	total_rows = cint(
		jingrow.db.sql(
			f"""select table_rows
			   from  information_schema.tables
			   where table_name = 'tab{pagetype}'"""
		)[0][0]
	)

	# fetch accurate cardinality for columns by query. WARN: This can take A LOT of time.
	for column in columns:
		cardinality = _get_column_cardinality(table, column)
		update_cardinality(column, cardinality)

	return {
		"table_name": table.strip("`"),
		"total_rows": total_rows,
		"schema": schema,
		"indexes": indexes,
	}


@redis_cache
def _get_column_cardinality(table, column):
	return jingrow.db.sql(f"select count(distinct {column}) from {table}")[0][0]


def get_pagetype_name(table_name: str) -> str:
	return table_name.removeprefix("tab")
