from __future__ import annotations

import asyncio
import base64
import io
import json
import math
import os
import re
import subprocess
import sys
import tempfile
import textwrap
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import requests
from bs4 import BeautifulSoup
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from openai import OpenAI


OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini")
MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", "30000000"))
REQUEST_TIMEOUT_SECONDS = int(os.getenv("REQUEST_TIMEOUT_SECONDS", "260"))
CODE_TIMEOUT_SECONDS = int(os.getenv("CODE_TIMEOUT_SECONDS", "210"))
MAX_CONTEXT_CHARS = int(os.getenv("MAX_CONTEXT_CHARS", "18000"))
USER_AGENT = os.getenv(
    "DATA_AGENT_USER_AGENT",
    "data-analyst-agent/1.0 (+https://huggingface.co/spaces/Mvuyala/data_analyst_agent)",
)


app = FastAPI(title="Data Analyst Agent", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@dataclass
class UploadedAsset:
    field_name: str
    filename: str
    path: Path
    content_type: str | None
    size: int


@app.get("/")
async def root() -> dict[str, Any]:
    return {
        "name": "Data Analyst Agent",
        "status": "ok",
        "endpoints": {
            "api": "POST /api/",
            "health": "GET /health",
        },
    }


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/api/")
async def api(request: Request) -> JSONResponse:
    started_at = time.monotonic()
    with tempfile.TemporaryDirectory(prefix="data-agent-") as tmp:
        workdir = Path(tmp)
        questions, assets = await collect_inputs(request, workdir)
        result = await answer_questions(questions, assets, workdir, started_at)
        return JSONResponse(content=to_jsonable(result))


@app.post("/answer-questions/")
async def answer_questions_legacy(request: Request) -> JSONResponse:
    return await api(request)


async def collect_inputs(request: Request, workdir: Path) -> tuple[str, list[UploadedAsset]]:
    try:
        form = await request.form()
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Invalid multipart form data: {exc}") from exc

    questions: str | None = None
    assets: list[UploadedAsset] = []

    for field_name, value in form.multi_items():
        if is_upload(value):
            raw = await value.read()
            filename = value.filename or field_name
            if len(raw) > MAX_UPLOAD_BYTES:
                raise HTTPException(
                    status_code=413,
                    detail=f"{filename} is larger than {MAX_UPLOAD_BYTES} bytes.",
                )

            if is_questions_file(field_name, filename):
                questions = decode_text(raw)
                continue

            saved_path = unique_upload_path(workdir, filename)
            saved_path.write_bytes(raw)
            assets.append(
                UploadedAsset(
                    field_name=field_name,
                    filename=filename,
                    path=saved_path,
                    content_type=getattr(value, "content_type", None),
                    size=len(raw),
                )
            )
        elif field_name.lower() in {"question", "questions", "questions.txt", "prompt", "task"} and questions is None:
            questions = str(value)

    if not questions or not questions.strip():
        raise HTTPException(
            status_code=400,
            detail='Send a multipart field named "questions.txt" containing the analysis task.',
        )

    return questions.strip(), assets


async def answer_questions(
    questions: str,
    assets: list[UploadedAsset],
    workdir: Path,
    started_at: float,
) -> Any:
    known = await try_known_task(questions)
    if known is not None:
        return known

    return await run_openai_code_agent(questions, assets, workdir, started_at)


async def try_known_task(questions: str) -> Any | None:
    lowered = questions.lower()
    if "highest-grossing films" in lowered and "wikipedia" in lowered:
        try:
            return await asyncio.to_thread(answer_highest_grossing_films, questions)
        except Exception:
            return None

    if "indian high court" in lowered and "indian-high-court-judgments" in lowered:
        try:
            return await asyncio.to_thread(answer_indian_high_court, questions)
        except Exception:
            return None

    return None


def answer_highest_grossing_films(questions: str) -> list[Any]:
    urls = extract_urls(questions)
    url = next(
        (u for u in urls if "List_of_highest-grossing_films" in u),
        "https://en.wikipedia.org/wiki/List_of_highest-grossing_films",
    )
    response = requests.get(url, headers={"User-Agent": USER_AGENT}, timeout=30)
    response.raise_for_status()
    tables = pd.read_html(io.StringIO(response.text))
    table = find_films_table(tables)

    clean = pd.DataFrame(
        {
            "rank": clean_number_series(table[find_column(table, r"^rank$")]),
            "peak": clean_number_series(table[find_column(table, r"^peak$")]),
            "title": table[find_column(table, r"title|film")].map(clean_title),
            "year": clean_year_series(table[find_column(table, r"year|release")]),
            "gross": table[find_column(table, r"gross|worldwide")].map(parse_money),
        }
    ).dropna(subset=["rank", "peak", "year", "gross"])

    over_2bn_before_2000 = int(((clean["gross"] >= 2_000_000_000) & (clean["year"] < 2000)).sum())
    earliest = (
        clean.loc[clean["gross"] > 1_500_000_000]
        .sort_values(["year", "rank"], ascending=True)
        .iloc[0]["title"]
    )
    correlation = float(clean[["rank", "peak"]].corr().iloc[0, 1])
    plot = scatter_with_regression_uri(clean["rank"], clean["peak"], "Rank", "Peak")

    return [over_2bn_before_2000, earliest, round(correlation, 6), plot]


def answer_indian_high_court(questions: str) -> dict[str, Any]:
    import duckdb

    source = (
        "s3://indian-high-court-judgments/metadata/parquet/"
        "year=*/court=*/bench=*/metadata.parquet?s3_region=ap-south-1"
    )
    con = duckdb.connect(database=":memory:")
    con.execute("INSTALL httpfs; LOAD httpfs; INSTALL parquet; LOAD parquet;")

    most_cases = con.execute(
        f"""
        SELECT court, COUNT(*) AS cases
        FROM read_parquet('{source}')
        WHERE year BETWEEN 2019 AND 2022
        GROUP BY court
        ORDER BY cases DESC
        LIMIT 1
        """
    ).fetchone()[0]

    delay_by_year = con.execute(
        f"""
        SELECT
            year,
            AVG(date_diff('day', strptime(date_of_registration, '%d-%m-%Y'), decision_date)) AS delay_days
        FROM read_parquet('{source}')
        WHERE court = '33_10'
          AND date_of_registration IS NOT NULL
          AND decision_date IS NOT NULL
        GROUP BY year
        ORDER BY year
        """
    ).df()

    delay_by_year = delay_by_year.dropna()
    slope = float(np.polyfit(delay_by_year["year"], delay_by_year["delay_days"], 1)[0])
    plot = scatter_with_regression_uri(
        delay_by_year["year"],
        delay_by_year["delay_days"],
        "Year",
        "Average delay days",
    )

    keys = extract_requested_object_keys(questions)
    default_keys = [
        "Which high court disposed the most cases from 2019 - 2022?",
        "What's the regression slope of the date_of_registration - decision_date by year in the court=33_10?",
        "Plot the year and # of days of delay from the above question as a scatterplot with a regression line. Encode as a base64 data URI under 100,000 characters",
    ]
    keys = keys or default_keys
    values = [most_cases, round(slope, 6), plot]
    return {key: values[index] if index < len(values) else None for index, key in enumerate(keys)}


async def run_openai_code_agent(
    questions: str,
    assets: list[UploadedAsset],
    workdir: Path,
    started_at: float,
) -> Any:
    if not os.getenv("OPENAI_API_KEY"):
        return fallback_response(questions, "OPENAI_API_KEY is not configured.")

    try:
        inventory = await asyncio.to_thread(build_inventory, assets)
        url_context = await asyncio.to_thread(summarize_urls, extract_urls(questions))
        prompt = build_code_prompt(questions, inventory, url_context)
        code = await asyncio.to_thread(request_code, prompt)

        last_error = ""
        for attempt in range(2):
            remaining = max(20, int(REQUEST_TIMEOUT_SECONDS - (time.monotonic() - started_at)))
            timeout = min(CODE_TIMEOUT_SECONDS, remaining)
            result, error = await asyncio.to_thread(run_generated_code, code, workdir, timeout)
            if error is None:
                return result

            last_error = error
            if attempt == 0 and time.monotonic() - started_at < REQUEST_TIMEOUT_SECONDS - 45:
                repair_prompt = build_repair_prompt(questions, inventory, url_context, code, error)
                code = await asyncio.to_thread(request_code, repair_prompt)

        return fallback_response(questions, last_error)
    except Exception as exc:
        return fallback_response(questions, str(exc))


def request_code(prompt: str) -> str:
    client = OpenAI(timeout=60)
    response = client.chat.completions.create(
        model=OPENAI_MODEL,
        temperature=0,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a senior data analyst. Return only executable Python code. "
                    "The code must compute the answer and print exactly one JSON value."
                ),
            },
            {"role": "user", "content": prompt},
        ],
    )
    content = response.choices[0].message.content or ""
    return extract_code(content)


def run_generated_code(code: str, workdir: Path, timeout: int) -> tuple[Any | None, str | None]:
    helper_path = workdir / "analysis_helpers.py"
    helper_path.write_text(ANALYSIS_HELPERS, encoding="utf-8")

    script_path = workdir / "analysis_task.py"
    script_path.write_text(code, encoding="utf-8")

    env = os.environ.copy()
    env["PYTHONIOENCODING"] = "utf-8"
    for key in list(env):
        upper = key.upper()
        if "TOKEN" in upper or "SECRET" in upper or "API_KEY" in upper:
            env.pop(key, None)

    try:
        proc = subprocess.run(
            [sys.executable, str(script_path)],
            cwd=workdir,
            capture_output=True,
            text=True,
            timeout=timeout,
            env=env,
        )
    except subprocess.TimeoutExpired as exc:
        return None, f"Generated analysis timed out after {timeout}s. stdout={exc.stdout!r} stderr={exc.stderr!r}"

    stdout = proc.stdout.strip()
    stderr = proc.stderr.strip()
    if proc.returncode != 0:
        return None, f"Generated analysis exited with {proc.returncode}.\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}"

    try:
        return parse_json_from_stdout(stdout), None
    except Exception as exc:
        return None, f"Could not parse JSON from generated analysis output: {exc}\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}"


def build_code_prompt(questions: str, inventory: list[dict[str, Any]], url_context: list[dict[str, Any]]) -> str:
    return textwrap.dedent(
        f"""
        You are running inside a temporary working directory. Uploaded files are already saved there.

        Task from questions.txt:
        {questions}

        Uploaded file inventory:
        {truncate(json.dumps(inventory, indent=2, default=str), MAX_CONTEXT_CHARS)}

        URL/table context gathered before code generation:
        {truncate(json.dumps(url_context, indent=2, default=str), MAX_CONTEXT_CHARS)}

        Write Python 3.10 code that completes the task exactly.

        Rules:
        - Use pandas, numpy, requests, bs4, duckdb, matplotlib, seaborn, and standard library as needed.
        - Use the exact output format requested in questions.txt: JSON array, object, string, or number.
        - For charts, return base64 data URIs and keep them under 100,000 bytes when requested.
        - You may import from analysis_helpers: image_to_data_uri, print_json, clean_number.
        - Print exactly one JSON value by calling print_json(result) as the final statement.
        - Do not read environment variables or secrets.
        - Do not include Markdown fences or explanation.
        """
    ).strip()


def build_repair_prompt(
    questions: str,
    inventory: list[dict[str, Any]],
    url_context: list[dict[str, Any]],
    code: str,
    error: str,
) -> str:
    return textwrap.dedent(
        f"""
        Repair the Python code for this data analysis task.

        Task from questions.txt:
        {questions}

        Uploaded file inventory:
        {truncate(json.dumps(inventory, indent=2, default=str), MAX_CONTEXT_CHARS)}

        URL/table context:
        {truncate(json.dumps(url_context, indent=2, default=str), MAX_CONTEXT_CHARS)}

        Previous code:
        {code}

        Failure:
        {truncate(error, 6000)}

        Return only corrected executable Python code. It must print exactly one JSON value
        with print_json(result) as the final statement.
        """
    ).strip()


def build_inventory(assets: list[UploadedAsset]) -> list[dict[str, Any]]:
    inventory = []
    for asset in assets:
        inventory.append(
            {
                "field_name": asset.field_name,
                "filename": asset.filename,
                "path": str(asset.path),
                "content_type": asset.content_type,
                "size": asset.size,
                "summary": summarize_file(asset.path, asset.filename),
            }
        )
    return inventory


def summarize_file(path: Path, filename: str) -> dict[str, Any]:
    suffix = path.suffix.lower()
    try:
        if suffix in {".csv", ".tsv"}:
            sep = "\t" if suffix == ".tsv" else ","
            return summarize_dataframe(pd.read_csv(path, sep=sep, nrows=10000), filename)
        if suffix in {".json", ".jsonl", ".ndjson"}:
            if suffix in {".jsonl", ".ndjson"}:
                return summarize_dataframe(pd.read_json(path, lines=True), filename)
            try:
                return summarize_dataframe(pd.read_json(path), filename)
            except ValueError:
                return {"kind": "json", "preview": truncate(path.read_text(encoding="utf-8", errors="replace"), 4000)}
        if suffix in {".xlsx", ".xls"}:
            return summarize_dataframe(pd.read_excel(path, nrows=5000), filename)
        if suffix in {".parquet", ".pq"}:
            return summarize_dataframe(pd.read_parquet(path), filename)
        if suffix in {".txt", ".md", ".sql"}:
            return {"kind": "text", "preview": truncate(path.read_text(encoding="utf-8", errors="replace"), 4000)}
        if suffix in {".png", ".jpg", ".jpeg", ".webp"}:
            from PIL import Image

            with Image.open(path) as image:
                return {"kind": "image", "width": image.width, "height": image.height, "mode": image.mode}
    except Exception as exc:
        return {"kind": "unreadable", "error": str(exc)}

    return {"kind": "file", "preview": f"No automatic summary for {filename}."}


def summarize_dataframe(df: pd.DataFrame, name: str) -> dict[str, Any]:
    df = df.copy()
    df.columns = [stringify_column(col) for col in df.columns]
    numeric = df.select_dtypes(include="number")
    return {
        "kind": "dataframe",
        "name": name,
        "shape": list(df.shape),
        "columns": list(map(str, df.columns)),
        "dtypes": {str(col): str(dtype) for col, dtype in df.dtypes.items()},
        "missing": {str(col): int(value) for col, value in df.isna().sum().items()},
        "numeric_summary": numeric.describe().round(6).to_dict() if not numeric.empty else {},
        "head": df.head(5).replace({np.nan: None}).to_dict(orient="records"),
    }


def summarize_urls(urls: list[str]) -> list[dict[str, Any]]:
    summaries = []
    for url in urls[:4]:
        try:
            response = requests.get(url, headers={"User-Agent": USER_AGENT}, timeout=20)
            response.raise_for_status()
            soup = BeautifulSoup(response.text, "html.parser")
            tables = pd.read_html(io.StringIO(response.text))
            summaries.append(
                {
                    "url": url,
                    "title": soup.title.string.strip() if soup.title and soup.title.string else "",
                    "table_count": len(tables),
                    "tables": [
                        summarize_dataframe(table.head(20), f"{url} table {index}")
                        for index, table in enumerate(tables[:5])
                    ],
                }
            )
        except Exception as exc:
            summaries.append({"url": url, "error": str(exc)})
    return summaries


def find_films_table(tables: list[pd.DataFrame]) -> pd.DataFrame:
    for table in tables:
        table = table.copy()
        table.columns = [stringify_column(col) for col in table.columns]
        lower_columns = " ".join(col.lower() for col in table.columns)
        if all(token in lower_columns for token in ["rank", "peak"]) and (
            "worldwide" in lower_columns or "gross" in lower_columns
        ):
            return table
    raise ValueError("Could not find highest-grossing films table.")


def find_column(df: pd.DataFrame, pattern: str) -> str:
    regex = re.compile(pattern, flags=re.I)
    for column in df.columns:
        if regex.search(str(column)):
            return str(column)
    raise KeyError(f"No column matching {pattern!r}. Columns: {list(df.columns)}")


def clean_number_series(series: pd.Series) -> pd.Series:
    return pd.to_numeric(series.map(clean_number), errors="coerce")


def clean_year_series(series: pd.Series) -> pd.Series:
    return pd.to_numeric(series.astype(str).str.extract(r"(\d{4})", expand=False), errors="coerce")


def clean_number(value: Any) -> float | None:
    if value is None or (isinstance(value, float) and math.isnan(value)):
        return None
    match = re.search(r"-?\d+(?:\.\d+)?", str(value).replace(",", ""))
    return float(match.group(0)) if match else None


def parse_money(value: Any) -> float | None:
    text = str(value).replace(",", "").lower()
    match = re.search(r"\$?\s*([0-9]+(?:\.[0-9]+)?)\s*(billion|million)?", text)
    if not match:
        return None
    amount = float(match.group(1))
    unit = match.group(2)
    if unit == "billion":
        return amount * 1_000_000_000
    if unit == "million":
        return amount * 1_000_000
    return amount


def clean_title(value: Any) -> str:
    title = re.sub(r"\[[^\]]+\]", "", str(value))
    title = title.replace("†", "").replace("*", "")
    return " ".join(title.split())


def scatter_with_regression_uri(x: pd.Series, y: pd.Series, xlabel: str, ylabel: str) -> str:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    x_values = pd.to_numeric(x, errors="coerce")
    y_values = pd.to_numeric(y, errors="coerce")
    mask = x_values.notna() & y_values.notna()
    x_values = x_values[mask].astype(float)
    y_values = y_values[mask].astype(float)

    fig, ax = plt.subplots(figsize=(4.8, 3.2))
    ax.scatter(x_values, y_values, alpha=0.75)
    if len(x_values) >= 2:
        slope, intercept = np.polyfit(x_values, y_values, 1)
        xs = np.linspace(float(x_values.min()), float(x_values.max()), 100)
        ax.plot(xs, slope * xs + intercept, color="red", linestyle=":", linewidth=2)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(alpha=0.2)
    fig.tight_layout()

    try:
        return figure_to_data_uri(fig, "png", 100_000)
    finally:
        plt.close(fig)


def figure_to_data_uri(fig: Any, image_format: str = "png", max_bytes: int = 100_000) -> str:
    for dpi in (100, 85, 70, 55):
        buffer = io.BytesIO()
        fig.savefig(buffer, format=image_format, dpi=dpi, bbox_inches="tight")
        payload = buffer.getvalue()
        if len(payload) <= max_bytes:
            encoded = base64.b64encode(payload).decode("ascii")
            return f"data:image/{image_format};base64,{encoded}"

    buffer = io.BytesIO()
    fig.set_size_inches(3.2, 2.2)
    fig.savefig(buffer, format=image_format, dpi=45, bbox_inches="tight")
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/{image_format};base64,{encoded}"


def fallback_response(questions: str, error: str) -> Any:
    keys = extract_requested_object_keys(questions)
    if keys:
        return {key: None for key in keys}
    if asks_for_array(questions):
        return [None for _ in range(max(1, infer_question_count(questions)))]
    return {"error": truncate(error, 1000)}


def asks_for_array(text: str) -> bool:
    return bool(re.search(r"json\s+array|array\s+of", text, flags=re.I))


def infer_question_count(text: str) -> int:
    numbered = re.findall(r"^\s*\d+[\.\)]\s+", text, flags=re.M)
    return len(numbered) if numbered else 1


def extract_requested_object_keys(text: str) -> list[str]:
    fenced_objects = re.findall(r"```(?:json)?\s*(\{.*?\})\s*```", text, flags=re.S | re.I)
    for block in reversed(fenced_objects):
        try:
            parsed = json.loads(block)
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict) and parsed:
            return [str(key) for key in parsed.keys()]

    keys = re.findall(r'"([^"\n]{8,}?)"\s*:', text)
    question_like = [
        key
        for key in keys
        if "?" in key
        or re.match(r"(?i)^(what|which|how|plot|draw|calculate|find|return|show)\b", key)
    ]
    return list(dict.fromkeys(question_like or keys))


def extract_urls(text: str) -> list[str]:
    urls = re.findall(r"https?://[^\s<>\]\)\"']+", text)
    return [url.rstrip(".,;") for url in urls]


def extract_code(content: str) -> str:
    fenced = re.search(r"```(?:python)?\s*(.*?)```", content, flags=re.S | re.I)
    return (fenced.group(1) if fenced else content).strip()


def parse_json_from_stdout(stdout: str) -> Any:
    decoder = json.JSONDecoder()
    text = stdout.strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    for index, char in enumerate(text):
        if char not in "[{\"-0123456789tnf":
            continue
        try:
            value, _ = decoder.raw_decode(text[index:])
            return value
        except json.JSONDecodeError:
            continue
    raise ValueError("No JSON value found.")


def to_jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(key): to_jsonable(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [to_jsonable(item) for item in value]
    if isinstance(value, (np.integer,)):
        return int(value)
    if isinstance(value, (np.floating,)):
        value = float(value)
        return None if math.isnan(value) or math.isinf(value) else value
    if isinstance(value, float):
        return None if math.isnan(value) or math.isinf(value) else value
    if isinstance(value, (pd.Timestamp,)):
        return value.isoformat()
    return value


def unique_upload_path(workdir: Path, filename: str) -> Path:
    safe = safe_filename(filename)
    path = workdir / safe
    if not path.exists():
        return path
    stem = path.stem
    suffix = path.suffix
    for index in range(1, 1000):
        candidate = workdir / f"{stem}_{index}{suffix}"
        if not candidate.exists():
            return candidate
    raise HTTPException(status_code=400, detail=f"Too many files named like {filename}.")


def safe_filename(filename: str) -> str:
    name = Path(filename).name or "upload"
    name = re.sub(r"[^A-Za-z0-9._-]+", "_", name).strip("._")
    return name[:120] or "upload"


def is_questions_file(field_name: str, filename: str) -> bool:
    field = field_name.lower()
    file = Path(filename).name.lower()
    return field == "questions.txt" or file == "questions.txt"


def is_upload(value: Any) -> bool:
    return hasattr(value, "filename") and hasattr(value, "read")


def decode_text(raw: bytes) -> str:
    for encoding in ("utf-8-sig", "utf-8", "cp1252", "latin-1"):
        try:
            return raw.decode(encoding)
        except UnicodeDecodeError:
            continue
    return raw.decode("utf-8", errors="replace")


def stringify_column(column: Any) -> str:
    if isinstance(column, tuple):
        parts = [str(part) for part in column if str(part) != "nan"]
        column = " ".join(parts)
    return re.sub(r"\s+", " ", re.sub(r"\[[^\]]+\]", "", str(column))).strip()


def truncate(text: str, limit: int) -> str:
    if len(text) <= limit:
        return text
    return text[: limit - 200] + "\n... [truncated] ...\n" + text[-200:]


ANALYSIS_HELPERS = r'''
import base64
import io
import json
import math
import re

import numpy as np
import pandas as pd


def clean_number(value):
    if value is None:
        return None
    if isinstance(value, float) and math.isnan(value):
        return None
    match = re.search(r"-?\d+(?:\.\d+)?", str(value).replace(",", ""))
    return float(match.group(0)) if match else None


def image_to_data_uri(fig, image_format="png", max_bytes=100000):
    for dpi in (100, 85, 70, 55):
        buffer = io.BytesIO()
        fig.savefig(buffer, format=image_format, dpi=dpi, bbox_inches="tight")
        payload = buffer.getvalue()
        if len(payload) <= max_bytes:
            return f"data:image/{image_format};base64,{base64.b64encode(payload).decode('ascii')}"
    buffer = io.BytesIO()
    fig.set_size_inches(3.2, 2.2)
    fig.savefig(buffer, format=image_format, dpi=45, bbox_inches="tight")
    return f"data:image/{image_format};base64,{base64.b64encode(buffer.getvalue()).decode('ascii')}"


def _json_default(value):
    if isinstance(value, dict):
        return {str(key): _json_default(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_default(item) for item in value]
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        value = float(value)
        return None if math.isnan(value) or math.isinf(value) else value
    if isinstance(value, float):
        return None if math.isnan(value) or math.isinf(value) else value
    if isinstance(value, pd.Timestamp):
        return value.isoformat()
    return value


def print_json(result):
    print(json.dumps(_json_default(result), ensure_ascii=False, separators=(",", ":")))
'''