"""
Lineage SDK
Provides methods to create, query, and delete lineage relationships for any object
(models, deployments, trainings, etc.)
"""

from typing import Dict, Any, List, Optional, Union
from requests import RequestException
from adxp_sdk.auth.credentials import BaseCredentials
from adxp_sdk.base import BaseClient


def _build_endpoint(path: str = "") -> str:
    """Build endpoint path for lineage API."""
    base = "api/v1/lineages"
    if not path or not path.strip():
        return base
    return f"{base}/{path.lstrip('/')}"


class LineageClient(BaseClient):
    """Client for managing lineage relationships (create, get, delete)"""

    def __init__(self, credentials: BaseCredentials):
        """
        Initialize LineageClient

        Args:
            credentials: Authentication credentials
        """
        super().__init__(credentials)

    def create_lineage(
        self,
        lineages: List[Dict[str, Any]],
    ) -> Union[Dict[str, Any], List[Any], Any]:
        """
        Create lineage relations via POST /api/v1/lineages

        Args:
            lineages: List of lineage items. Each item must have:
                - source_key: str
                - source_type: str (e.g., KNOWLEDGE, MODEL, VECTOR_DB)
                - target_key: str
                - target_type: str
                - action: str (USE or CREATE)

        Returns:
            API response (201 No Content on success)

        Example:
            client.create_lineage([
                {
                    "source_key": "repo-uuid",
                    "source_type": "KNOWLEDGE",
                    "target_key": "vector-db-uuid",
                    "target_type": "VECTOR_DB",
                    "action": "USE",
                },
            ])
        """
        payload = {"lineages": lineages}
        return self._make_request("POST", _build_endpoint(), json_data=payload)

    def delete_lineage(
        self,
        source_key: str,
        target_key: Optional[str] = None,
    ) -> Union[Dict[str, Any], List[Any], Any]:
        """
        Delete lineage via DELETE /api/v1/lineages/{source_key}

        Args:
            source_key: Source object key (UUID)
            target_key: Optional target key. If provided, deletes specific relation.
                       If None, deletes all relations from source_key.

        Returns:
            API response (204 No Content on success)
        """
        params: Dict[str, Any] = {}
        if target_key is not None:
            params["target_key"] = target_key
        return self._make_request(
            "DELETE",
            _build_endpoint(source_key),
            params=params if params else None,
        )

    def get_lineage(
        self, 
        object_key: str, 
        direction: str, 
        action: str = "USE", 
        max_depth: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Get lineage information for an object via GET /api/v1/lineages/{object_key}/{direction}
        
        Args:
            object_key (str): The object key (UUID or name, e.g., model UUID, deployment name, training ID)
            direction (str): Traversal direction ('downstream' or 'upstream')
            action (str): Type of relationship to query ('USE' or 'CREATE') (default: "USE")
            max_depth (int): Maximum depth to traverse (default: 5, min: 1, max: 10)
            
        Returns:
            list: The API response containing lineage relationships (list of dicts)
                
        Example response:
            [
                {
                    "source_key": "0328b11c-291a-428b-b346-3c5a1c47c689",
                    "target_key": "66387e06-817b-4d4d-8ca1-e6b9383009f1",
                    "action": "USE",
                    "depth": 1,
                    "source_type": "KNOWLEDGE",
                    "target_type": "SERVING_MODEL}"
                }
            ]
        
        Example usage:
            # Check downstream dependencies (what uses this object)
            downstream = client.get_lineage(
                object_key="model-uuid-123",
                direction="downstream",
                action="USE",
                max_depth=5
            )
            
            # Check upstream dependencies (what this object uses)
            upstream = client.get_lineage(
                object_key="model-uuid-123",
                direction="upstream",
                action="USE",
                max_depth=3
            )
        """
        # Validate direction
        if direction not in ["downstream", "upstream"]:
            raise ValueError(f"direction must be 'downstream' or 'upstream', got: {direction}")
        
        # Validate action
        if action not in ["USE", "CREATE"]:
            raise ValueError(f"action must be 'USE' or 'CREATE', got: {action}")
        
        # Validate max_depth
        if not isinstance(max_depth, int) or max_depth < 1 or max_depth > 10:
            raise ValueError(f"max_depth must be an integer between 1 and 10, got: {max_depth}")
        
        params = {
            "action": action,
            "max_depth": max_depth
        }
        
        try:
            response = self._make_request(
                method="GET",
                endpoint=_build_endpoint(f"{object_key}/{direction}"),
                params=params,
            )
            # Lineage API returns list directly (response_model=list[LineageRelationWithTypes])
            # _make_request returns response.json() which could be a list or dict
            # Handle both cases
            if isinstance(response, list):
                return response
            # If response is a dict with 'data' field, extract it
            if isinstance(response, dict) and "data" in response:
                data = response["data"]
                # Ensure data is a list
                if isinstance(data, list):
                    return data
                # If data is not a list, wrap it
                return [data] if data else []
            # If response is a dict but no 'data' field, this shouldn't happen
            # but return empty list as fallback
            return []
        except RequestException as e:
            raise RequestException(f"Failed to get lineage: {str(e)}")

