from datetime import datetime
from types import SimpleNamespace
from typing import Any, Dict, Optional

import mlflow

from mmp.factory import Factory
from mmp.settings.factory import SettingsFactory
from mmp.utils.core.console import Console
from mmp.utils.core.logger import log_data, log_infer, log_mlflow, log_pipe
from mmp.utils.core.reproducibility import set_global_seeds
from mmp.utils.data.data_io import format_predictions, load_inference_data, save_output
from mmp.utils.integrations import mlflow_integration as mlflow_utils
from mmp.utils.integrations.version_checker import log_version_warnings


def run_inference_pipeline(
    run_id: str,
    recipe_path: Optional[str] = None,
    config_path: Optional[str] = None,
    data_path: Optional[str] = None,
    context_params: Optional[Dict[str, Any]] = None,
    output_path: Optional[str] = None,
):
    """
    배치 추론 파이프라인을 실행합니다.

    Args:
        run_id: 학습된 모델의 MLflow Run ID
        recipe_path: Override할 Recipe 파일 경로 (None이면 artifact에서 복원)
        config_path: Override할 Config 파일 경로 (None이면 artifact에서 복원)
        data_path: 추론할 데이터 경로 (None이면 artifact의 SQL 사용)
        context_params: SQL 렌더링에 사용할 파라미터
        output_path: 결과 저장 전체 경로 (Config override, 파일명+확장자 포함)
    """
    console = Console()
    context_params = context_params or {}

    # 설정 로드: SettingsFactory를 통해 Artifact 복원 + Override 처리
    log_pipe("설정 로드 중")
    settings = SettingsFactory.for_inference(
        run_id=run_id,
        config_path=config_path,
        recipe_path=recipe_path,
        data_path=data_path,
        context_params=context_params,
    )

    # 재현성을 위한 전역 시드 설정
    seed = getattr(settings.recipe.model, "computed", {}).get("seed", 42)
    set_global_seeds(seed)

    # 파이프라인 시작 로깅
    log_pipe("========== 배치 추론 파이프라인 시작 ==========")
    log_pipe(f"Model Run ID: {run_id} | 환경: {settings.config.environment.name}")

    # MLflow 실행 컨텍스트 시작
    run_name = f"batch_inference_{run_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    with mlflow_utils.start_run(settings, run_name=run_name) as run:
        inference_run_id = run.info.run_id

        # Factory 생성
        factory = Factory(settings)

        # 1. 모델 로드
        log_infer("모델 로드 시작")
        model_uri = f"runs:/{run_id}/model"
        with console.progress_tracker(
            "model_loading", 100, f"Loading model from {model_uri}"
        ) as update:
            model = mlflow.pyfunc.load_model(model_uri)
            update(100)

        log_infer(f"모델 로드 완료 - {model_uri}")

        # 버전 호환성 체크 (모델 로드 성공 후 경고 표시)
        log_version_warnings(run_id)

        # 2. 데이터 준비
        log_data("데이터 로드 시작")
        data_adapter = factory.create_data_adapter()
        df = load_inference_data(
            data_adapter=data_adapter,
            data_path=data_path,
            model=model,
            run_id=run_id,
            context_params=context_params,
            console=console,
        )

        log_data(f"로드 완료 - {len(df):,}행, {len(df.columns)}열")
        mlflow.log_metric("inference_input_rows", len(df))
        mlflow.log_metric("inference_input_columns", len(df.columns))

        # 3. 예측 실행
        log_infer("모델 추론 시작")
        with console.progress_tracker("inference", 100, "Running model prediction") as update:
            # run_mode="batch"로 Offline Store 사용
            predictions_result = model.predict(
                df, params={"run_mode": "batch", "return_dataframe": True}
            )

            # PyfuncWrapper에서 data_interface_schema 정보 가져오기
            wrapped_model = model.unwrap_python_model()
            data_interface_schema = getattr(wrapped_model, "data_interface_schema", {}) or {}

            # data_interface_schema를 사용하여 format
            predictions_df = format_predictions(predictions_result, df, data_interface_schema)
            update(100)

        log_infer(f"추론 완료 - {len(predictions_df):,}개 예측 생성")
        mlflow.log_metric("inference_output_rows", len(predictions_df))

        # 4. 결과 저장
        log_mlflow("결과 저장 시작")
        save_output(
            df=predictions_df,
            settings=settings,
            output_type="inference",
            factory=factory,
            run_id=inference_run_id,
            console=console,
            additional_metadata={"model_run_id": run_id},
            output_path_override=output_path,
        )

        log_pipe("========== 배치 추론 파이프라인 완료 ==========")
        return SimpleNamespace(
            run_id=inference_run_id, model_uri=model_uri, prediction_count=len(predictions_df)
        )
