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 v0.relationalai import debugging
from v0.relationalai.clients.result_helpers import format_columns, format_value, merge_columns, convert_hash_value, \
    sort_data_frame_result
from v0.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
def format_results(results, result_cols:List[str]|None = None, key_locations:List[int]|None = None)  -> 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 len(results.results):
            ret_cols = result_cols or [] # output column names
            key_locations = key_locations or [] # where output keys are located in outputs
            out_keys_n = len(key_locations) # number of keys in output
            assert out_keys_n <= len(ret_cols)
            out_vals_n = len(ret_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"]
                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 v0.relationalai.experimental.inspect import _print_watch_frame
                    _print_watch_frame(result_frame)

                # Process other results
                else:
                    result_frame = format_columns(result_frame, types, Generation.QB)
                    result["table"] = result_frame
                    if "/:output" in relation_id \
                            and "_cols_col" in relation_id:
                        # Match rows with an id like "/:output.*_cols_col[0-9]+"
                        # These should be all of the GNF value outputs
                        matched = re.search(r"_cols_col([0-9]+)", relation_id)
                        assert matched, f"Column id not found for: {relation_id}"
                        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)
                        result_frame.columns = [*key_cols, f"v{col_ix}"]

                        if has_cols[col_ix].empty:
                            has_cols[col_ix] = result_frame
                        else:
                            has_cols[col_ix] = pd.concat([has_cols[col_ix], result_frame], ignore_index=True)
                    elif ":output" in relation_id \
                        and "_keys" in relation_id:
                        # 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 ":output" in relation_id: # 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 len(ret_cols) and len(data_frame.columns) <= len(ret_cols):
                if result_cols is not None:
                    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
