from metadata.generated.schema.entity.data.table import Table
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
    OpenMetadataConnection,
)
from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
    OpenMetadataJWTClientConfig,
)

server_config = OpenMetadataConnection(
    hostPort="https://your-company.getcollate.io/api",
    authProvider="openmetadata",
    securityConfig=OpenMetadataJWTClientConfig(
        jwtToken="<YOUR-JWT-TOKEN>"
    ),
)
metadata = OpenMetadata(server_config)

# Get lineage by fully qualified name
lineage = metadata.get_lineage_by_name(
    entity=Table,
    fqn="sample_data.ecommerce_db.shopify.dim_customer",
    up_depth=1,
    down_depth=1,
)

print(f"Entity: {lineage['entity']['fullyQualifiedName']}")
print(f"Upstream edges: {len(lineage.get('upstreamEdges', []))}")
print(f"Downstream edges: {len(lineage.get('downstreamEdges', []))}")

# Get lineage by ID
lineage = metadata.get_lineage_by_id(
    entity=Table,
    entity_id="455e3d9d-dbbf-455e-b3be-7191daa825f3",
    up_depth=3,
    down_depth=2,
)

import org.openmetadata.sdk.api.Lineage;

// Get lineage by entity type and ID
Lineage.LineageGraph graph = Lineage.of("table", tableId)
    .upstream(3)
    .downstream(2)
    .fetch();

String raw = graph.getRaw();

// Get lineage with default depth (1 upstream, 1 downstream)
Lineage.LineageGraph graph = Lineage.of("table", tableId).fetch();

# Get lineage by entity type and ID
curl "{base_url}/api/v1/lineage/table/455e3d9d-dbbf-455e-b3be-7191daa825f3" \
  -H "Authorization: Bearer {access_token}"

# Get lineage with custom depth
curl "{base_url}/api/v1/lineage/table/455e3d9d-dbbf-455e-b3be-7191daa825f3?upstreamDepth=3&downstreamDepth=2" \
  -H "Authorization: Bearer {access_token}"

# Get lineage by fully qualified name
curl "{base_url}/api/v1/lineage/table/name/sample_data.ecommerce_db.shopify.dim_customer" \
  -H "Authorization: Bearer {access_token}"

{
  "entity": {
    "id": "455e3d9d-dbbf-455e-b3be-7191daa825f3",
    "type": "table",
    "name": "dim_customer",
    "fullyQualifiedName": "sample_data.ecommerce_db.shopify.dim_customer",
    "deleted": false,
    "href": "http://localhost:8585/api/v1/tables/455e3d9d-dbbf-455e-b3be-7191daa825f3"
  },
  "nodes": [
    {
      "id": "800caa0f-a149-48d2-a0ce-6ca84501767e",
      "type": "table",
      "name": "raw_customer",
      "fullyQualifiedName": "sample_data.ecommerce_db.shopify.raw_customer",
      "deleted": false,
      "href": "http://localhost:8585/api/v1/tables/800caa0f-a149-48d2-a0ce-6ca84501767e"
    },
    {
      "id": "c3d4e5f6-a1b2-7890-abcd-ef1234567890",
      "type": "dashboard",
      "name": "customer_dashboard",
      "fullyQualifiedName": "sample_superset.customer_dashboard",
      "deleted": false,
      "href": "http://localhost:8585/api/v1/dashboards/c3d4e5f6-a1b2-7890-abcd-ef1234567890"
    }
  ],
  "upstreamEdges": [
    {
      "fromEntity": "800caa0f-a149-48d2-a0ce-6ca84501767e",
      "toEntity": "455e3d9d-dbbf-455e-b3be-7191daa825f3",
      "lineageDetails": {
        "sqlQuery": "INSERT INTO dim_customer SELECT id, name FROM raw_customer",
        "columnsLineage": [
          {
            "fromColumns": [
              "sample_data.ecommerce_db.shopify.raw_customer.id"
            ],
            "toColumn": "sample_data.ecommerce_db.shopify.dim_customer.customer_id"
          }
        ]
      }
    }
  ],
  "downstreamEdges": [
    {
      "fromEntity": "455e3d9d-dbbf-455e-b3be-7191daa825f3",
      "toEntity": "c3d4e5f6-a1b2-7890-abcd-ef1234567890"
    }
  ]
}