#!/usr/bin/env python3
"""
DataMerase MCP Server

Exposes the DataMerase hybrid synthetic data generator
(https://datamerase.io) as MCP tools: upload a dataset, configure column
roles (target / continuous / categorical / text / ignore), kick off a
generation job, poll it, and download the resulting synthetic CSV.

Configuration (environment variables):
    DATAMERASE_API_URL   Base URL of the DataMerase API.
                          Defaults to https://datamerase.io (production).
    DATAMERASE_API_KEY   Required. Your DataMerase API key (X-API-Key header).
                          No default — the server refuses to start without it.
"""

import csv
import io
import os
import sys
import time
from pathlib import Path
from typing import Any

import requests
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver.exceptions import ToolError

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

API_BASE_URL = os.environ.get("DATAMERASE_API_URL", "https://datamerase.io").rstrip("/")

API_KEY = os.environ.get("DATAMERASE_API_KEY")
if not API_KEY:
    raise RuntimeError(
        "DATAMERASE_API_KEY environment variable is not set. "
        "The DataMerase MCP server requires a valid API key — set "
        "DATAMERASE_API_KEY in your environment (or in the 'env' block of "
        "claude_desktop_config.json) before starting this server. There is "
        "no built-in default key."
    )

REQUEST_TIMEOUT = 30    # seconds — /health, /docs, /generate, /status
UPLOAD_TIMEOUT = 60     # seconds — /upload
DOWNLOAD_TIMEOUT = 60   # seconds — /download

VALID_UPLOAD_EXTENSIONS = {
    ".csv": "text/csv",
    ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ".xls": "application/vnd.ms-excel",
}

mcp = MCPServer("DataMerase")


# ---------------------------------------------------------------------------
# HTTP helpers — every DataMerase JSON endpoint returns the same envelope:
#   {"success": bool, "data": <payload or null>, "error": {"message","code"} or null}
# These helpers centralize the three failure modes the caller asked us to
# handle explicitly: an unreachable API, a non-JSON response, and a clean
# API-level failure (401 / 409 / 429 / column_not_found / etc.) — in every
# case the underlying message is surfaced to the caller, never swallowed.
# ---------------------------------------------------------------------------


def _headers() -> dict:
    return {"X-API-Key": API_KEY}


def _handle_json_response(resp: requests.Response, context: str) -> Any:
    try:
        body = resp.json()
    except ValueError:
        raise ToolError(
            f"{context}: DataMerase API returned a non-JSON response "
            f"(HTTP {resp.status_code}): {resp.text[:500]!r}"
        )

    if not isinstance(body, dict) or "success" not in body:
        raise ToolError(
            f"{context}: DataMerase API returned an unexpected JSON shape "
            f"(HTTP {resp.status_code}): {body!r}"
        )

    if not body.get("success"):
        err = body.get("error") or {}
        message = err.get("message", "Unknown error")
        code = err.get("code", resp.status_code)
        raise ToolError(f"{context} failed: [{code}] {message} (HTTP {resp.status_code})")

    return body.get("data")


def _request(method: str, path: str, context: str, **kwargs) -> Any:
    url = f"{API_BASE_URL}{path}"
    timeout = kwargs.pop("timeout", REQUEST_TIMEOUT)
    try:
        resp = requests.request(method, url, headers=_headers(), timeout=timeout, **kwargs)
    except requests.exceptions.ConnectTimeout:
        raise ToolError(f"{context}: connection to {API_BASE_URL} timed out.")
    except requests.exceptions.ConnectionError as e:
        raise ToolError(f"{context}: could not reach DataMerase API at {API_BASE_URL} ({e}).")
    except requests.exceptions.Timeout:
        raise ToolError(f"{context}: request to {url} timed out after {timeout}s.")
    except requests.exceptions.RequestException as e:
        raise ToolError(f"{context}: request error — {e}")

    return _handle_json_response(resp, context)


# ---------------------------------------------------------------------------
# Internal implementations (plain functions, reused by generate_and_wait so
# the full pipeline doesn't have to go through the MCP tool-call machinery)
# ---------------------------------------------------------------------------


def _upload(file_path: str) -> dict:
    path = Path(file_path).expanduser().resolve()
    if not path.is_file():
        raise ToolError(f"upload_dataset: no such file: {path}")

    ext = path.suffix.lower()
    if ext not in VALID_UPLOAD_EXTENSIONS:
        raise ToolError(
            f"upload_dataset: unsupported file type {ext!r} for {path.name} — "
            f"expected one of {sorted(VALID_UPLOAD_EXTENSIONS)}"
        )

    if path.stat().st_size == 0:
        raise ToolError(f"upload_dataset: {path} is empty (0 bytes).")

    with path.open("rb") as f:
        try:
            resp = requests.post(
                f"{API_BASE_URL}/upload",
                headers=_headers(),
                files={"file": (path.name, f, VALID_UPLOAD_EXTENSIONS[ext])},
                timeout=UPLOAD_TIMEOUT,
            )
        except requests.exceptions.ConnectionError as e:
            raise ToolError(f"upload_dataset: could not reach DataMerase API at {API_BASE_URL} ({e}).")
        except requests.exceptions.Timeout:
            raise ToolError(f"upload_dataset: upload to {API_BASE_URL} timed out after {UPLOAD_TIMEOUT}s.")
        except requests.exceptions.RequestException as e:
            raise ToolError(f"upload_dataset: request error — {e}")

    return _handle_json_response(resp, "upload_dataset")


def _generate(
    job_id: str,
    target_col: str | None,
    continuous_cols: list[str] | None,
    categorical_cols: list[str] | None,
    text_cols: list[str] | None,
    ignore_cols: list[str] | None,
    n_samples: int,
    draws: int,
    tune: int,
    family: str,
    text_epochs: int,
) -> dict:
    continuous_cols = list(continuous_cols or [])
    categorical_cols = list(categorical_cols or [])
    text_cols = list(text_cols or [])
    ignore_cols = list(ignore_cols or [])

    # DataMerase's own generation API has no "ignore" field — a column is
    # excluded from the synthetic output simply by not appearing in any of
    # target_col / continuous_cols / categorical_cols / text_cols. The
    # bundled web UI's "ignore" role works the same way under the hood, so
    # ignore_cols here is purely a client-side bookkeeping/validation aid:
    # it is never sent to the API, but any column listed there is checked
    # against the other role lists so a caller can't accidentally both
    # "ignore" and "use" the same column.
    role_map: dict[str, list[str]] = {}
    if target_col:
        role_map.setdefault(target_col, []).append("target")
    for c in continuous_cols:
        role_map.setdefault(c, []).append("continuous")
    for c in categorical_cols:
        role_map.setdefault(c, []).append("categorical")
    for c in text_cols:
        role_map.setdefault(c, []).append("text")
    for c in ignore_cols:
        role_map.setdefault(c, []).append("ignore")

    conflicts = {c: roles for c, roles in role_map.items() if len(roles) > 1}
    if conflicts:
        detail = "; ".join(f"{c!r} -> {roles}" for c, roles in conflicts.items())
        raise ToolError(
            f"generate_synthetic_data: each column may have only one role, "
            f"but found conflicting assignments: {detail}"
        )

    body = {
        "job_id": job_id,
        "continuous_cols": continuous_cols,
        "categorical_cols": categorical_cols,
        "text_cols": text_cols,
        "n_samples": n_samples,
        "draws": draws,
        "tune": tune,
        "family": family,
        "text_epochs": text_epochs,
    }
    if target_col:
        body["target_col"] = target_col

    return _request("POST", "/generate", "generate_synthetic_data", json=body)


def _status(job_id: str) -> dict:
    return _request("GET", f"/status/{job_id}", "check_job_status")


def _download(job_id: str, save_path: str) -> dict:
    url = f"{API_BASE_URL}/download/{job_id}"
    try:
        resp = requests.get(url, headers=_headers(), timeout=DOWNLOAD_TIMEOUT)
    except requests.exceptions.ConnectionError as e:
        raise ToolError(f"download_results: could not reach DataMerase API at {API_BASE_URL} ({e}).")
    except requests.exceptions.Timeout:
        raise ToolError(f"download_results: request to {url} timed out after {DOWNLOAD_TIMEOUT}s.")
    except requests.exceptions.RequestException as e:
        raise ToolError(f"download_results: request error — {e}")

    if resp.status_code != 200:
        # Error responses (404 invalid_job / no_result, 401, etc.) use the
        # same JSON envelope as every other endpoint.
        _handle_json_response(resp, "download_results")
        raise ToolError(f"download_results: unexpected HTTP {resp.status_code} with no error payload")

    content_type = resp.headers.get("Content-Type", "")
    if "csv" not in content_type and "octet-stream" not in content_type:
        raise ToolError(
            f"download_results: expected a CSV file but got Content-Type {content_type!r} "
            f"(HTTP {resp.status_code})"
        )

    out_path = Path(save_path).expanduser().resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_bytes(resp.content)

    text = resp.content.decode("utf-8", errors="replace")
    reader = csv.reader(io.StringIO(text))
    rows = list(reader)
    header = rows[0] if rows else []
    n_rows = max(len(rows) - 1, 0)
    preview_lines = text.splitlines()[:6]

    return {
        "saved_to": str(out_path),
        "columns": header,
        "n_rows": n_rows,
        "size_bytes": len(resp.content),
        "preview": "\n".join(preview_lines),
    }


# ---------------------------------------------------------------------------
# MCP tools
# ---------------------------------------------------------------------------


@mcp.tool()
def check_server_health() -> dict:
    """Check whether the DataMerase API is up and get basic capacity info.

    Returns server status, GPU availability, and the number of currently
    running / total tracked jobs. Useful before starting a generation job
    to see if the server is already near its concurrency cap.
    """
    return _request("GET", "/health", "check_server_health")


@mcp.tool()
def upload_dataset(file_path: str) -> dict:
    """Upload a local CSV or Excel file to DataMerase.

    Args:
        file_path: Absolute or user-relative path to a local .csv, .xlsx,
            or .xls file (max 50MB).

    Returns a job_id (needed by every other tool), the list of detected
    columns, an auto-detected role guess per column (continuous /
    categorical / text), the row count, and any data-quality warnings
    (e.g. columns with missing values).
    """
    return _upload(file_path)


@mcp.tool()
def generate_synthetic_data(
    job_id: str,
    target_col: str | None = None,
    continuous_cols: list[str] | None = None,
    categorical_cols: list[str] | None = None,
    text_cols: list[str] | None = None,
    ignore_cols: list[str] | None = None,
    n_samples: int = 3000,
    draws: int = 1000,
    tune: int = 1000,
    family: str = "auto",
    text_epochs: int = 3,
) -> dict:
    """Start an asynchronous synthetic-data generation job for an uploaded dataset.

    Assign every column from the upload to exactly one role:
        target_col        the single numeric column to model as the outcome
                           (optional if you are only generating text_cols)
        continuous_cols    other numeric columns to synthesize via vine copula
        categorical_cols   categorical columns to synthesize
        text_cols          free-text columns to synthesize via GPT-2
        ignore_cols        columns to leave out of the synthetic output
                           entirely (mirrors the "ignore" role in the
                           DataMerase web UI) — list them here explicitly so
                           conflicting assignments are caught early; they are
                           never sent to the API, a column is excluded simply
                           by not appearing in any other list.

    Args:
        job_id: job_id returned by upload_dataset.
        n_samples: number of synthetic rows to generate (default 3000).
        draws, tune: PyMC sampler draws/tuning steps for the numeric model
            (defaults 1000/1000; lower values generate faster but with a
            noisier posterior).
        family: statistical family for target_col ("auto" to auto-detect,
            or one of the supported families, e.g. "gaussian", "poisson",
            "bernoulli").
        text_epochs: fine-tuning epochs per text column (default 3).

    Returns {"status": "started", "job_id": ...} — poll check_job_status
    for progress. Starting a job on one already running/done raises a
    clean error (409 job_conflict), and starting a 4th job while 3 are
    already running raises a clean error (429 capacity_exceeded) — both
    are surfaced verbatim from the API.
    """
    return _generate(
        job_id, target_col, continuous_cols, categorical_cols, text_cols,
        ignore_cols, n_samples, draws, tune, family, text_epochs,
    )


@mcp.tool()
def check_job_status(job_id: str) -> dict:
    """Poll the status of a generation job started by generate_synthetic_data.

    Args:
        job_id: the job to check.

    Returns status ("uploaded" | "running" | "done" | "error"), a log of
    progress messages, the evaluation metrics dict once done (correlation
    diff, KS statistic, TSTR RMSE/R^2, duplicate rate, etc.), and an error
    string if the job failed.
    """
    return _status(job_id)


@mcp.tool()
def download_results(job_id: str, save_path: str) -> dict:
    """Download the synthetic_data.csv for a completed job to a local path.

    Args:
        job_id: a job whose status is "done" (check with check_job_status first).
        save_path: local file path to write the CSV to; parent directories
            are created if needed.

    Returns the saved path, detected columns, row count, file size, and a
    short text preview of the first few lines.
    """
    return _download(job_id, save_path)


@mcp.tool()
def generate_and_wait(
    file_path: str,
    save_path: str,
    target_col: str | None = None,
    continuous_cols: list[str] | None = None,
    categorical_cols: list[str] | None = None,
    text_cols: list[str] | None = None,
    ignore_cols: list[str] | None = None,
    n_samples: int = 3000,
    draws: int = 1000,
    tune: int = 1000,
    family: str = "auto",
    text_epochs: int = 3,
    poll_interval_seconds: float = 3.0,
    timeout_seconds: float = 900.0,
) -> dict:
    """Run the entire DataMerase pipeline in one call: upload, generate, wait, download.

    Uploads file_path, starts a generation job with the given column
    configuration (see generate_synthetic_data for what target_col /
    continuous_cols / categorical_cols / text_cols / ignore_cols mean),
    polls check_job_status until the job reaches "done" or "error", then
    downloads the resulting synthetic CSV to save_path.

    Args:
        file_path: local CSV/Excel file to upload.
        save_path: local path to save the synthetic CSV to on success.
        poll_interval_seconds: delay between status polls (default 3s).
        timeout_seconds: give up and raise if the job hasn't finished
            within this many seconds (default 900s / 15min).

    Returns a dict with the upload info, the final status payload
    (including metrics), and the download info (including a preview of
    the generated rows) — everything needed to report the result back to
    the user in one shot.
    """
    upload_info = _upload(file_path)
    job_id = upload_info["job_id"]

    _generate(
        job_id, target_col, continuous_cols, categorical_cols, text_cols,
        ignore_cols, n_samples, draws, tune, family, text_epochs,
    )

    deadline = time.monotonic() + timeout_seconds
    status_info = _status(job_id)
    while status_info["status"] not in ("done", "error"):
        if time.monotonic() > deadline:
            raise ToolError(
                f"generate_and_wait: job {job_id} did not finish within "
                f"{timeout_seconds:.0f}s (last status: {status_info['status']!r}). "
                f"It may still complete — poll check_job_status('{job_id}') later."
            )
        time.sleep(poll_interval_seconds)
        status_info = _status(job_id)

    if status_info["status"] == "error":
        raise ToolError(f"generate_and_wait: job {job_id} failed — {status_info.get('error')}")

    download_info = _download(job_id, save_path)

    return {
        "job_id": job_id,
        "upload": upload_info,
        "final_status": status_info,
        "metrics": status_info.get("metrics"),
        "download": download_info,
    }


if __name__ == "__main__":
    mcp.run(transport="stdio")