import pandas as pd
import json
import json
import numpy as np
from numbers import Number
import json, re
import inspect
import gwaslab as gl
from gwaslab_agent.g_tools import FILTERER_SET

def _resolve_df_value(self, v):
    """Resolve df reference expressions in arguments.

    Supports strings like `df_3.Column[:10]` to pull column values from a
    previously stored DataFrame in `self.DATA_REGISTRY`. If the resolved
    value is a Series, it will be converted to a Python list.

    Also recursively resolves lists and dicts of such expressions.
    """
    try:
        import re
        if isinstance(v, str):
            m = re.fullmatch(r"(df_\d+)\.([A-Za-z0-9_]+)(?:\[:?(\d+)\])?", v)
            if m:
                df_id = m.group(1)
                col = m.group(2)
                n = m.group(3)
                df = self.DATA_REGISTRY.get(df_id)
                s = df[col]
                if isinstance(n, str):
                    s = s.iloc[:int(n)]
                try:
                    import pandas as pd
                    if isinstance(s, pd.Series):
                        return s.tolist()
                except Exception:
                    pass
                try:
                    return list(s)
                except Exception:
                    return s
        if isinstance(v, list):
            return [_resolve_df_value(self, x) for x in v]
        if isinstance(v, dict):
            return {k: _resolve_df_value(self, vv) for k, vv in v.items()}
        return v
    except Exception:
        return v

def wrap_loader_method(self, name, method):
    """
    Wrap a loader method and serialize its output into an LLM-safe JSON string.

    Behavior
    - Executes the given loader `method(**kwargs)` and captures trailing log text.
    - Normalizes the return value into a compact JSON payload with `status`, `method`, `type`, `data`, and `log`.
    - For `gl.Sumstats` results:
      - Synchronizes `self.sumstats.data`, `self.sumstats.meta`, and `self.sumstats.build`.
      - Merges logs `self.log.combine(result.log, pre=False)`.
      - Returns `type="gl.Sumstats"` with a confirmation string in `data`.
      - Note: assignment-style logging (e.g., `sumstats = gl.Sumstats(...)`) is handled upstream in the print layer when `type == "gl.Sumstats"`.

    Type Mapping
    - `gl.Sumstats` → `type="gl.Sumstats"`; `data` is a confirmation string.
    - `pandas.DataFrame` → `type="DataFrame"`; `data` includes Markdown preview, headers, row/col counts, and truncation flags.
    - `pandas.Series` → `type="Series"`; small series as dict; large series include preview dict, size, truncated flag.
    - `numpy.ndarray` → `type="ndarray"`; small arrays as list; large arrays include preview list, shape, dtype, size, truncated flag.
    - `list`/`tuple` → `type="list"/"tuple"`; large sequences include preview and truncation metadata.
    - `dict` → `type="dict"` (original mapping).
    - numeric (`numbers.Number`) → `type="number"` (value in `data`).
    - `str` → `type="string"` (original text in `data`).
    - `None` → `type="none"` with message "Executed successfully (no return value).".
    - Other JSON-serializable objects → `type="unknown_jsonable"` via JSON conversion.
    - Otherwise → `type="string_fallback"` using `str(result)`.

    Error Handling
    - On exception, returns `status="error"`, `method`, `error`, and scrubbed `log` text when available.

    Parameters
    - `name` (str): Original method name for payload bookkeeping.
    - `method` (callable): Loader function to invoke; must accept keyword arguments.

    Returns
    - `str`: JSON payload with normalized fields described above.
    """
    def wrapped(**kwargs):
        try:
            previous_log_end = len(self.log.log_text)
            result = method(**kwargs)
            self.archive.append(result)
            new_log = self.log.log_text[previous_log_end:]
            
            if isinstance(result, gl.Sumstats):
                out_type = "gl.Sumstats"
                data_string = "Sumstats has been successfully loaded."
                self.sumstats.data = result.data
                self.sumstats.meta = result.meta
                self.sumstats.build = result.build
                self.log.combine(result.log,pre=False)
                new_log = self.log.log_text[previous_log_end + 1:]
                return json.dumps({
                    "status": "success",
                    "method": name,
                    "type": out_type,
                    "data": data_string,
                    "log": new_log.strip() if new_log else ""
                }, ensure_ascii=False)

            elif isinstance(result, pd.DataFrame):
                out_type = "DataFrame"
                max_rows = 20
                max_cols = 20
                df = result
                if isinstance(df.index, pd.MultiIndex):
                    df = df.reset_index()
                total_rows = int(len(df))
                total_cols = int(df.shape[1])
                columns = df.columns.to_list()
                preview = df.iloc[:max_rows, :max_cols]
                data = {
                    "preview": preview.to_markdown(index=False),
                    "columns": columns,
                    "rows": total_rows,
                    "cols": total_cols,
                    "truncated_rows": total_rows > max_rows,
                    "truncated_cols": total_cols > max_cols,
                }

            elif isinstance(result, pd.Series):
                out_type = "Series"
                max_items = 1000
                total_items = int(len(result))
                if total_items > max_items:
                    preview = result.iloc[:max_items]
                    data = {
                        "preview": preview.to_dict(),
                        "size": total_items,
                        "truncated": True,
                    }
                else:
                    data = result.to_dict()

            elif isinstance(result, np.ndarray):
                max_size = 10000
                arr = result
                if arr.size > max_size:
                    preview_count = min(arr.size, 1000)
                    data = {
                        "preview": arr.ravel()[:preview_count].tolist(),
                        "shape": list(arr.shape),
                        "dtype": str(arr.dtype),
                        "size": int(arr.size),
                        "truncated": True,
                    }
                else:
                    data = arr.tolist()
                out_type = "ndarray"

            # --- Tuple or list ---
            elif isinstance(result, (list, tuple)):
                out_type = type(result).__name__
                max_items = 1000
                try:
                    total_items = int(len(result))
                except Exception:
                    total_items = None
                if isinstance(total_items, int) and total_items > max_items:
                    data = {
                        "preview": list(result[:max_items]),
                        "size": total_items,
                        "truncated": True,
                    }
                else:
                    data = result

            # --- Dict ---
            elif isinstance(result, dict):
                data = result
                out_type = "dict"

            # --- Single number ---
            elif isinstance(result, Number):
                data = result
                out_type = "number"

            # --- String-like ---
            elif isinstance(result, str):
                data = result
                out_type = "string"

            # --- None or void return ---
            elif result is None:
                data = "Executed successfully (no return value)."
                out_type = "none"

            # --- Fallback: attempt JSON serialization ---
            else:
                try:
                    data = json.loads(json.dumps(result, default=str))
                    out_type = "unknown_jsonable"
                except Exception:
                    data = str(result)
                    out_type = "string_fallback"

            # --- Return consistent JSON ---
            return json.dumps({
                "status": "success",
                "method": name,
                "type": out_type,
                "data": data,
                "log": new_log.strip() if new_log else ""
            }, ensure_ascii=False)

        except Exception as e:
            # In case the method itself fails
            err_log = ""
            if hasattr(self.log, "getvalue"):
                err_log = self.log.getvalue()
            return json.dumps({
                "status": "error",
                "method": name,
                "error": str(e),
                "log": err_log.strip()
            }, ensure_ascii=False)
    return wrapped


from gwaslab_agent.g_image import _scrub_log
from gwaslab_agent.g_image import _show_locally
from gwaslab_agent.g_image import _is_figure
def wrap_main_agent_method(self, name, method):
    """Wrap a method for LLM-safe, structured output serialization."""
    def wrapped(**kwargs):
        try:
            previous_log_end = len(self.log.log_text)
            if (
                isinstance(name, str)
                and (name in FILTERER_SET or name.startswith("filter"))
                and ("inplace" not in kwargs)
            ):
                try:
                    sig = inspect.signature(method)
                    if "inplace" in sig.parameters:
                        kwargs["inplace"] = False
                except Exception:
                    pass
            try:
                kwargs = _resolve_df_value(self, kwargs)
            except Exception:
                pass
            result = method(**kwargs)
            
            self.archive.append(result)
            new_log = self.log.log_text[previous_log_end:]
            new_log = _scrub_log(new_log)

            # Resolve your handle → real object (unchanged)
            if isinstance(result, dict) and "subset_id" in result:
                obj_id = result["subset_id"]
                obj = self.FILTERED_SUMSTATS.get(obj_id)
                result = obj

            # --- If it's a figure/image: SHOW LOCALLY but NEVER return the image ---
            if _is_figure(result):
                _show_locally(result)  # renders in Jupyter
                return json.dumps({
                    "status": "success",
                    "method": name,
                    "type": "image_redacted",
                    "data": "Image/figure creation finished.",
                    "log": new_log
                }, ensure_ascii=False)

            if isinstance(result, gl.Sumstats):
                if isinstance(name, str) and name in FILTERER_SET:
                    obj_id = self.FILTERED_SUMSTATS.put(result)
                    new_log = self.FILTERED_SUMSTATS.get(obj_id).log.log_text[previous_log_end + 1:]
                    return {
                        "status": "success",
                        "type": "filtered Sumstats object",
                        "instructions": (
                            "Access using "
                            "`run_on_filtered` for visualization and processing."
                        ),
                        "subset_id": obj_id,
                        "log": _scrub_log(new_log),
                    }
                else:
                    return json.dumps({
                        "status": "success",
                        "method": name,
                        "type": "none",
                        "data": "Executed successfully (no return value).",
                        "log": new_log
                    }, ensure_ascii=False)

            if isinstance(result, pd.DataFrame):
                out_type = "DataFrame"
                max_rows = 100
                max_cols = 50
                df = result
                if isinstance(df.index, pd.MultiIndex):
                    df = df.reset_index()
                total_rows = int(len(df))
                total_cols = int(df.shape[1])
                preview = df.iloc[:max_rows, :max_cols]
                try:
                    df_id = self.DATA_REGISTRY.put(result)
                except Exception:
                    df_id = None
                data = {
                    "preview": preview.to_markdown(index=False),
                    "rows": total_rows,
                    "cols": total_cols,
                    "truncated_rows": total_rows > max_rows,
                    "truncated_cols": total_cols > max_cols,
                    "df_id": df_id,
                }

            elif isinstance(result, pd.Series):
                out_type = "Series"
                max_items = 1000
                total_items = int(len(result))
                if total_items > max_items:
                    preview = result.iloc[:max_items]
                    data = {
                        "preview": preview.to_dict(),
                        "size": total_items,
                        "truncated": True,
                    }
                else:
                    data = result.to_dict()

            elif isinstance(result, np.ndarray):
                # If ndarray looks like an image, still render locally and redact
                if _is_figure(result):
                    _show_locally(result)
                    return json.dumps({
                        "status": "success",
                        "method": name,
                        "type": "image_redacted",
                        "data": "Image/figure shown in notebook; content withheld from LLM.",
                        "log": new_log
                    }, ensure_ascii=False)
                out_type = "ndarray"
                max_size = 10000
                arr = result
                if arr.size > max_size:
                    preview_count = min(arr.size, 1000)
                    data = {
                        "preview": arr.ravel()[:preview_count].tolist(),
                        "shape": list(arr.shape),
                        "dtype": str(arr.dtype),
                        "size": int(arr.size),
                        "truncated": True,
                    }
                else:
                    data = arr.tolist()

            elif isinstance(result, (list, tuple)):
                out_type = type(result).__name__
                max_items = 1000
                try:
                    total_items = int(len(result))
                except Exception:
                    total_items = None
                if isinstance(total_items, int) and total_items > max_items:
                    data = {
                        "preview": list(result[:max_items]),
                        "size": total_items,
                        "truncated": True,
                    }
                else:
                    data = result

            elif isinstance(result, dict):
                data = result
                out_type = "dict"

            elif isinstance(result, Number):
                data = result
                out_type = "number"

            elif isinstance(result, str):
                data = result
                out_type = "string"

            elif result is None:
                data = "Executed successfully (no return value)."
                out_type = "none"

            else:
                try:
                    data = json.loads(json.dumps(result, default=str))
                    out_type = "unknown_jsonable"
                except Exception:
                    data = str(result)
                    out_type = "string_fallback"

            return json.dumps({
                "status": "success",
                "method": name,
                "type": out_type,
                "data": data,
                "log": _scrub_log(new_log)
            }, ensure_ascii=False)

        except Exception as e:
            err_log = ""
            if hasattr(self.log, "getvalue"):
                err_log = self.log.getvalue()
            return json.dumps({
                "status": "error",
                "method": name,
                "error": str(e),
                "log": _scrub_log(err_log)
            }, ensure_ascii=False)
    return wrapped
