26Y03d2

"""
[2단계-Polars 최종완결판] 멀티 코어 기반 일별 자원 마감 정산 및 AIStor 전송 파이프라인
(CLUSTER_NODE_PATTERNS 기반 노드 화이트리스트 필터링 기능 복원 완료)
"""
import os
import argparse
import boto3
from botocore.exceptions import ClientError
import polars as pl
from datetime import datetime, timedelta, timezone
from pathlib import Path

import config
from config import (
    RAW_DIR, MERGED_DIR, CLUSTER_NODE_PATTERNS, 
    get_workload_type, classify_node_infrastructure
)

KST = timezone(timedelta(hours=9))
W_CPU, W_MEM, W_PV = 1.0, 0.11, 0.02

def parse_arguments():
    parser = argparse.ArgumentParser(description="FinOps Polars-Powered Reprocessing Engine")
    parser.add_argument("--days", type=int, default=None)
    parser.add_argument("--start-date", type=str, default=None)
    parser.add_argument("--end-date", type=str, default=None)
    parser.add_argument("--cluster", type=str, default="ALL", choices=["COMPUTE", "STORAGE", "ALL"])
    parser.add_argument("--workload-domain", type=str, default=None)
    return parser.parse_args()

def get_minio_client():
    return boto3.client(
        "s3", endpoint_url=os.getenv("MINIO_ENDPOINT"), aws_access_key_id=os.getenv("MINIO_ACCESS_KEY"),
        aws_secret_access_key=os.getenv("MINIO_SECRET_KEY"), region_name="us-east-1",
        config=boto3.session.Config(signature_version="s3v4")
    )

def sync_down_raw_from_minio(s3_client, bucket_name, start_dt, end_dt):
    chunk_delta = timedelta(hours=1)
    current_start = start_dt
    print(f"\n📡 [AIStor 스마트 캐시 엔진] 원천 데이터 가용성 체크 시작...")

    while current_start < end_dt:
        chunk_str = current_start.strftime("%Y%m%d_%H")
        f_name = f"prom_raw_{chunk_str}.parquet"
        object_key = f"raw/{f_name}"
        local_path = RAW_DIR / f_name
        
        if local_path.exists() and local_path.stat().st_size > 0:
            print(f"   ⚡ [Cache Hit] 로컬 캐시 자산 활용 (다운로드 스킵): {f_name}")
        else:
            try:
                s3_client.head_object(Bucket=bucket_name, Key=object_key)
                print(f"   📥 [Cache Miss] AIStor 백업으로부터 다운로드: {object_key}")
                s3_client.download_file(bucket_name, object_key, str(local_path))
            except ClientError as e:
                if e.response['Error']['Code'] != "404":
                    print(f"   ❌ AIStor 통신 실패 에러 ({f_name}): {str(e)}")
            except Exception as e:
                print(f"   ⚠️ 일반 다운로드 예외 발생 ({f_name}): {str(e)}")
                
        current_start += chunk_delta

# ─── 🛡️ [복원] CLUSTER_NODE_PATTERNS 기반 엄격한 노드 필터링 함수 ───
def filter_target_nodes(df: pl.DataFrame, cluster_type: str) -> pl.DataFrame:
    """
    config의 CLUSTER_NODE_PATTERNS 정규식 명세와 매칭되는 소속 노드만 화이트리스트 필터링합니다.
    """
    pattern = CLUSTER_NODE_PATTERNS.get(cluster_type)
    if not pattern:
        print(f"   ⚠️  [경고] {cluster_type}에 대한 CLUSTER_NODE_PATTERNS 명세가 config에 없어 필터를 건너뜁니다.")
        return df
    
    # 패턴 명세가 리스트 형태일 경우 파이프(|)로 묶어 정규식 OR 연산 처리 안정성 확보
    if isinstance(pattern, list):
        regex_pattern = "|".join(pattern)
    else:
        regex_pattern = pattern
        
    initial_len = len(df)
    filtered_df = df.filter(pl.col("node").str.contains(regex_pattern))
    drop_count = initial_len - len(filtered_df)
    
    if drop_count > 0:
        print(f"   🧹 [노드 필터 가동] {cluster_type} 규격 외 타 클러스터 오염 노드 데이터 {drop_count:,}행 제거 완료.")
        
    return filtered_df

def main():
    args = parse_arguments()
    s3_client = get_minio_client()
    bucket_name = os.getenv("MINIO_RAW_BUCKET", "devops-test")
    now_kst = datetime.now(KST)
    
    if args.start_date and args.end_date:
        start_dt = datetime.strptime(args.start_date, "%Y-%m-%d").replace(tzinfo=KST)
        end_dt = (datetime.strptime(args.end_date, "%Y-%m-%d") + timedelta(days=1)).replace(tzinfo=KST)
    elif args.days:
        start_dt = (now_kst - timedelta(days=args.days)).replace(hour=0, minute=0, second=0, microsecond=0)
        end_dt = now_kst
    else:
        yesterday = (now_kst - timedelta(days=1)).date()
        start_dt = datetime.combine(yesterday, datetime.min.time()).replace(tzinfo=KST)
        end_dt = datetime.combine(now_kst.date(), datetime.min.time()).replace(tzinfo=KST)

    sync_down_raw_from_minio(s3_client, bucket_name, start_dt, end_dt)

    raw_files = list(RAW_DIR.glob("prom_raw_*.parquet"))
    if not raw_files:
        print("\n🚨 [중단] 정산 가동할 물리 원천 파일이 존재하지 않습니다.")
        return

    print(f"\n📊 [Polars Lazy Engine 기동] {len(raw_files)}개 소스 파티션 스캔 링킹...")
    lf_raw = pl.scan_parquet(str(RAW_DIR / "prom_raw_*.parquet"))

    start_naive = start_dt.replace(tzinfo=None)
    end_naive = end_dt.replace(tzinfo=None)
    
    lf_filtered = lf_raw.filter(
        (pl.col("timestamp") >= start_naive) & (pl.col("timestamp") < end_naive)
    ).with_columns(
        pl.col("timestamp").dt.strftime("%Y%m%d").alias("date")
    )

    print("⏳ 대용량 인프라 정산 메트릭 다차원 분산 집계 중 (Rust Core 스레드 가동)...")
    df_raw = lf_filtered.collect()
    
    if df_raw.is_empty():
        print("⚠️  해당 기간 내 정산할 데이터 로우가 없습니다.")
        return

    df_raw = df_raw.with_columns([
        pl.col("node").map_elements(
            lambda x: {"cluster_type": classify_node_infrastructure(x)[0], "workload_domain": classify_node_infrastructure(x)[1]},
            return_dtype=pl.Struct({"cluster_type": pl.String, "workload_domain": pl.String})
        ).alias("infra_struct"),
        pl.col("pod").map_elements(get_workload_type, return_dtype=pl.String).alias("workload_type")
    ]).with_columns([
        pl.col("infra_struct").struct.field("cluster_type").alias("cluster_type"),
        pl.col("infra_struct").struct.field("workload_domain").alias("workload_domain")
    ]).drop("infra_struct")

    target_dates = df_raw["date"].unique().to_list()
    target_clusters = ["COMPUTE", "STORAGE"] if args.cluster.upper() == "ALL" else [args.cluster.upper()]

    print(f"🎯 분할 연산 타겟 스케줄러 가동 -> 일자 풀: {target_dates} | 클러스터 풀: {target_clusters}")

    GB_DIV = 1024 ** 3

    for date_chunk in target_dates:
        for cluster_chunk in target_clusters:
            
            df_slice = df_raw.filter((pl.col("date") == date_chunk) & (pl.col("cluster_type") == cluster_chunk))
            
            # 💡 [복원 로직 주입] 2차 마진 가드레일: 정규식 패턴 기반 대상 노드만 2차 필터링 수행
            df_slice = filter_target_nodes(df_slice, cluster_chunk)
            
            if df_slice.is_empty():
                continue

            if args.workload_domain:
                df_slice = df_slice.filter(pl.col("workload_domain") == args.workload_domain)
                if df_slice.is_empty(): continue

            print(f"   ⏳ [Polars 집계 연산] {date_chunk} ➡️ {cluster_chunk} 고성능 그룹바이 컴파일...")

            target_fields = ["cpu_request", "cpu_limit", "cpu_usage", "cpu_throttled", "mem_request", "mem_limit", "mem_usage", "mem_rss", "oom_event", "pv_capacity", "pv_used"]
            for col in target_fields:
                if col not in df_slice.columns:
                    df_slice = df_slice.with_columns(pl.lit(0.0).alias(col))

            # [테이블 1] 상세 메트릭 원부 집계
            df_pod = df_slice.group_by([
                "date", "cluster_type", "workload_domain", "namespace", "workload_type", "node", "pod", "container"
            ]).agg([
                pl.len().alias("minutes_running"),
                pl.col("cpu_request").max().alias("cpu_request_max"),
                pl.col("cpu_limit").max().alias("cpu_limit_max"),
                pl.col("cpu_usage").quantile(0.95).alias("cpu_usage_p95"),
                pl.col("cpu_throttled").max().alias("cpu_throttled_max"),
                (pl.col("mem_request").max() / GB_DIV).alias("mem_request_max"),
                (pl.col("mem_limit").max() / GB_DIV).alias("mem_limit_max"),
                (pl.col("mem_usage").quantile(0.95) / GB_DIV).alias("mem_usage_p95"),
                (pl.col("mem_rss").quantile(0.95) / GB_DIV).alias("mem_rss_p95"),
                pl.col("oom_event").sum().alias("oom_strike_sum"),
                (pl.col("pv_capacity").max() / GB_DIV).alias("pv_capacity_max"),
                (pl.col("pv_used").quantile(0.95) / GB_DIV).alias("pv_usage_p95")  # ◀ 네이밍 완치 완료
            ]).fill_null(0.0)

            df_pod = df_pod.with_columns([
                (pl.col("cpu_request_max") * (pl.col("minutes_running") / 60.0)).alias("cpu_allocated_core_hours"),
                (pl.col("cpu_usage_p95") * (pl.col("minutes_running") / 60.0)).alias("cpu_usage_core_hours"),
                (pl.col("mem_request_max") * (pl.col("minutes_running") / 60.0)).alias("mem_allocated_gb_hours"),
                (pl.col("mem_usage_p95") * (pl.col("minutes_running") / 60.0)).alias("mem_usage_gb_hours"),
                (pl.col("pv_capacity_max") * (pl.col("minutes_running") / 60.0)).alias("pv_allocated_gb_hours"),
                (pl.col("pv_usage_p95") * (pl.col("minutes_running") / 60.0)).alias("pv_usage_gb_hours"),
                (pl.col("oom_strike_sum") > 0).alias("is_oom_killed"),
                (pl.col("cpu_request_max") == 0).alias("has_no_request"),
                (pl.col("cpu_limit_max") == 0).alias("has_no_limit"),
                ((pl.col("cpu_usage_p95") - pl.col("cpu_request_max")).clip(lower_bound=0)).alias("cpu_shortage_cores")
            ]).with_columns([
                ((pl.col("cpu_allocated_core_hours") - pl.col("cpu_usage_core_hours")).clip(lower_bound=0)).alias("cpu_waste_core_hours"),
                ((pl.col("mem_allocated_gb_hours") - pl.col("mem_usage_gb_hours")).clip(lower_bound=0)).alias("mem_waste_gb_hours"),
                ((pl.col("pv_allocated_gb_hours") - pl.col("pv_usage_gb_hours")).clip(lower_bound=0)).alias("pv_waste_gb_hours"),
            ])

            df_pod = df_pod.with_columns(
                pl.when(pl.col("is_oom_killed")).then(pl.lit("💥 OOM장애발생"))
                .when((pl.col("cpu_shortage_cores") > 0.5) | (pl.col("cpu_throttled_max") > 0.2)).then(pl.lit("⚠️ Request부족"))
                .when((pl.col("cpu_waste_core_hours") > 10) | (pl.col("pv_waste_gb_hours") > 50)).then(pl.lit("📉 과다할당"))
                .otherwise(pl.lit("✅ 최적화완료")).alias("status")
            )

            # 📅 [테이블 2] Namespace 사용량 요약 빌드
            df_daily_ns = df_pod.group_by(["date", "namespace"]).agg([
                pl.col("cpu_usage_core_hours").sum().alias("cpu_used_ch"),
                pl.col("cpu_allocated_core_hours").sum().alias("cpu_alloc_ch"),
                pl.col("cpu_waste_core_hours").sum().alias("cpu_waste_ch"),
                pl.col("mem_usage_gb_hours").sum().alias("mem_used_gh"),
                pl.col("mem_allocated_gb_hours").sum().alias("mem_alloc_gh"),
                pl.col("pv_usage_gb_hours").sum().alias("pv_used_gh"),
                pl.col("pv_allocated_gb_hours").sum().alias("pv_alloc_gh")
            ]).with_columns(
                ((pl.col("cpu_alloc_ch") * W_CPU) + (pl.col("mem_alloc_gh") * W_MEM) + (pl.col("pv_alloc_gh") * W_PV)).round(1).alias("final_usage_score")
            )

            # 📉 [테이블 3] 파레토 원부 롤업 빌드 
            df_ns = df_pod.group_by("namespace").agg([
                pl.col("cpu_waste_core_hours").sum().alias("total_waste_core_hours"),
                pl.col("minutes_running").sum().alias("minutes_running_sum"),
                pl.col("container").count().alias("container_cnt"),
                pl.col("cpu_allocated_core_hours").sum().alias("total_allocated_core_hours")
            ]).sort("total_waste_core_hours", descending=True)
            
            global_total_waste = df_ns["total_waste_core_hours"].sum() or 0.1
            df_ns = df_ns.with_columns([
                pl.lit(date_chunk).alias("date"),
                ((pl.col("total_waste_core_hours") / global_total_waste) * 100).round(2).alias("waste_share_pct")
            ]).with_columns(
                pl.col("waste_share_pct").cum_sum().round(2).alias("waste_cumsum_pct")
            )

            # ─── 💾 [로컬 저장 명세 고정 바인딩] ───
            cluster_partition_dir = MERGED_DIR / cluster_chunk
            cluster_partition_dir.mkdir(parents=True, exist_ok=True)
            
            target_output_file = cluster_partition_dir / f"daily_enriched_{cluster_chunk}_{date_chunk}.parquet"
            ns_output_file     = cluster_partition_dir / f"daily_ns_usage_{cluster_chunk}_{date_chunk}.parquet"
            pareto_output_file = cluster_partition_dir / f"pareto_ns_{cluster_chunk}_{date_chunk}.parquet"
            
            df_pod.write_parquet(target_output_file)
            df_daily_ns.write_parquet(ns_output_file)
            df_ns.write_parquet(pareto_output_file)
            
            print(f"     💾 [로컬 저장] {cluster_chunk} 파티션 3대 정산 자산화 성료.")

            # ─── 📡 [AIStor 영구 적재 가드레일 주입] ───
            print(f"     Submitting artifacts to AIStor Tables storage system...")
            try:
                s3_client.upload_file(str(target_output_file), bucket_name, f"merged/{cluster_chunk}/{target_output_file.name}")
                s3_client.upload_file(str(ns_output_file), bucket_name, f"merged/{cluster_chunk}/{ns_output_file.name}")
                s3_client.upload_file(str(pareto_output_file), bucket_name, f"merged/{cluster_chunk}/{pareto_output_file.name}")
                print(f"     ✅ [AIStor 업로드 완수] 테이블 자산 적재 완료 -> {pareto_output_file.name}")
            except Exception as e:
                print(f"     ❌ [AIStor 통신장애] 오브젝트 전송 실패: {str(e)}")

    print("\n🏁 === [Step2 Polars 배치 완료] 로컬 및 AIStor 테이블 연동이 완벽히 마감되었습니다. ===")

if __name__ == "__main__":
    main()