#!/usr/bin/env python3
"""
모델 배포 API 생성 스크립트

FastAPI를 사용하여 모델을 REST API로 배포합니다.

설치:
    cd plugins/model-deployment/skills/deployment
    uv pip install -r requirements.txt

사용법:
    python deploy_api.py --model-path "./models/model.pkl" --feature-names "V1,V2,V3,Amount"
    python deploy_api.py --model-path "./models/model.pkl" --sample-data "./data/train.csv" --target-column "Class"

실행:
    uvicorn app:app --reload --host 0.0.0.0 --port 8000

필요 패키지:
    - fastapi
    - uvicorn
    - pydantic
    - joblib
"""

import argparse
import os
import sys
from pathlib import Path
from typing import List

import joblib
import pandas as pd


def print_header(text):
    """헤더 출력"""
    print(f"\n{'=' * 60}")
    print(text)
    print('=' * 60)


def print_section(text):
    """섹션 출력"""
    print(f"\n{'-' * 60}")
    print(text)
    print('-' * 60)


def load_model(model_path):
    """모델 로드"""
    print(f"\n✓ 모델 로드 중: {model_path}")

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"모델 파일을 찾을 수 없습니다: {model_path}")

    model = joblib.load(model_path)
    print(f"✓ 모델 로드 완료: {type(model).__name__}")

    return model


def get_feature_names(sample_data_path, target_column):
    """샘플 데이터에서 특성 이름 추출"""
    print(f"\n✓ 특성 이름 추출 중: {sample_data_path}")

    df = pd.read_csv(sample_data_path)

    if target_column and target_column in df.columns:
        features = [col for col in df.columns if col != target_column]
    else:
        features = df.columns.tolist()

    print(f"✓ 특성 개수: {len(features)}개")

    return features


def generate_api_code(model_path, feature_names, output_dir, task_type='classification'):
    """FastAPI 코드 생성"""
    print_section("FastAPI 코드 생성")

    model_name = Path(model_path).stem

    # Pydantic 모델 정의 (입력 검증)
    features_str = ',\n        '.join([f"{feat}: float" for feat in feature_names])

    # 예측 타입
    if task_type == 'classification':
        prediction_response = """class PredictionResponse(BaseModel):
    prediction: int
    probability: Optional[List[float]] = None"""
    else:
        prediction_response = """class PredictionResponse(BaseModel):
    prediction: float"""

    api_code = f'''"""
FastAPI Model Serving
Generated by model-deployment plugin

실행:
    uvicorn app:app --reload --host 0.0.0.0 --port 8000

API 문서:
    http://localhost:8000/docs
"""

from typing import List, Optional

import joblib
import numpy as np
import pandas as pd
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# FastAPI 앱
app = FastAPI(
    title="{model_name} API",
    description="Machine Learning Model Serving API",
    version="1.0.0"
)

# 모델 로드 (서버 시작 시 1회)
MODEL_PATH = "{model_path}"
model = joblib.load(MODEL_PATH)

# 특성 이름
FEATURE_NAMES = {feature_names}


# Pydantic 모델 (입력 검증)
class PredictionRequest(BaseModel):
    {features_str}

    class Config:
        json_schema_extra = {{
            "example": {{
                {', '.join([f'"{feat}": 1.0' for feat in feature_names[:3]])}
            }}
        }}


# Pydantic 모델 (응답)
{prediction_response}


class HealthResponse(BaseModel):
    status: str
    model_type: str
    feature_count: int


@app.get("/", response_model=dict)
async def root():
    """API 루트"""
    return {{
        "message": "Model API is running",
        "docs": "/docs",
        "health": "/health",
        "predict": "/predict"
    }}


@app.get("/health", response_model=HealthResponse)
async def health():
    """헬스 체크"""
    return {{
        "status": "healthy",
        "model_type": type(model).__name__,
        "feature_count": len(FEATURE_NAMES)
    }}


@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """예측 수행"""
    try:
        # 입력 데이터를 DataFrame으로 변환
        input_data = pd.DataFrame([request.dict()], columns=FEATURE_NAMES)

        # 예측
        prediction = model.predict(input_data)[0]

        # 확률 (분류 모델인 경우)
        if hasattr(model, 'predict_proba'):
            probabilities = model.predict_proba(input_data)[0].tolist()
            return {{
                "prediction": int(prediction),
                "probability": probabilities
            }}
        else:
            return {{
                "prediction": float(prediction)
            }}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/batch_predict", response_model=List[PredictionResponse])
async def batch_predict(requests: List[PredictionRequest]):
    """배치 예측"""
    try:
        # 입력 데이터를 DataFrame으로 변환
        input_data = pd.DataFrame([req.dict() for req in requests], columns=FEATURE_NAMES)

        # 예측
        predictions = model.predict(input_data)

        # 결과 생성
        results = []
        if hasattr(model, 'predict_proba'):
            probabilities = model.predict_proba(input_data)
            for pred, prob in zip(predictions, probabilities):
                results.append({{
                    "prediction": int(pred),
                    "probability": prob.tolist()
                }})
        else:
            for pred in predictions:
                results.append({{"prediction": float(pred)}})

        return results

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
'''

    # 파일 저장
    api_path = os.path.join(output_dir, 'app.py')
    with open(api_path, 'w', encoding='utf-8') as f:
        f.write(api_code)

    print(f"✓ FastAPI 코드 저장: {api_path}")

    return api_path


def generate_dockerfile(output_dir, requirements_path=None):
    """Dockerfile 생성"""
    print_section("Dockerfile 생성")

    dockerfile_content = '''FROM python:3.10-slim

WORKDIR /app

# 시스템 패키지 업데이트
RUN apt-get update && apt-get install -y \\
    build-essential \\
    && rm -rf /var/lib/apt/lists/*

# Python 패키지 설치
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 애플리케이션 코드 복사
COPY app.py .
COPY model.pkl .

# 포트 노출
EXPOSE 8000

# 실행 명령어
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
'''

    dockerfile_path = os.path.join(output_dir, 'Dockerfile')
    with open(dockerfile_path, 'w', encoding='utf-8') as f:
        f.write(dockerfile_content)

    print(f"✓ Dockerfile 저장: {dockerfile_path}")

    return dockerfile_path


def generate_docker_compose(output_dir):
    """docker-compose.yml 생성"""
    print_section("docker-compose.yml 생성")

    compose_content = '''version: '3.8'

services:
  model-api:
    build: .
    ports:
      - "8000:8000"
    environment:
      - PYTHONUNBUFFERED=1
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3
'''

    compose_path = os.path.join(output_dir, 'docker-compose.yml')
    with open(compose_path, 'w', encoding='utf-8') as f:
        f.write(compose_content)

    print(f"✓ docker-compose.yml 저장: {compose_path}")

    return compose_path


def generate_requirements_txt(output_dir):
    """requirements.txt 생성"""
    requirements_content = '''# FastAPI & Server
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
pydantic>=2.5.0

# ML
scikit-learn>=1.3.0
pandas>=2.0.0
numpy>=1.24.0
joblib>=1.3.0

# Optional: Advanced models
xgboost>=2.0.0
lightgbm>=4.0.0
'''

    requirements_path = os.path.join(output_dir, 'requirements.txt')
    with open(requirements_path, 'w', encoding='utf-8') as f:
        f.write(requirements_content)

    print(f"✓ requirements.txt 저장: {requirements_path}")

    return requirements_path


def generate_readme(output_dir, model_name):
    """README.md 생성"""
    print_section("README.md 생성")

    readme_content = f'''# {model_name} API

FastAPI 기반 Machine Learning 모델 서빙 API입니다.

## 빠른 시작

### 1. 로컬 실행

```bash
# 의존성 설치
pip install -r requirements.txt

# API 서버 실행
uvicorn app:app --reload --host 0.0.0.0 --port 8000
```

API 문서: http://localhost:8000/docs

### 2. Docker 실행

```bash
# Docker 이미지 빌드
docker build -t {model_name}-api .

# 컨테이너 실행
docker run -p 8000:8000 {model_name}-api
```

### 3. Docker Compose 실행

```bash
docker-compose up -d
```

## API 엔드포인트

### GET /
API 정보 확인

### GET /health
헬스 체크
- 상태: healthy/unhealthy
- 모델 타입
- 특성 개수

### POST /predict
단일 예측
```json
{{
  "feature1": 1.0,
  "feature2": 2.0,
  ...
}}
```

**응답**:
```json
{{
  "prediction": 1,
  "probability": [0.2, 0.8]
}}
```

### POST /batch_predict
배치 예측
```json
[
  {{"feature1": 1.0, "feature2": 2.0}},
  {{"feature1": 3.0, "feature2": 4.0}}
]
```

## 테스트

```bash
# 헬스 체크
curl http://localhost:8000/health

# 예측 테스트
curl -X POST "http://localhost:8000/predict" \\
  -H "Content-Type: application/json" \\
  -d '{{"feature1": 1.0, "feature2": 2.0}}'
```

## 프로덕션 배포

### 환경 변수
- `PORT`: API 포트 (기본값: 8000)
- `WORKERS`: Uvicorn 워커 수 (기본값: 1)

### 성능 튜닝
```bash
uvicorn app:app --host 0.0.0.0 --port 8000 --workers 4
```

## 모니터링

- Prometheus 메트릭: `/metrics` (추가 설정 필요)
- 로그: stdout/stderr

## 보안

- API 키 인증 구현 권장
- HTTPS 사용
- Rate limiting 설정

## 라이선스

MIT License
'''

    readme_path = os.path.join(output_dir, 'README.md')
    with open(readme_path, 'w', encoding='utf-8') as f:
        f.write(readme_content)

    print(f"✓ README.md 저장: {readme_path}")

    return readme_path


def main():
    parser = argparse.ArgumentParser(description='모델 배포 API 생성 스크립트')
    parser.add_argument('--model-path', type=str, required=True,
                        help='학습된 모델 파일 경로 (.pkl)')
    parser.add_argument('--feature-names', type=str, default=None,
                        help='특성 이름 (쉼표로 구분, 예: V1,V2,V3)')
    parser.add_argument('--sample-data', type=str, default=None,
                        help='샘플 데이터 경로 (특성 이름 자동 추출)')
    parser.add_argument('--target-column', type=str, default=None,
                        help='타겟 컬럼명 (샘플 데이터 사용 시)')
    parser.add_argument('--task-type', type=str, choices=['classification', 'regression', 'auto'],
                        default='auto', help='태스크 타입')
    parser.add_argument('--output-dir', type=str, default=None,
                        help='출력 디렉토리')

    args = parser.parse_args()

    print_header("모델 배포 API 생성 시작")

    # 출력 디렉토리 설정
    if args.output_dir:
        output_dir = args.output_dir
    else:
        model_path = Path(args.model_path)
        if 'projects' in model_path.parts:
            project_idx = model_path.parts.index('projects')
            project_name = model_path.parts[project_idx + 1]
            output_dir = f"projects/{project_name}/deployment"
        else:
            output_dir = "deployment"

    os.makedirs(output_dir, exist_ok=True)
    print(f"✓ 출력 디렉토리: {output_dir}")

    # 모델 로드
    model = load_model(args.model_path)
    model_name = Path(args.model_path).stem

    # 태스크 타입 추정
    if args.task_type == 'auto':
        if hasattr(model, 'predict_proba'):
            task_type = 'classification'
        else:
            task_type = 'regression'
        print(f"\n✓ 자동 태스크 타입 감지: {task_type}")
    else:
        task_type = args.task_type

    # 특성 이름 추출
    if args.feature_names:
        feature_names = [f.strip() for f in args.feature_names.split(',')]
        print(f"\n✓ 특성 이름 (수동): {len(feature_names)}개")
    elif args.sample_data:
        feature_names = get_feature_names(args.sample_data, args.target_column)
    else:
        raise ValueError("--feature-names 또는 --sample-data 중 하나는 필수입니다.")

    # FastAPI 코드 생성
    api_path = generate_api_code(args.model_path, feature_names, output_dir, task_type)

    # Dockerfile 생성
    dockerfile_path = generate_dockerfile(output_dir)

    # docker-compose.yml 생성
    compose_path = generate_docker_compose(output_dir)

    # requirements.txt 생성
    requirements_path = generate_requirements_txt(output_dir)

    # README.md 생성
    readme_path = generate_readme(output_dir, model_name)

    # 모델 복사
    import shutil
    model_dest = os.path.join(output_dir, 'model.pkl')
    shutil.copy(args.model_path, model_dest)
    print(f"\n✓ 모델 복사: {model_dest}")

    print_header("모델 배포 API 생성 완료")
    print(f"\n📁 모든 파일이 생성되었습니다: {output_dir}/")
    print(f"   - app.py: FastAPI 애플리케이션")
    print(f"   - Dockerfile: Docker 이미지 빌드")
    print(f"   - docker-compose.yml: Docker Compose 설정")
    print(f"   - requirements.txt: Python 패키지")
    print(f"   - README.md: 사용 가이드")
    print(f"   - model.pkl: 학습된 모델")

    print(f"\n🚀 API 실행:")
    print(f"   cd {output_dir}")
    print(f"   uvicorn app:app --reload --host 0.0.0.0 --port 8000")
    print(f"\n   API 문서: http://localhost:8000/docs")

    print(f"\n🐳 Docker 실행:")
    print(f"   cd {output_dir}")
    print(f"   docker-compose up -d")

    return 0


if __name__ == '__main__':
    sys.exit(main())