# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import os
import platform
import re
import socket
import sys
from concurrent.futures import ThreadPoolExecutor
from math import isfinite
from pathlib import Path
from time import sleep, time

from ultralytics.utils import (
    ENVIRONMENT,
    GIT,
    LOGGER,
    PLATFORM_URL,
    PYTHON_VERSION,
    RANK,
    SETTINGS,
    TESTS_RUNNING,
    Retry,
    colorstr,
)
from ultralytics.utils.events import events

PREFIX = colorstr("Platform: ")
PLATFORM_API_URL = os.getenv("PLATFORM_API_URL", f"{PLATFORM_URL}/api/webhooks")


def slugify(text):
    """Convert text to URL-safe slug (e.g., 'My Project 1' -> 'my-project-1')."""
    if not text:
        return text
    return re.sub(r"-+", "-", re.sub(r"[^a-z0-9\s-]", "", str(text).lower()).replace(" ", "-")).strip("-")[:128]


try:
    assert not TESTS_RUNNING  # do not log pytest
    assert SETTINGS.get("platform", False) is True or os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
    _api_key = os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
    assert _api_key  # verify API key is present

    import requests

    from ultralytics.utils.logger import ConsoleLogger, SystemLogger
    from ultralytics.utils.torch_utils import model_info_for_loggers

    _executor = ThreadPoolExecutor(max_workers=10)  # Bounded thread pool for async operations

except (AssertionError, ImportError):
    _api_key = None


def resolve_platform_uri(uri, hard=True):
    """Resolve ul:// URIs to signed URLs by authenticating with Ultralytics Platform.

    Formats:
        ul://username/datasets/slug  -> Returns signed URL to NDJSON file
        ul://username/project/model  -> Returns signed URL to .pt file

    Args:
        uri (str): Platform URI starting with "ul://".
        hard (bool): Whether to raise an error if resolution fails.

    Returns:
        (str | None): Signed URL on success, None if not found and hard=False.

    Raises:
        ValueError: If API key is missing/invalid or URI format is wrong.
        PermissionError: If access is denied.
        RuntimeError: If resource is not ready (e.g., dataset still processing).
        FileNotFoundError: If resource not found and hard=True.
        ConnectionError: If network request fails and hard=True.
    """
    import requests

    path = uri[5:]  # Remove "ul://"
    parts = path.split("/")

    api_key = os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
    if not api_key:
        raise ValueError(f"ULTRALYTICS_API_KEY required for '{uri}'. Get key at {PLATFORM_URL}/settings")

    base = PLATFORM_API_URL
    headers = {"Authorization": f"Bearer {api_key}"}

    # ul://username/datasets/slug
    if len(parts) == 3 and parts[1] == "datasets":
        username, _, slug = parts
        url = f"{base}/datasets/{username}/{slug}/export"

    # ul://username/project/model
    elif len(parts) == 3:
        username, project, model = parts
        url = f"{base}/models/{username}/{project}/{model}/download"

    else:
        raise ValueError(f"Invalid platform URI: {uri}. Use ul://user/datasets/name or ul://user/project/model")

    # (connect_timeout, read_timeout) — short connect so retries are fast, long read for server-side generation
    timeout = (10, 3600) if "/datasets/" in url else (10, 90)

    try:
        for attempt in range(5):
            try:
                # GET not HEAD: a HEAD response carries no body, so every platform error message was lost.
                r = requests.get(url, headers=headers, allow_redirects=False, timeout=timeout)
                if r.status_code in {408, 429} or r.status_code >= 500:
                    raise requests.exceptions.HTTPError(f"HTTP {r.status_code}", response=r)
                break
            except (
                requests.exceptions.ConnectionError,
                requests.exceptions.ReadTimeout,
                requests.exceptions.HTTPError,
            ) as e:
                if attempt >= 4:
                    raise
                delay = 2 * (2**attempt)  # 2s, 4s, 8s, 16s backoff
                LOGGER.warning(f"Retry {attempt + 1}/5 for {uri} in {delay}s: {e}")
                sleep(delay)
    except Exception as e:
        if hard:
            raise ConnectionError(f"Failed to resolve {uri}: {e}") from e
        LOGGER.warning(f"Failed to resolve {uri}: {e}")
        return None

    # Handle redirect responses (301, 302, 303, 307, 308)
    if 300 <= r.status_code < 400 and "location" in r.headers:
        return r.headers["location"]  # Return signed URL

    # Echo only the platform's own JSON `error`: a proxy or WAF error page can quote our Authorization
    # header, and it would land in the user's console.
    try:
        detail = str(r.json().get("error", "")).strip()
    except Exception:
        detail = ""
    detail = f" {detail[:500]}" if detail else ""

    if r.status_code == 401:
        raise ValueError(f"Invalid ULTRALYTICS_API_KEY for '{uri}'.{detail}")
    if r.status_code == 403:
        raise PermissionError(f"Access denied for '{uri}'. Check dataset/model visibility settings.{detail}")
    if r.status_code == 404:
        if hard:
            raise FileNotFoundError(f"Not found on platform: {uri}.{detail}")
        LOGGER.warning(f"Not found on platform: {uri}.{detail}")
        return None
    if r.status_code == 409:
        raise RuntimeError(f"Resource not ready: {uri}. Dataset may still be processing.{detail}")

    raise RuntimeError(f"Platform error for '{uri}' (HTTP {r.status_code}).{detail or f' {r.reason}'}")


def _interp_plot(plot, n=101):
    """Interpolate plot curve data to n points to reduce storage size."""
    import numpy as np

    if not plot.get("x") or not plot.get("y"):
        return plot  # No interpolation needed (e.g., confusion_matrix)

    x, y = np.array(plot["x"]), np.array(plot["y"])
    if len(x) <= n:
        return plot  # Already small enough

    # New x values (101 points gives clean 0.01 increments: 0, 0.01, 0.02, ..., 1.0)
    x_new = np.linspace(x[0], x[-1], n)

    # Interpolate y values (handle both 1D and 2D arrays)
    if y.ndim == 1:
        y_new = np.interp(x_new, x, y)
    else:
        y_new = np.array([np.interp(x_new, x, yi) for yi in y])

    # Also interpolate ap if present (for PR curves)
    result = {**plot, "x": x_new.tolist(), "y": y_new.tolist()}
    if "ap" in plot:
        result["ap"] = plot["ap"]  # Keep AP values as-is (per-class scalars)

    return result


def _validation_payload(image_metrics, sample_limit=5_000, extremes_limit=100):
    """Return exact F1 extremes and an evenly ranked sample for correlation analysis."""
    ranked = sorted(image_metrics.items(), key=lambda item: (item[1]["f1"], item[0]))
    if len(ranked) > sample_limit:
        sample = [ranked[round(i * (len(ranked) - 1) / (sample_limit - 1))] for i in range(sample_limit)]
    else:
        sample = ranked

    def rows(items):
        return [[Path(name).stem.split("_", 1)[0], metric["tp"], metric["fp"], metric["fn"]] for name, metric in items]

    return {
        "population": len(ranked),
        "sampling": "f1_rank",
        "rows": rows(sample),
        "extremes": {"worst": rows(ranked[:extremes_limit]), "best": rows(reversed(ranked[-extremes_limit:]))},
    }


def _sanitize_json_value(value):
    """Replace non-finite floats in payloads with None so requests JSON encoding succeeds."""
    if isinstance(value, dict):
        return {k: _sanitize_json_value(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_sanitize_json_value(v) for v in value]
    if isinstance(value, float):
        return value if isfinite(value) else None  # avoid "Out of range float values are not JSON compliant" warnings
    return value


def _send(event, data, project, name, model_id=None, retry=2, timeout=30):
    """Send event to Platform endpoint with retry logic."""
    if not _api_key:
        return None
    payload = {"event": event, "project": project, "name": name, "data": _sanitize_json_value(data)}
    if model_id:
        payload["modelId"] = model_id

    def send_once():
        global _api_key
        r = requests.post(
            f"{PLATFORM_API_URL}/training/metrics",
            json=payload,
            headers={"Authorization": f"Bearer {_api_key}"},
            timeout=timeout,
        )
        if 400 <= r.status_code < 500 and r.status_code not in {408, 429}:
            try:
                msg = r.json().get("error", r.reason)
            except Exception:
                msg = r.reason
            # Only 401 is credential-scoped; 403/404 concern one run and must not disable the process.
            if r.status_code == 401:
                _api_key = None
            # A console_output failure must not be logged: ConsoleLogger flushes the warning back as the
            # next chunk, which fails again. 401 is safe — the cleared key short-circuits _send.
            if event != "console_output" or r.status_code == 401:
                LOGGER.warning(f"{PREFIX}{msg}")
            return None  # Don't retry client errors (except 408 timeout, 429 rate limit)
        r.raise_for_status()
        return r.json()

    # Same loop as above, so a console_output send stays silent at every level — including Retry's
    # per-attempt warning. It must still retry: _flush_buffer clears the buffer before calling us.
    quiet = event == "console_output"
    try:
        return Retry(times=retry, delay=1, verbose=not quiet)(send_once)()
    except Exception as e:
        if not quiet:
            LOGGER.debug(f"{PREFIX}Failed to send {event}: {e}")
        return None


def _send_async(event, data, project, name, model_id=None):
    """Send event asynchronously using bounded thread pool."""
    _executor.submit(_send, event, data, project, name, model_id)


def _handle_control_response(trainer, ctx, response):
    """Apply centralized stop signals returned by Platform webhook responses.

    Notes:
        ``ctx["cancelled"]`` is the durable cancellation signal. During startup, trainer setup later resets
        ``trainer.stop``, so early stop requests still rely on ``on_pretrain_routine_end()`` to reapply the flag after
        setup completes.
    """
    if response and response.get("cancelled"):
        ctx["cancelled"] = True
        trainer.stop = True
        LOGGER.info(f"{PREFIX}Training cancelled from Platform ⚠️")


def _upload_model(model_path, project, name, progress=False, retry=1, model_id=None, run_id=None):
    """Publish a model checkpoint to its configured Platform storage location."""
    from ultralytics.utils.uploads import safe_upload

    if not _api_key:
        return None
    model_path = Path(model_path)
    if not model_path.exists():
        LOGGER.warning(f"{PREFIX}Model file not found: {model_path}")
        return None
    model_size = model_path.stat().st_size
    if os.getenv("PLATFORM_API_URL"):
        return {"modelPath": str(model_path.resolve()), "modelSize": model_size}

    # Get signed upload URL from Platform (server sanitizes filename for storage safety)
    @Retry(times=3, delay=2)
    def get_signed_url():
        payload = {"project": project, "name": name, "filename": model_path.name}
        if model_id:
            payload["modelId"] = model_id  # Direct lookup avoids slug mismatch from auto-increment
        if run_id:
            payload["runId"] = run_id
        r = requests.post(
            f"{PLATFORM_API_URL}/models/upload",
            json=payload,
            headers={"Authorization": f"Bearer {_api_key}"},
            timeout=30,
        )
        r.raise_for_status()
        return r.json()

    try:
        data = get_signed_url()
    except Exception as e:
        LOGGER.warning(f"{PREFIX}Failed to get upload URL: {e}")
        return None

    # Upload to GCS using safe_upload with retry logic and optional progress bar
    if safe_upload(file=model_path, url=data["uploadUrl"], retry=retry, progress=progress):
        gcs_path = data.get("gcsPath")
        if gcs_path and run_id:
            saved = _send(
                "checkpoint_saved",
                {
                    "modelPath": gcs_path,
                    "runId": run_id,
                    "uploadPath": data.get("uploadPath"),
                },
                project,
                name,
                model_id,
                timeout=90,
            )
            return {"modelPath": gcs_path, "modelSize": model_size} if saved else None
        return {"modelPath": gcs_path, "modelSize": model_size}
    return None


def _get_environment_info():
    """Collect comprehensive environment info using existing ultralytics utilities."""
    import shutil

    import psutil
    import torch

    from ultralytics import __version__
    from ultralytics.utils.torch_utils import get_cpu_info, get_gpu_info

    # Get RAM and disk totals
    memory = psutil.virtual_memory()
    disk_usage = shutil.disk_usage("/")

    env = {
        "ultralyticsVersion": __version__,
        "hostname": socket.gethostname(),
        "os": platform.platform(),
        "environment": ENVIRONMENT,
        "pythonVersion": PYTHON_VERSION,
        "pythonExecutable": sys.executable,
        "cpuCount": os.cpu_count() or 0,
        "cpu": get_cpu_info(),
        "command": " ".join(sys.argv),
        "totalRamGb": round(memory.total / (1 << 30), 1),  # Total RAM in GB
        "totalDiskGb": round(disk_usage.total / (1 << 30), 1),  # Total disk in GB
    }

    # Git info using cached GIT singleton (no subprocess calls)
    try:
        if GIT.is_repo:
            if GIT.origin:
                env["gitRepository"] = GIT.origin
            if GIT.branch:
                env["gitBranch"] = GIT.branch
            if GIT.commit:
                env["gitCommit"] = GIT.commit[:12]  # Short hash
            if GIT.message:
                env["gitCommitMessage"] = GIT.message
    except Exception:
        pass

    # GPU info
    try:
        if torch.cuda.is_available():
            env["gpuCount"] = torch.cuda.device_count()
            env["gpuType"] = get_gpu_info(0) if torch.cuda.device_count() > 0 else None
    except Exception:
        pass

    return env


def _get_project_name(trainer):
    """Get slugified project and name from trainer args."""
    raw = str(trainer.args.project)
    parts = raw.split("/", 1)
    project = f"{parts[0]}/{slugify(parts[1])}" if len(parts) == 2 else slugify(raw)
    return project, slugify(str(trainer.args.name or "train"))


def on_pretrain_routine_start(trainer):
    """Initialize Platform logging at training start."""
    if RANK not in {-1, 0} or not trainer.args.project:
        return

    project, name = _get_project_name(trainer)
    LOGGER.info(f"{PREFIX}Streaming training metrics to Platform")

    # Single dict for all platform callback state (like trainer.hub_session for HUB callbacks)
    ctx = {
        "model_id": None,
        "run_id": None,
        "last_upload": time(),
        "checkpoint_upload": None,
        "cancelled": False,
        "console_logger": None,
        "system_logger": None,
    }
    trainer.platform = ctx

    # Create callback to send console output to Platform
    def send_console_output(content, line_count, chunk_id):
        """Send batched console output to Platform webhook."""
        _send_async(
            "console_output",
            {"chunkId": chunk_id, "content": content, "lineCount": line_count},
            project,
            name,
            ctx["model_id"],
        )

    # Console capture with batching (5 lines or 5 seconds). Built here, but not started until Platform
    # has accepted the run below: capturing first left the user's stdout redirected through a dead
    # integration whenever training_started failed, and its final flush would post a console chunk
    # carrying no model_id.
    ctx["console_logger"] = ConsoleLogger(batch_size=5, flush_interval=5.0, on_flush=send_console_output)

    # Collect environment info (W&B-style metadata)
    environment = _get_environment_info()

    # Build trainArgs - callback runs before get_dataset() so args.data is still original (e.g., ul:// URIs)
    # Note: model_info is sent later in on_fit_epoch_end (epoch 0) when the model is actually loaded
    train_args = {k: str(v) for k, v in vars(trainer.args).items()}

    # Send synchronously to get modelId for subsequent webhooks (critical, more retries)
    response = _send(
        "training_started",
        {
            "trainArgs": train_args,
            "epochs": trainer.epochs,
            "device": str(trainer.device),
            "environment": environment,
        },
        project,
        name,
        retry=4,
    )
    if response and response.get("modelId"):
        ctx["model_id"] = response["modelId"]
        ctx["run_id"] = response.get("runId")
        # Server returns actual slug (may differ from requested name due to auto-increment, e.g. "train" → "train-2")
        if response.get("modelSlug"):
            ctx["model_slug"] = response["modelSlug"]
            url = f"{PLATFORM_URL}/{project}/{ctx['model_slug']}"
            LOGGER.info(f"{PREFIX}View model at {url}")
        ctx["console_logger"].start_capture()  # only now: the run is tracked and model_id is known
        # Note: trainer.stop is set in on_pretrain_routine_end (after _setup_train resets it)
        _handle_control_response(trainer, ctx, response)
    else:
        LOGGER.warning(f"{PREFIX}Training will not be tracked on Platform")
        trainer.platform = None  # Disable further callbacks


def on_pretrain_routine_end(trainer):
    """Apply pre-start cancellation after _setup_train resets trainer.stop."""
    ctx = getattr(trainer, "platform", None)
    if ctx and ctx["cancelled"]:
        LOGGER.info(f"{PREFIX}Training cancelled from Platform before starting ✅")
        trainer.stop = True


def on_fit_epoch_end(trainer):
    """Log training and system metrics at epoch end."""
    ctx = getattr(trainer, "platform", None)
    if not ctx or RANK not in {-1, 0} or not trainer.args.project:
        return

    project, name = _get_project_name(trainer)
    metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics}

    if trainer.optimizer and trainer.optimizer.param_groups:
        metrics["lr"] = trainer.optimizer.param_groups[0]["lr"]

    # Extract model info at epoch 0 (sent as separate field, not in metrics)
    model_info = None
    if trainer.epoch == 0:
        try:
            info = model_info_for_loggers(trainer)
            model_info = {
                "parameters": info.get("model/parameters", 0),
                "gflops": info.get("model/GFLOPs", 0),
                "speedMs": info.get("model/speed_PyTorch(ms)", 0),
            }
        except Exception:
            pass

    # Get system metrics (cache SystemLogger in platform context for efficiency)
    system = {}
    try:
        if not ctx["system_logger"]:
            ctx["system_logger"] = SystemLogger(all_drives=True)
        system = ctx["system_logger"].get_metrics(rates=True)
    except Exception:
        pass

    payload = {
        "epoch": trainer.epoch,
        "metrics": metrics,
        "system": system,
        "fitness": trainer.fitness,
        "best_fitness": trainer.best_fitness,
    }
    if model_info:
        payload["modelInfo"] = model_info

    def _send_and_check_cancel():
        """Send epoch_end and check response for cancellation (runs in background thread)."""
        response = _send("epoch_end", payload, project, name, ctx["model_id"], retry=1)
        _handle_control_response(trainer, ctx, response)

    _executor.submit(_send_and_check_cancel)


def on_model_save(trainer):
    """Upload model checkpoint (rate limited to every 15 min)."""
    ctx = getattr(trainer, "platform", None)
    if not ctx or RANK not in {-1, 0} or not trainer.args.project:
        return
    # Rate limit to every 15 minutes (900 seconds)
    if time() - ctx["last_upload"] < 900:
        return
    if ctx["checkpoint_upload"] and not ctx["checkpoint_upload"].done():
        return

    model_path = trainer.best if trainer.best and Path(trainer.best).exists() else trainer.last
    if not model_path:
        return

    project, name = _get_project_name(trainer)
    ctx["checkpoint_upload"] = _executor.submit(
        _upload_model, model_path, project, name, model_id=ctx["model_id"], run_id=ctx["run_id"]
    )
    ctx["last_upload"] = time()


def on_train_end(trainer):
    """Run events on train end once fitness and duration are known, then log final results and upload best model."""
    events(trainer.args, trainer.device, trainer)
    ctx = getattr(trainer, "platform", None)  # set only by on_pretrain_routine_start, so unset without an API key
    if not ctx or RANK not in {-1, 0} or not trainer.args.project:
        return

    project, name = _get_project_name(trainer)

    if ctx["cancelled"]:
        LOGGER.info(f"{PREFIX}Uploading partial results for cancelled training")

    # Stop console capture
    if ctx["console_logger"]:
        ctx["console_logger"].stop_capture()
        ctx["console_logger"] = None

    # Upload best model (blocking with progress bar to ensure it completes)
    artifact = None
    if trainer.best and Path(trainer.best).exists():
        if ctx["checkpoint_upload"]:
            ctx["checkpoint_upload"].result()
        artifact = _upload_model(
            trainer.best,
            project,
            name,
            progress=True,
            retry=3,
            model_id=ctx["model_id"],
            run_id=ctx["run_id"],
        )
        if not artifact:
            LOGGER.warning(f"{PREFIX}Model will not be available for download on Platform (upload failed)")

    # Collect plots from trainer and validator, deduplicating by type
    plots_by_type = {}
    for info in getattr(trainer, "plots", {}).values():
        if info.get("data") and info["data"].get("type"):
            plots_by_type[info["data"]["type"]] = info["data"]
    for info in getattr(getattr(trainer, "validator", None), "plots", {}).values():
        if info.get("data") and info["data"].get("type"):
            plots_by_type.setdefault(info["data"]["type"], info["data"])  # Don't overwrite trainer plots
    plots = [_interp_plot(p) for p in plots_by_type.values()]  # Interpolate curves to reduce size

    # Get class names
    names = getattr(getattr(trainer, "validator", None), "names", None) or (trainer.data or {}).get("names")
    class_names = list(names.values()) if isinstance(names, dict) else list(names) if names else None

    # stopper.best_epoch is 1-indexed; -1 aligns with the 0-indexed `epoch` field
    best_epoch = max(0, getattr(getattr(trainer, "stopper", None), "best_epoch", trainer.epoch + 1) - 1)

    image_metrics = trainer.validator.metrics.box.image_metrics if trainer.args.task == "detect" else {}
    validation = _validation_payload(image_metrics)
    _send(
        "training_complete",
        {
            "results": {
                "metrics": {**trainer.metrics, "fitness": trainer.fitness},
                "bestEpoch": best_epoch,
                "bestFitness": trainer.best_fitness,
                **({"calibration": c} if (c := getattr(trainer, "depth_calibration", None)) else {}),
                **({"validation": validation} if validation["rows"] else {}),
                **(artifact or {}),
            },
            "classNames": class_names,
            "plots": plots,
            "runId": ctx["run_id"],
        },
        project,
        name,
        ctx["model_id"],
        retry=4,  # Critical, more retries
    )
    url = f"{PLATFORM_URL}/{project}/{ctx.get('model_slug', name)}"
    LOGGER.info(f"{PREFIX}View results at {url}")


def on_val_start(validator):
    """Run events on validation start, when the user asked for validation.

    A trainer runs its own final validation on a copy of the trainer's args, whose mode is still 'train'. Since an event
    is named for its mode, firing here would send a second 'train' event indistinguishable from the training one.
    """
    if validator.args.mode == "val":
        events(validator.args, validator.device)


def on_predict_end(predictor):
    """Run events on predict end, once per-image speeds are known."""
    events(predictor.args, predictor.device, predictor)


def on_export_start(exporter):
    """Run events on export start."""
    events(exporter.args, exporter.device)


callbacks = {
    # Anonymous analytics, gated only by the documented sync setting inside Events
    "on_val_start": on_val_start,
    "on_predict_end": on_predict_end,
    "on_export_start": on_export_start,
    "on_train_end": on_train_end,  # sends the train event, then uploads results if a Platform run is in flight
    **(
        {
            "on_pretrain_routine_start": on_pretrain_routine_start,
            "on_pretrain_routine_end": on_pretrain_routine_end,
            "on_fit_epoch_end": on_fit_epoch_end,
            "on_model_save": on_model_save,
        }
        if _api_key
        else {}
    ),
}
