"""
MLFlow API
Integration with MLFlow tracking server through backend API
"""

from typing import Optional, Dict, Any, List
from .base import BaseAPI


class MLFlowAPI(BaseAPI):
    """MLFlow tracking endpoints"""
    
    def get_ui_url(self) -> Dict[str, Any]:
        """
        Get MLFlow UI URL.
        
        Returns:
            Dict with ui_url and tracking_uri
        """
        return self.client._request("GET", "mlflow/ui-url")
    
    def get_run(self, run_id: str) -> Dict[str, Any]:
        """
        Get MLFlow run details.
        
        Args:
            run_id: MLFlow run ID
            
        Returns:
            Run data with run_id, experiment_id, status, start_time, end_time, tags, params, metrics
        """
        return self.client._request("GET", f"mlflow/runs/{run_id}")
    
    def get_run_by_kws_id(self, kws_run_id: int) -> Dict[str, Any]:
        """
        Get MLFlow run by KWS experiment run ID.
        
        Args:
            kws_run_id: KWS experiment run ID
            
        Returns:
            Run data with run_id, experiment_id, status, start_time, end_time, tags, params, metrics
        """
        return self.client._request("GET", f"mlflow/runs/kws/{kws_run_id}")
    
    def list_runs_by_experiment_name(
        self,
        experiment_name: str,
        filter_string: Optional[str] = None,
        max_results: int = 100,
    ) -> List[Dict[str, Any]]:
        """
        List MLFlow runs for an experiment by name.
        
        Args:
            experiment_name: MLFlow experiment name
            filter_string: Optional MLFlow filter string
            max_results: Maximum number of results (1-1000)
            
        Returns:
            List of run data
        """
        params = {"max_results": min(max_results, 1000)}
        if filter_string:
            params["filter_string"] = filter_string
        
        return self.client._request(
            "GET",
            f"mlflow/experiments/{experiment_name}/runs",
            params=params
        )
    
    def list_runs_by_experiment_id(
        self,
        experiment_id: int,
        filter_string: Optional[str] = None,
        max_results: int = 100,
    ) -> List[Dict[str, Any]]:
        """
        List MLFlow runs for a KWS experiment.
        
        Args:
            experiment_id: KWS experiment ID
            filter_string: Optional MLFlow filter string
            max_results: Maximum number of results (1-1000)
            
        Returns:
            List of run data
        """
        params = {"max_results": min(max_results, 1000)}
        if filter_string:
            params["filter_string"] = filter_string
        
        return self.client._request(
            "GET",
            f"mlflow/experiments/{experiment_id}/runs",
            params=params
        )
    
    def get_artifacts(
        self,
        run_id: str,
        artifact_path: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Get MLFlow run artifacts.
        
        Args:
            run_id: MLFlow run ID
            artifact_path: Optional artifact path
            
        Returns:
            Dict with artifact_uri and artifacts list
        """
        params = {}
        if artifact_path:
            params["artifact_path"] = artifact_path
        
        return self.client._request(
            "GET",
            f"mlflow/runs/{run_id}/artifacts",
            params=params
        )
    
    def log_metrics(
        self,
        run_id: str,
        metrics: Dict[str, float],
        step: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        Log metrics to MLFlow run.
        
        Args:
            run_id: MLFlow run ID
            metrics: Dictionary of metric name -> value (must be numeric)
            step: Optional step number for time-series metrics
            
        Returns:
            Success message
            
        Example:
            >>> client.mlflow.log_metrics(
            ...     run_id="abc123",
            ...     metrics={"accuracy": 0.95, "f1_score": 0.92},
            ...     step=1
            ... )
        """
        data = {"metrics": metrics}
        if step is not None:
            data["step"] = step
        
        return self.client._request(
            "POST",
            f"mlflow/runs/{run_id}/metrics",
            json=data
        )
    
    def log_params(
        self,
        run_id: str,
        params: Dict[str, str],
    ) -> Dict[str, Any]:
        """
        Log parameters to MLFlow run.
        
        Args:
            run_id: MLFlow run ID
            params: Dictionary of parameter name -> value (must be strings)
            
        Returns:
            Success message
            
        Example:
            >>> client.mlflow.log_params(
            ...     run_id="abc123",
            ...     params={"learning_rate": "0.001", "batch_size": "32"}
            ... )
        """
        return self.client._request(
            "POST",
            f"mlflow/runs/{run_id}/params",
            json={"params": params}
        )
    
    def update_run_status(
        self,
        run_id: str,
        status: str,
    ) -> Dict[str, Any]:
        """
        Update MLFlow run status.
        
        Args:
            run_id: MLFlow run ID
            status: Run status ("RUNNING", "FINISHED", "FAILED", "KILLED")
            
        Returns:
            Updated run data
        """
        return self.client._request(
            "PUT",
            f"mlflow/runs/{run_id}/status",
            json={"status": status}
        )
