"""
post_lineage_event.py - AWS DataZone에 Tableau 대시보드와 데이터 소스 간의 계보 정보 등록

이 스크립트는 다음 기능을 수행합니다:
1. Tableau 대시보드 자산 정보 조회
2. 대시보드가 사용하는 데이터 소스 정보 수집
3. OpenLineage 형식의 계보 이벤트 생성
4. DataZone에 계보 이벤트 등록
    
계보 정보는 데이터의 흐름과 의존성을 시각화하여 데이터 자산 간의 관계를 보여줍니다.

사용법:
    python post_lineage_event.py --asset-id <asset_id>
    
    필수 인자:
    --asset-id : DataZone 자산 ID
"""     
        
import boto3
import json
import uuid
import argparse
import sys
from datetime import datetime
from botocore.exceptions import ClientError

# 설정 상수 정의
AWS_REGION = '{AWS Region 입력}'
DATAZONE_DOMAIN_ID = '{DataZone 도메인 ID 입력}' 
DATAZONE_PROJECT_ID = '{DataZone 프로젝트 ID 입력}' 
DATASOURCE_PROJECT_ID = '{DataZone의 DataSource가 저장된 프로젝트 ID 입력}'

# AWS 클라이언트 초기화 
datazone_client = boto3.client('datazone', region_name=AWS_REGION)

def parse_arguments():
    """
    명령줄 인자를 파싱합니다.
    
    Returns:
        argparse.Namespace: 파싱된 명령줄 인자
    """
    parser = argparse.ArgumentParser(
        description='AWS DataZone에 Tableau 대시보드와 데이터 소스 간의 계보 정보를 등록합니다.'
    )

    parser.add_argument(
        '--asset-id',
        type=str,
        help='계보 정보를 등록할 DataZone 자산 ID (필수)'
    )

    # 인자가 없으면 도움말 출력
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    return parser.parse_args()

def get_asset(asset_id):
    """
    DataZone에서 자산 정보를 조회합니다.
    
    Args:
        asset_id (str): 조회할 자산 ID
        
    Returns:
        dict: 자산 정보 딕셔너리
    """
    response = datazone_client.get_asset(
        domainIdentifier=DATAZONE_DOMAIN_ID,
        identifier=asset_id
    )
    return response

def search_asset_in_project(asset_name, project_id):
    """
    특정 프로젝트 내에서 자산을 검색합니다.
    
    Args:
        asset_name (str): 검색할 자산 이름
        
    Returns:
        list: 검색된 자산 목록, 없거나 오류 발생 시 None
    """
    try:
        response = datazone_client.search(
            domainIdentifier=DATAZONE_DOMAIN_ID,
            owningProjectIdentifier=project_id,
            searchScope="ASSET",
            searchText=asset_name
        )   
        return response['items']
    except ClientError as e:
        if e.response['Error']['Code'] == 'AccessDeniedException':
            print(f"접근 권한이 없어 '{asset_name}' 자산을 검색할 수 없습니다.")
            return None
        else:
            raise e

def get_asset_id_by_source_identifier(source_identifier, project_id):
    """
    소스 식별자로 자산 ID를 찾습니다.
    
    Args:
        source_identifier (str): 검색할 소스 식별자
        
    Returns:
        str: 자산 ID, 없으면 None
    """
    # 소스 식별자에서 이름 부분 추출
    try:
        name = source_identifier.split('/', 1)[1]
    except IndexError:
        print(f"잘못된 소스 식별자 형식: {source_identifier}")
        return None

    # 이름으로 자산 검색
    assets = search_asset_in_project(name, project_id)
    if not assets:
        print(f"'{name}' 이름의 자산을 찾을 수 없습니다.")
        return None
        
    # 검색된 자산 중 일치하는 소스 식별자를 가진 자산 찾기
    for asset in assets:
        asset_id = asset['assetItem']['identifier']
        asset_info = get_asset(asset_id)

        # AssetCommonDetailsForm에서 소스 식별자 확인
        filtered_forms = [form for form in asset_info['formsOutput']
                         if form['formName'] == 'AssetCommonDetailsForm']

        if not filtered_forms:
            continue

        content = json.loads(filtered_forms[0]['content'])
        if content.get('sourceIdentifier') == source_identifier:
            return asset_id

    print(f"소스 식별자 '{source_identifier}'와 일치하는 자산을 찾을 수 없습니다.")
    return None

def get_asset_columns(asset_id):
    """
    자산의 컬럼 정보를 가져옵니다.
    
    Args:
        asset_id (str): 자산 ID
        
    Returns:
        list: 컬럼 정보 목록, 없으면 빈 리스트
    """
    # 자산 정보 조회
    response = get_asset(asset_id)

     # Glue 테이블/뷰 양식 찾기
    glue_form_index = next((index for (index, d) in enumerate(response['formsOutput'])
                        if d['formName'] in ['GlueViewForm', 'GlueTableForm']), None)

    if glue_form_index is None:
        print(f"자산 ID {asset_id}에 Glue 테이블/뷰 양식이 없습니다.")
        return []

    # 컬럼 정보 추출
    glue_data = response['formsOutput'][glue_form_index]['content']
    data_dict = json.loads(glue_data)

    if 'columns' not in data_dict:
        print(f"자산 ID {asset_id}에 컬럼 정보가 없습니다.")
        return []

    return data_dict['columns']

def get_source_identifier(form_output):
    """
    자산의 양식 출력에서 소스 식별자를 찾습니다.
    
    Args:
        form_output (list): 자산의 양식 출력 목록
        
    Returns:
        str: 소스 식별자, 없으면 None
    """
    for item in form_output:
        try:
            content = json.loads(item['content'])
            if 'sourceIdentifier' in content:
                return content['sourceIdentifier']
        except json.JSONDecodeError:
            continue
            
    print("자산에 소스 식별자가 없습니다.")
    return None

def create_lineage_event(asset_id, datasources):
    """
    대시보드와 데이터 소스 간의 데이터 계보를 생성하는 OpenLineage 이벤트를 생성합니다.
    
    Args:
        asset_id (str): 대시보드 자산 ID
        datasources (list): 데이터 소스 정보 목록
        
    Returns:
        dict: 계보 이벤트 등록 응답
    """
    print(f"자산 ID {asset_id}의 계보 이벤트 생성 시작...")

    # 자산 정보 로드
    try:
        asset_info = get_asset(asset_id)
    except ClientError as e:
        print(f"자산 정보 조회 중 오류 발생: {str(e)}")
        return None

    # 현재 시간 ISO 형식으로 생성
    event_time = datetime.utcnow().isoformat() + "Z"

    # 이벤트 UUID 생성
    run_id = str(uuid.uuid4())

    # 대시보드 자산 정보를 담은 출력 facet 정의
    output_facets = {
        "dataZone": {
            "_producer": "custom",
            "_schemaURL": "https://raw.githubusercontent.com/OpenLineage/OpenLineage/main/spec/OpenLineage.json#/definitions/DataZoneFacet",
            "assetId": asset_info['id'],
            "assetName": asset_info['name'],
            "assetDescription": asset_info.get('description', '')
        }
    }

    # 출력 데이터셋 정의 (대시보드의 소스 식별자 구성)
    source_identifier = get_source_identifier(asset_info['formsOutput'])
    if not source_identifier:
        print(f"자산 ID {asset_id}에 소스 식별자가 없어 계보 이벤트를 생성할 수 없습니다.")
        return None

    # 소스 식별자 파싱
    try:
        parts = source_identifier.split('/', 1)
        namespace = parts[0]
        name = parts[1]
    except IndexError:
        print(f"잘못된 소스 식별자 형식: {source_identifier}")
        return None

    # 출력 데이터셋 정의
    outputs = [{
        "namespace": namespace,
        "name": name,
        "facets": output_facets
    }]

    # 입력 데이터셋 정의 (데이터 소스)
    inputs = []
    for source in datasources:
        # 기본 입력 데이터셋 정보
        input_dataset = {
            "namespace": source["namespace"],
            "name": source["name"]
        }

        # 소스 식별자 생성
        source_identifier = f'{source["namespace"]}/{source["name"]}'
        
        # 자산 ID 조회
        print(f"소스 식별자 '{source_identifier}'에 해당하는 자산 검색 중...")
        source_asset_id = get_asset_id_by_source_identifier(source_identifier, DATASOURCE_PROJECT_ID)

        # 자산 ID가 있는 경우 스키마 정보 추가
        if source_asset_id:
            print(f"자산 ID {source_asset_id}의 컬럼 정보 조회 중...")
            asset_columns = get_asset_columns(source_asset_id)

            if asset_columns:
                # 컬럼 정보 변환 (OpenLineage 스키마에 맞게)
                formatted_columns = []
                for column in asset_columns:
                    formatted_column = {
                        'type': column.get('dataType', 'unknown'),
                        'name': column.get('columnName', 'unknown')
                    }
                    # 설명 정보가 있으면 추가
                    if 'description' in column and column['description']:
                        formatted_column['description'] = column['description']

                    formatted_columns.append(formatted_column)
                
                # 스키마 정보 추가
                input_dataset['facets'] = {
                    "schema": {
                        "_producer": "custom/glue-integration",
                        "_schemaURL": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/SchemaFacet",
                        'fields': formatted_columns
                    },
                    "sourceDetails": {
                        "_producer": "custom/glue-integration",
                        "_schemaURL": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/BaseFacet",
                    }
                }

        inputs.append(input_dataset)

    # OpenLineage 이벤트 구성
    lineage_event = {
        "eventType": "COMPLETE",
        "producer": "custom/tableau-integration",
        "schemaURL": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunEvent",
        "eventTime": event_time,
        "run": {
            "runId": run_id
        },
        "job": {
            "namespace": "tableau",
            "name": f"tableau_dashboard_job_{asset_id}"
        },
        "inputs": inputs,
        "outputs": outputs
    }
    
    # print("생성된 OpenLineage 이벤트:")
    # print(json.dumps(lineage_event, indent=2, ensure_ascii=False))

    # 계보 이벤트 제출
    print(f"DataZone에 계보 이벤트 등록 중...")
    try:
        response = datazone_client.post_lineage_event(
            domainIdentifier=DATAZONE_DOMAIN_ID,
            event=json.dumps(lineage_event, ensure_ascii=False)
        )
        print(f"자산 ID {asset_id}에 대한 계보 이벤트가 성공적으로 등록되었습니다. (Run ID: {run_id})")
        return response
    except Exception as e:
        print(f"계보 이벤트 등록 중 오류 발생: {str(e)}")
        return None

def main():
    """
    메인 실행 함수 - 명령줄 인자를 파싱하고 자산의 계보 이벤트를 생성합니다.
    """
    try:
        # 명령줄 인자 파싱
        args = parse_arguments()
        
        # 필수 인자 확인
        if not args.asset_id:
            print("오류: --asset-id 인자가 필요합니다.")
            sys.exit(1)
            
        asset_id = args.asset_id
        print(f"Tableau 대시보드 계보 이벤트 생성 시작 (자산 ID: {asset_id})...")
        
        # 데이터 소스 정보 (예시)
        datasources = [
            {
                'namespace': 'glue',
                'name': 'biz.sales_analysis'
            },
            {
                'namespace': 'glue',
                'name': 'biz.sales_target'
            },
            {
                'namespace': 'dw',
                'name': 'product_info'
            }
        ]
            
        # 계보 이벤트 생성 및 등록
        response = create_lineage_event(asset_id, datasources)
        
        if response:
            print("계보 이벤트 등록 완료")
            # 추후 확장을 위해 응답 반환
            return response
        else:
            print("계보 이벤트 등록 실패")
            return None

    except Exception as e:
        print(f"계보 이벤트 생성 중 오류 발생: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

if __name__ == "__main__":
    main()