from __future__ import annotations

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

import numpy as np
from pandas import DataFrame
import pandas as pd
import decimal

from relationalai import debugging
from .. import metamodel as m

#--------------------------------------------------
# Constants
#--------------------------------------------------

UNIXEPOCH = 62135683200000
MILLISECONDS_PER_DAY = 24 * 60 * 60 * 1000

#--------------------------------------------------
# Result formatting
#--------------------------------------------------

def format_value(value, type):
    if "UInt128" in type:
        return decode_uint128(value)
    elif "FixedDecimal" in type:
        decimal_info = re.search(r"FixedDecimal\{Int(\d+), (\d+)\}", type)
        if decimal_info:
            bits = int(decimal_info.group(1))
            scale = int(decimal_info.group(2))
            if bits == 128:
                return (decimal.Decimal(str((int(value[1]) << 64) + int(value[0]))) if value[1] > 0 else decimal.Decimal(str(value[0]))) / (10 ** scale)
            else:
                return decimal.Decimal(str(value)) / (10 ** scale)
    elif "Int128" in type:
        return (int(value[1]) << 64) + int(value[0]) if value[1] > 0 else value[0]
    elif "Missing" in type:
        return None
    return value

def format_columns(result_frame:DataFrame, types:List[str]):
    for i, col in enumerate(result_frame.columns):
        if "UInt128" in types[i]:
            result_frame[col] = result_frame[col].apply(decode_uint128)
        elif "FixedDecimal" in types[i]:
            decimal_info = re.search(r"FixedDecimal\{Int(\d+), (\d+)\}", types[i])
            if decimal_info:
                bits = int(decimal_info.group(1))
                scale = int(decimal_info.group(2))
                if bits == 128:
                    result_frame[col] = result_frame[col].apply(lambda x: decode_fixed_decimal_128(x, scale))
                else:
                    result_frame[col] = result_frame[col].apply(lambda x: decimal.Decimal(str(x)) / (10 ** scale))
        elif "Int128" in types[i]:
            result_frame[col] = result_frame[col].apply(decode_int128)
        elif "Missing" in types[i]:
            result_frame[col] = result_frame[col].apply(lambda _: None)
        elif types[i] == "Dates.DateTime":
            result_frame[col] = pd.to_datetime(result_frame[col] - UNIXEPOCH, unit="ms")
        elif types[i] == "Dates.Date":
            result_frame[col] = pd.to_datetime(result_frame[col] * MILLISECONDS_PER_DAY - UNIXEPOCH, unit="ms")
    return result_frame

def decode_int128(x):
    high, low = int(x[1]), int(x[0])
    decoded = (high << 64) + low if high != 0 else low
    # For negative numbers, we need to apply two's complement
    if high & (1 << 63): # MSB is sign bit
        decoded = decoded - (1 << 128)
    return decoded

dec128ctx = decimal.Context(prec=38, rounding=decimal.ROUND_HALF_EVEN)
def decode_fixed_decimal_128(x: Any, scale: int) -> Any:
    high, low = int(x[1]), int(x[0])
    value = (high << 64) + low

    # For negative numbers, we need to apply two's complement
    if high & (1 << 63): # MSB is sign bit
        value = value - (1 << 128)

    return decimal.Decimal(value, context=dec128ctx).scaleb(-scale, context=dec128ctx)

def decode_uint128(x):
    if isinstance(x, dict):
        keys = sorted(x.keys())
        result = []
        for key in keys:
            value = x[key]
            if isinstance(value, np.ndarray):
                result.extend(value)
            elif isinstance(value, (int, np.integer)):
                result.append(value)
            else:
                raise ValueError(f"Unsupported type for key {key}: {type(value)}")

        return decode_uint128(np.array(result, dtype=np.uint64))

    elif isinstance(x, np.ndarray):
        return base64.b64encode(x.tobytes()).decode()[:-2]

    else:
        raise ValueError(f"Unsupported input type: {type(x)}")

def convert_hash_value(value):
    if isinstance(value, np.ndarray):
        return tuple(value)
    return value

def merge_columns(left:DataFrame, right:DataFrame, on:List[str]):
    if not len(left.columns):
        for col in on:
            left.insert(loc=len(left.columns), column=col, value=np.nan)
    if left.empty:
        col_count = len(left.columns)
        left.insert(loc=col_count, column=f'new_column{col_count}', value=np.nan)
    if right.empty:
        col_count = len(left.columns)
        left.insert(loc=col_count, column=f'new_column{col_count}', value=np.nan)
        return left
    if not on:
        return pd.merge(left, right, how='cross')
    return pd.merge(left, right, on=on, how='outer')

def format_results(results, task:m.Task|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 = task.return_cols() if task else []
            has_cols:List[DataFrame] = [DataFrame() for _ in range(0, len(ret_cols))]
            key_len = 0
            for result in results.results:
                relation_id = result["relationId"]
                result_frame = result["table"].to_pandas()

                types = [
                    t
                    for t in result["relationId"].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" in relation_id:
                    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)
                    from relationalai.experimental.inspect import _print_watch_frame
                    _print_watch_frame(result_frame)

                # Process other results
                else:
                    result_frame = format_columns(result_frame, types)
                    result["table"] = result_frame
                    if ":output/:cols" in result["relationId"]:
                        matched = re.search(r":col([0-9]+)", result["relationId"])
                        assert matched, f"Could not determine column mapping for {result['relationaId']}"
                        col_ix = int(matched.group(1))
                        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 result["relationId"]:
                        data_frame = pd.concat(
                            [data_frame, result_frame], ignore_index=True
                        )

            if any(not col.empty for col in has_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)
                data_frame = df_wide_reset.drop(columns=key_cols)

            try:
                data_frame.sort_values(by=[str(c) for c in data_frame.columns], ascending=[True] * len(data_frame.columns), inplace=True)
            except Exception:
                pass
            data_frame = data_frame.reset_index(drop=True)

            if (
                len(ret_cols)
                and len(data_frame.columns) == len(ret_cols)
                and task is not None
            ):
                data_frame.columns = task.return_cols()[0 : len(data_frame.columns)]
        return (data_frame, list(problems.values()))
