Monitor

@synalinks_export("synalinks.callbacks.Monitor")
class Monitor(Callback):
    """Monitor callback for logging training metrics to MLflow.

    This callback logs training progress and evaluation metrics to MLflow
    for experiment tracking and visualization.

    Args:
        experiment_name (str): Name of the MLflow experiment. If None, uses
            the experiment of `synalinks.enable_observability()` when it was
            called (so runs land next to their traces), else the program name.
        run_name (str): Name of the MLflow run. If None, auto-generated.
        tracking_uri (str): MLflow tracking server URI. If None, uses the
            value from `synalinks.enable_observability()` or the default
            (local ./mlruns directory or MLFLOW_TRACKING_URI env var).
        log_batch_metrics (bool): Whether to log metrics at batch level
            (default: False).
        log_epoch_metrics (bool): Whether to log metrics at epoch level
            (default: True).
        log_program_plot (bool): Whether to log the program plot as an artifact
            at the beginning of training (default: True).
        log_program_model (bool): Whether to log the program as an MLflow model
            at the end of training (default: True).
        tags (dict): Optional tags to add to the MLflow run.
        run_id (str): Optional. The id of an existing MLflow run to resume
            instead of starting a new one. Metrics keep being appended to it,
            the step counter continuing after the last step already logged.
        resume (bool): Whether to look up an existing run named `run_name`
            in the experiment and resume it (default: False). Creates the
            run the first time. This is how repeated `evaluate()` calls
            draw a chart over time: each evaluation adds one point to the
            metrics of the same run.
        log_assessments (bool): Whether to log each evaluated sample's reward
            as a `reward` feedback assessment on the sample's trace during
            `evaluate()` (default: True). Requires the traces of
            `synalinks.enable_observability()`; the assessments show up on
            the traces of the evaluation run in the MLflow UI.

    Example:

    ```python
    import synalinks

    # Basic usage - uses local MLflow storage
    monitor = synalinks.callbacks.Monitor(experiment_name="my_experiment")

    # With custom MLflow tracking server
    monitor = synalinks.callbacks.Monitor(
        tracking_uri="http://localhost:5000",
        experiment_name="my_experiment",
        run_name="training_run_1",
        log_program_plot=True,
        log_program_model=True,
        tags={"model_type": "chain_of_thought"}
    )

    # Use in training
    program.fit(
        x=train_data,
        y=train_labels,
        epochs=10,
        callbacks=[monitor]
    )

    # Track evaluation results over time: every evaluate() call (in this
    # process or a later one) appends one point to the same run's charts
    monitor = synalinks.callbacks.Monitor(
        experiment_name="my_experiment",
        run_name="nightly_eval",
        resume=True,
    )
    program.evaluate(x=test_data, y=test_labels, callbacks=[monitor])
    ```

    Note:
        For tracing module calls along with training metrics, use
        `synalinks.enable_observability()` at the beggining of your script
        which configures the Monitor hook & callback:

        ```python
        synalinks.enable_observability(
            tracking_uri="http://localhost:5000",
            experiment_name="my_traces"
        )
        ```
    """

    def __init__(
        self,
        experiment_name=None,
        run_name=None,
        tracking_uri=None,
        log_batch_metrics=False,
        log_epoch_metrics=True,
        log_program_plot=True,
        log_program_model=True,
        tags=None,
        run_id=None,
        resume=False,
        log_assessments=True,
    ):
        super().__init__()
        if not MLFLOW_AVAILABLE:
            raise ImportError(
                "mlflow is required for the Monitor callback. "
                "Install it with: pip install mlflow"
            )

        self.experiment_name = experiment_name
        self.run_name = run_name
        self.tracking_uri = tracking_uri or mlflow_tracking_uri()
        self.log_batch_metrics = log_batch_metrics
        self.log_epoch_metrics = log_epoch_metrics
        self.log_program_plot = log_program_plot
        self.log_program_model = log_program_model
        self.tags = tags or {}
        self.run_id = run_id
        self.resume = resume
        self.log_assessments = log_assessments
        self.logger = logging.getLogger(__name__)

        self._run = None
        self._trace_mark = None
        self._step = 0
        self._epoch = 0
        # Track if we're inside fit() to avoid ending run during validation
        self._in_training = False

    def _setup_mlflow(self):
        """Configure MLflow tracking."""
        if self.tracking_uri:
            mlflow.set_tracking_uri(self.tracking_uri)

        experiment_name = self.experiment_name
        if experiment_name is None and is_observability_enabled():
            experiment_name = mlflow_experiment_name()
        if experiment_name is None and self.program is not None:
            experiment_name = self.program.name or "synalinks_experiment"

        self._experiment_id = mlflow.set_experiment(experiment_name).experiment_id

    def _start_run(self, run_name_suffix=""):
        """Start a new MLflow run, or resume one (`run_id` / `resume`)."""
        run_name = self.run_name
        if run_name and run_name_suffix:
            run_name = f"{run_name}_{run_name_suffix}"
        elif run_name_suffix:
            run_name = run_name_suffix

        run_id = self.run_id
        if run_id is None and self.resume:
            client = mlflow.MlflowClient()
            found = client.search_runs(
                experiment_ids=[self._experiment_id],
                filter_string=f"tags.mlflow.runName = '{run_name}'",
                order_by=["attributes.start_time DESC"],
                max_results=1,
            )
            if found:
                run_id = found[0].info.run_id

        if run_id is not None:
            self._run = mlflow.start_run(run_id=run_id)
            self.run_id = run_id
        else:
            self._run = mlflow.start_run(run_name=run_name)
            self.run_id = self._run.info.run_id

        tags = dict(self.tags)
        if self.program is not None:
            if self.program.name:
                tags["program_name"] = self.program.name
            if self.program.description:
                tags["program_description"] = self.program.description

        if tags:
            mlflow.set_tags(tags)

        self._step = 0
        self._epoch = 0
        if run_id is not None:
            client = mlflow.MlflowClient()
            for key in client.get_run(run_id).data.metrics:
                history = client.get_metric_history(run_id, key)
                self._step = max(self._step, *(m.step + 1 for m in history))

    def _end_run(self):
        """End the current MLflow run."""
        if self._run is not None:
            mlflow.end_run()
            self._run = None

    async def _log_metrics(self, logs, step=None):
        """Log metrics to MLflow asynchronously."""
        if logs is None or self._run is None:
            return

        metrics = {}
        for key, value in logs.items():
            if isinstance(value, (int, float)):
                metrics[key] = value

        if metrics:
            # Explicit run_id: MLflow's active run is thread-local, not seen by the worker
            await asyncio.to_thread(
                mlflow.log_metrics, metrics, step=step, run_id=self._run.info.run_id
            )

    async def _upload_artifact_via_http(self, local_path, artifact_path, run_id):
        """Upload artifact via HTTP to MLflow server asynchronously.

        This method uses the MLflow REST API to upload artifacts directly,
        bypassing local filesystem artifact repo issues. Requires the MLflow
        server to be started with --serve-artifacts flag.
        """
        import requests

        if not self.tracking_uri:
            raise ValueError("tracking_uri is required for HTTP artifact upload")

        # Get the run's artifact URI to determine the correct upload path
        client = mlflow.MlflowClient(tracking_uri=self.tracking_uri)
        run = await asyncio.to_thread(client.get_run, run_id)
        artifact_uri = run.info.artifact_uri

        filename = os.path.basename(local_path)
        if artifact_path:
            full_artifact_path = f"{artifact_path}/{filename}"
        else:
            full_artifact_path = filename

        # Parse the artifact URI to construct the correct upload URL
        # artifact_uri can be:
        #   - mlflow-artifacts:/<experiment_id>/<run_id>/artifacts
        #   - mlflow-artifacts://host:port/<experiment_id>/<run_id>/artifacts
        #   - /mlflow/artifacts/<experiment_id>/<run_id>/artifacts (server local path)
        if artifact_uri.startswith("mlflow-artifacts:"):
            # Extract the path part after the scheme
            uri_path = artifact_uri.replace("mlflow-artifacts://", "").replace(
                "mlflow-artifacts:/", ""
            )
            # Remove host:port if present (will use tracking_uri instead)
            if "/" in uri_path and not uri_path.startswith("/"):
                parts = uri_path.split("/", 1)
                if ":" in parts[0] or "." in parts[0]:
                    # First part looks like host:port, skip it
                    uri_path = parts[1] if len(parts) > 1 else ""
        elif artifact_uri.startswith("/"):
            # Server-side local path like /mlflow/artifacts/<exp_id>/<run_id>/artifacts
            # Extract the relative path: <exp_id>/<run_id>/artifacts
            # Find the pattern after the base artifacts directory
            parts = artifact_uri.split("/")
            # Look for 'artifacts' in the path and take everything after the first one
            try:
                artifacts_idx = parts.index("artifacts")
                uri_path = "/".join(parts[artifacts_idx + 1 :])
            except ValueError:
                # Fallback: use experiment_id/run_id/artifacts pattern
                uri_path = f"0/{run_id}/artifacts"
        else:
            # Fallback for other URI schemes
            uri_path = f"0/{run_id}/artifacts"

        # Construct the full upload URL
        base = f"{self.tracking_uri}/api/2.0/mlflow-artifacts/artifacts"
        url = f"{base}/{uri_path}/{full_artifact_path}"

        with open(local_path, "rb") as f:
            content = f.read()

        # Determine content type based on file extension
        content_type = "application/octet-stream"
        if local_path.endswith(".png"):
            content_type = "image/png"
        elif local_path.endswith(".json"):
            content_type = "application/json"

        headers = {"Content-Type": content_type}
        response = await asyncio.to_thread(
            requests.put, url, data=content, headers=headers
        )

        if response.status_code not in (200, 201, 204):
            raise Exception(
                f"Failed to upload artifact: {response.status_code} {response.text}"
            )

    async def _log_program_plot_artifact(self):
        """Log the program plot as an MLflow artifact asynchronously."""
        if self._run is None:
            self.logger.warning("No MLflow run active, skipping plot logging")
            return

        if self.program is None:
            self.logger.warning("No program set, skipping plot logging")
            return

        if not self.program.built:
            self.logger.warning("Program not built, skipping plot logging")
            return

        try:
            from synalinks.src.utils.program_visualization import check_graphviz
            from synalinks.src.utils.program_visualization import check_pydot
            from synalinks.src.utils.program_visualization import plot_program

            if not check_pydot() or not check_graphviz():
                self.logger.warning(
                    "pydot or graphviz not available, skipping program plot"
                )
                return

            run_id = self._run.info.run_id

            with tempfile.TemporaryDirectory() as tmpdir:
                plot_filename = f"{self.program.name or 'program'}.png"
                plot_path = os.path.join(tmpdir, plot_filename)

                # Run plot generation in thread pool
                await asyncio.to_thread(
                    plot_program,
                    self.program,
                    to_file=plot_filename,
                    to_folder=tmpdir,
                    show_schemas=True,
                    show_module_names=True,
                    show_trainable=True,
                    dpi=96,  # Lower DPI for smaller file size
                )

                if os.path.exists(plot_path):
                    # Use HTTP upload if tracking_uri is set (remote server),
                    # otherwise fall back to direct artifact logging (local)
                    if self.tracking_uri:
                        await self._upload_artifact_via_http(
                            plot_path, artifact_path="program_plots", run_id=run_id
                        )
                    else:
                        await asyncio.to_thread(
                            mlflow.log_artifact,
                            plot_path,
                            artifact_path="program_plots",
                            run_id=run_id,
                        )
                    self.logger.info(f"Logged program plot: {plot_filename}")
                else:
                    self.logger.warning(f"Plot file not created: {plot_path}")

        except Exception as e:
            self.logger.warning(f"Failed to log program plot: {e}")

    async def _log_params(self):
        """Log training hyperparameters to MLflow asynchronously."""
        if self._run is None or self.params is None:
            return

        try:
            params_to_log = {}
            for key, value in self.params.items():
                if isinstance(value, (str, int, float, bool)):
                    params_to_log[key] = value

            if params_to_log:
                await asyncio.to_thread(
                    mlflow.log_params, params_to_log, run_id=self._run.info.run_id
                )
                self.logger.debug(f"Logged params: {params_to_log}")
        except Exception as e:
            self.logger.warning(f"Failed to log params: {e}")

    async def _log_program_model(self):
        """Log the program trainable state as an MLflow artifact asynchronously.

        This saves only the trainable variables (state), not the full
        program architecture. This is useful for checkpointing the learned
        parameters like few-shot examples, optimized prompts, etc.
        """
        if self._run is None or self.program is None:
            self.logger.warning("No run or program, skipping model logging")
            return

        try:
            import orjson

            # Get the state tree (trainable, non-trainable, optimizer variables)
            state_tree = self.program.get_state_tree()

            # Create model info
            model_info = {
                "program_name": self.program.name or "program",
                "program_description": self.program.description or "",
                "framework": "synalinks",
                "num_trainable_variables": len(self.program.trainable_variables),
            }

            run_id = self._run.info.run_id

            # Write to temp files and log as artifacts
            with tempfile.TemporaryDirectory() as tmpdir:
                # Save state tree
                state_path = os.path.join(tmpdir, "state_tree.json")
                with open(state_path, "wb") as f:
                    f.write(orjson.dumps(state_tree, option=orjson.OPT_INDENT_2))

                # Save model info
                info_path = os.path.join(tmpdir, "model_info.json")
                with open(info_path, "wb") as f:
                    f.write(orjson.dumps(model_info, option=orjson.OPT_INDENT_2))

                # Upload artifacts
                if self.tracking_uri:
                    await self._upload_artifact_via_http(
                        state_path, artifact_path="model", run_id=run_id
                    )
                    await self._upload_artifact_via_http(
                        info_path, artifact_path="model", run_id=run_id
                    )
                else:
                    await asyncio.to_thread(
                        mlflow.log_artifact,
                        state_path,
                        artifact_path="model",
                        run_id=run_id,
                    )
                    await asyncio.to_thread(
                        mlflow.log_artifact,
                        info_path,
                        artifact_path="model",
                        run_id=run_id,
                    )

            self.logger.info(
                f"Logged program state: {self.program.name} "
                f"({len(self.program.trainable_variables)} trainable variables)"
            )

        except Exception as e:
            self.logger.warning(f"Failed to log program model: {e}")

    def on_train_begin(self, logs=None):
        """Called at the beginning of training."""
        self._in_training = True
        self._setup_mlflow()
        self._start_run(run_name_suffix="train")
        self.logger.debug("MLflow run started for training")

        # Log hyperparameters
        run_maybe_nested(self._log_params())

        # Log program plot
        if self.log_program_plot:
            run_maybe_nested(self._log_program_plot_artifact())

    def on_train_end(self, logs=None):
        """Called at the end of training."""
        run_maybe_nested(self._log_metrics(logs, step=self._step))

        # Log program as model at end of training
        if self.log_program_model:
            run_maybe_nested(self._log_program_model())

        self._end_run()
        self._in_training = False
        self.logger.debug("MLflow run ended for training")

    def on_epoch_begin(self, epoch, logs=None):
        """Called at the start of an epoch."""
        self._epoch = epoch

    def on_epoch_end(self, epoch, logs=None):
        """Called at the end of an epoch."""
        if not self.log_epoch_metrics:
            return

        self._epoch = epoch
        run_maybe_nested(self._log_metrics(logs, step=epoch))
        self.logger.debug(f"Logged metrics for epoch {epoch}")

    def on_train_batch_begin(self, batch, logs=None):
        """Called at the beginning of a training batch."""
        pass

    def on_train_batch_end(self, batch, logs=None):
        """Called at the end of a training batch."""
        if not self.log_batch_metrics:
            return

        self._step += 1
        run_maybe_nested(self._log_metrics(logs, step=self._step))

    def on_test_begin(self, logs=None):
        """Called at the beginning of evaluation or validation."""
        # Only start a new run if we're not already in a training run
        if self._run is None and not self._in_training:
            self._setup_mlflow()
            self._start_run(run_name_suffix="test")
            # Same run type as `mlflow.genai.evaluate()` runs
            mlflow.set_tag("mlflow.runType", "genai_evaluate")
            self.logger.debug("MLflow run started for testing")

    def on_test_end(self, logs=None):
        """Called at the end of evaluation or validation."""
        run_maybe_nested(self._log_metrics(logs, step=self._step))
        # Only end the run if we're not in training (standalone evaluate() call)
        if self._run is not None and not self._in_training:
            self._step += 1
            self._end_run()
            self.logger.debug("MLflow run ended for testing")

    def on_test_batch_begin(self, batch, logs=None):
        """Called at the beginning of a test batch."""
        self._trace_mark = monitor_hook.root_trace_mark()

    def on_test_batch_end(self, batch, logs=None):
        """Called at the end of a test batch."""
        if self.log_assessments:
            run_maybe_nested(self._log_batch_assessments())

        if not self.log_batch_metrics:
            return

        self._step += 1
        run_maybe_nested(self._log_metrics(logs, step=self._step))

    async def _log_batch_assessments(self):
        """Log the per-sample rewards of the batch just evaluated as `reward`
        feedback assessments on the samples' traces.

        Sample i is matched with the i-th root trace started since
        `on_test_batch_begin`; when the counts differ (tracing disabled, or
        the batch's predictions came from the auto-build pass that ran before
        the run started) nothing is logged.
        """
        if self._run is None or self.program is None:
            return
        rewards = getattr(self.program, "_per_sample_rewards", None)
        trace_ids = monitor_hook.root_trace_ids_since(self._trace_mark)
        if not rewards or len(rewards) != len(trace_ids):
            self.logger.debug(
                "Skipping assessments: %s rewards for %s traces",
                None if rewards is None else len(rewards),
                len(trace_ids),
            )
            return

        reward_fn = getattr(self.program, "_compile_reward", None)
        reward_fn = getattr(reward_fn, "_user_reward", reward_fn)
        source_id = getattr(reward_fn, "name", None) or "reward"
        source = mlflow.entities.AssessmentSource(
            source_type=mlflow.entities.AssessmentSourceType.CODE,
            source_id=source_id,
        )
        run_id = self._run.info.run_id

        def log_one(trace_id, value):
            try:
                mlflow.log_feedback(
                    trace_id=trace_id,
                    name="reward",
                    value=float(value),
                    source=source,
                    metadata={"mlflow.assessment.sourceRunId": run_id},
                )
            except Exception as e:
                self.logger.warning(f"Failed to log assessment on {trace_id}: {e}")

        await asyncio.gather(
            *(
                asyncio.to_thread(log_one, trace_id, value)
                for trace_id, value in zip(trace_ids, rewards)
            )
        )

    def on_predict_begin(self, logs=None):
        """Called at the beginning of prediction."""
        pass

    def on_predict_end(self, logs=None):
        """Called at the end of prediction."""
        pass

    def on_predict_batch_begin(self, batch, logs=None):
        """Called at the beginning of a prediction batch."""
        pass

    def on_predict_batch_end(self, batch, logs=None):
        """Called at the end of a prediction batch."""
        pass

    def __del__(self):
        """End our MLflow run if it was left open and is still the active one.

        Guarded by a run-id check: ``mlflow.end_run()`` always ends whatever run
        is *globally* active, so a finalizer firing at GC time must not end an
        unrelated run (this also keeps a leaked finalizer from polluting other
        code's, or another test's, active run).
        """
        run = getattr(self, "_run", None)
        if run is None:
            return
        try:
            active = mlflow.active_run()
            if active is not None and active.info.run_id == run.info.run_id:
                mlflow.end_run()
        except Exception:
            pass