from __future__ import annotations

from collections import defaultdict
from functools import reduce
import re
from typing import Any, Dict, List, Tuple, cast

from pandas import DataFrame
import pandas as pd

from relationalai import debugging
from relationalai.clients.result_helpers import format_columns, format_value, merge_columns, convert_hash_value, \
    sort_data_frame_result
from relationalai.tools.constants import Generation

# Convert LQP results into the expected single wide table dataframe for the end user
# - Requires identifying and unrolling all GNF relations and populating any resulting nulls
# - Relies on expected ordering of results, as we do not have IDs or names associated with columns here.
# - At the end, col names are rewritten with the expected output names blindly; we trust the stitching
#   has been done correctly for the names to line up with the correct columns
# `result_cols`: the expected output column names in order
# `key_locations`: where output keys are located in outputs
def format_results(results, result_cols:List[str], key_locations:List[int])  -> Tuple[DataFrame, List[Any]]:
    with debugging.span("format_results"):
        data_frame = DataFrame()
        problems = defaultdict(
            lambda: {
                "message": "",
                "path": "",
                "start_line": None,
                "end_line": None,
                "report": "",
                "code": "",
                "severity": "",
                "decl_id": "",
                "name": "",
                "output": "",
                "end_character": None,
                "start_character": None,
                "props": {},
            }
        )

        # Check if there are any results to process
        if not results.results:
            return (data_frame, list(problems.values()))

        out_keys_n = len(key_locations) # number of keys in output
        assert out_keys_n <= len(result_cols)
        out_vals_n = len(result_cols) - out_keys_n # number of values in output

        # only create cols for values, we handle keys separately as they are not GNF
        has_cols:List[DataFrame] = [DataFrame() for _ in range(0, out_vals_n)]
        keys_data_frame = DataFrame()
        key_len = 0

        for result in results.results:
            relation_id = result["relationId"]
            assert isinstance(relation_id, str)
            result_frame = result["table"].to_pandas()
            types = [
                t
                for t in relation_id.split("/")
                if t != "" and not t.startswith(":")
            ]

            # Process diagnostics
            if "/:rel/:catalog/:diagnostic/" in relation_id:
                # Handle different types of diagnostics based on relation_id
                if "/:message/" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[row.iloc[0]]["message"] = row.iloc[1]
                elif "/:report/" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[row.iloc[0]]["report"] = row.iloc[1]
                elif "/:start/:line" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 2:
                            problems[row.iloc[0]]["start_line"] = row.iloc[2]
                elif "/:end/:line" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 2:
                            problems[row.iloc[0]]["end_line"] = row.iloc[2]
                elif "/:model" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[row.iloc[0]]["path"] = row.iloc[1]
                elif "/:severity" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[row.iloc[0]]["severity"] = row.iloc[1]
                elif "/:code" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[row.iloc[0]]["code"] = row.iloc[1]

            # Process integrity constraint violations
            elif "/:rel/:catalog/:ic_violation" in relation_id:
                if "/:decl_id" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])]["decl_id"] = (
                                row.iloc[1]
                            )
                            problems[convert_hash_value(row.iloc[0])]["code"] = (
                                "IC_VIOLATION"
                            )
                elif "/:model" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])]["path"] = (
                                row.iloc[1]
                            )
                elif "/:name" in relation_id:
                    # Get the last segment of the relation_id as the name
                    segments = [
                        segment[1:]
                        for segment in relation_id.split("/")[4:]
                        if segment.startswith(":")
                    ]
                    for _, row in result_frame.iterrows():
                        problems[convert_hash_value(row.iloc[0])]["name"] = segments[-1]
                elif "/:output" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])]["message"] = (
                                row.iloc[1]
                            )
                elif "/:range/:end/:character" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])][
                                "end_character"
                            ] = row.iloc[1]
                elif "/:range/:end/:line" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])]["end_line"] = (
                                row.iloc[1]
                            )
                elif "/:range/:start/:character" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])][
                                "start_character"
                            ] = row.iloc[1]
                elif "/:range/:start/:line" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])]["start_line"] = (
                                row.iloc[1]
                            )
                elif "/:report" in relation_id:
                    for _, row in result_frame.iterrows():
                        if len(row) > 1:
                            problems[convert_hash_value(row.iloc[0])]["report"] = (
                                row.iloc[1]
                            )

            elif "/:pyrel_error_attrs" in relation_id:
                assert len(result_frame.columns) == 3, \
                    f"Error attribute with {len(result_frame.columns)} cols, expected 3"
                for _, row in result_frame.iterrows():
                    id = convert_hash_value(row.iloc[0])
                    problems[id]["code"] = "PYREL_ERROR"
                    if row.iloc[1] == "message":
                        problems[id]["message"] = row.iloc[2]
                    elif row.iloc[1] == "severity":
                        problems[id]["severity"] = row.iloc[2]
                    else:
                        props = cast(Dict, problems[id]["props"])
                        props[row.iloc[1]] = format_value(row.iloc[2], relation_id.split("/")[-1])

            elif "/:__pyrel_debug_watch" in relation_id:
                result_frame = format_columns(result_frame, types, Generation.QB)
                from relationalai.experimental.inspect import _print_watch_frame
                _print_watch_frame(result_frame)

            # Process outputs
            else:
                result_frame = format_columns(result_frame, types, Generation.QB)
                result["table"] = result_frame
                matched = re.search(r"^/:output_cols_col([0-9]+)/", relation_id)
                if matched:
                    col_ix = int(matched.group(1))

                    # Generate col names and write them into the df (idn for keys, vn for cols)
                    key_cols = [f"id{i}" for i in range(0, len(result_frame.columns) - 1)]
                    key_len = len(key_cols)
                    # Don't use the actual name as users are free to use our generated names
                    result_col_name = f"v{col_ix}"
                    result_frame.columns = [*key_cols, result_col_name]

                    has_cols[col_ix] = result_frame
                elif relation_id.startswith("/:output_keys"):
                    # data for all keys (wide), to merge in later
                    keys_data_frame = result_frame

                    # Rename wide key col names to match key cols in df_wide_reset
                    keys_data_frame.columns = pd.RangeIndex(len(keys_data_frame.columns))
                    keys_data_frame = keys_data_frame.rename(columns=lambda c: f"id{c}")

                elif relation_id.startswith("/:output"): # wide outputs case
                    data_frame = pd.concat(
                        [data_frame, result_frame], ignore_index=True
                    )

        # GNF values case: stitch together output vals and keys into one wide dataframe
        if any(not col.empty for col in has_cols):
            # Merge value cols together by their key cols
            key_cols = [f"id{i}" for i in range(0, key_len)]
            df_wide_reset = reduce(lambda left, right: merge_columns(left, right, key_cols), has_cols)

            # Join wide keys with wide vals (keys all at the front; still needs reordering)
            data_frame = pd.merge(keys_data_frame, df_wide_reset, on=key_cols, how='outer')

            # Reorder outputs
            if key_locations:
                data_frame = _shift_keys(data_frame, keys_data_frame, out_keys_n, key_locations)

            else: # if no keys in output, just drop all of the key cols
                data_frame = data_frame.drop(columns=key_cols)

        # Empty values case: reorder/drop keys as needed
        elif not keys_data_frame.empty:
            if key_locations: # Reorder outputs
                data_frame = _shift_keys(keys_data_frame, keys_data_frame, out_keys_n, key_locations)
            else: # if there are no keys to output, we may still need to populate nulls for output values
                # Take into account the cols that could contain values (even though they're empty)
                key_cols = [f"id{i}" for i in range(0, len(keys_data_frame.columns))]
                has_cols.append(keys_data_frame) # include the keys so we know how many nulls to generate
                df_wide_reset = reduce(lambda left, right: merge_columns(left, right, key_cols), has_cols)
                data_frame = df_wide_reset.drop(columns=key_cols)

        data_frame = sort_data_frame_result(data_frame)

        # Overwrite column names with user-defined names
        # The assumption is that the extra keys have been chopped off the front, and the
        # remaining columns are in the correct order and require renaming
        if 0 < len(data_frame.columns) <= len(result_cols):
            data_frame.columns = result_cols[:len(data_frame.columns)]

        return (data_frame, list(problems.values()))

# Reorder `res` df to match user-specified output order and drop non-output key cols
# E.g., Current df looks like:
#   [out_key_pos1, out_key_pos4, hidden_key_1, value_pos2, value_pos3, value_pos5]
# Target df looks like:
#   [out_key_pos1, value_pos2, value_pos3, out_key_pos4, value_pos5]
def _shift_keys(res:DataFrame, keys_data_frame:DataFrame, out_keys_n:int, key_locations:List[int]):
    offset = len(keys_data_frame.columns) # index of first value in data frame
    assert out_keys_n <= offset
    extra_keys_n = offset - out_keys_n # number of keys to drop (those not in output)

    # Shift output keys into the correct spots and drop the rest
    for i in key_locations:
        res = res[res.columns[1:].insert(i + offset - 1, res.columns[0])]
        offset -= 1
    res = res.drop(columns=res.columns[:extra_keys_n])
    return res
