#!/usr/bin/env python3
"""
API 테스트 스크립트
각 엔드포인트가 제대로 작동하는지 테스트합니다.
"""
import os
import sys
import base64
import time
import requests
import json
from pathlib import Path

# 프로젝트 루트를 Python 경로에 추가
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

class APITester:
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url
        self.test_images_dir = os.path.join(project_root, "tests", "images")
        self.results_dir = os.path.join(project_root, "tests", "results")
        
        # 결과 디렉토리 생성
        os.makedirs(self.results_dir, exist_ok=True)
        
        print(f"🚀 API 테스터 초기화 완료")
        print(f"   서버 URL: {base_url}")
        print(f"   테스트 이미지: {self.test_images_dir}")
        print(f"   결과 저장: {self.results_dir}")
    
    def image_to_base64(self, image_path):
        """이미지를 base64로 인코딩"""
        try:
            with open(image_path, "rb") as image_file:
                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            return encoded_string
        except Exception as e:
            print(f"❌ 이미지 인코딩 실패: {e}")
            return None
    
    def save_base64_image(self, base64_string, filename):
        """base64 이미지를 파일로 저장"""
        try:
            image_data = base64.b64decode(base64_string)
            filepath = os.path.join(self.results_dir, filename)
            with open(filepath, "wb") as f:
                f.write(image_data)
            print(f"   💾 결과 이미지 저장: {filename}")
            return True
        except Exception as e:
            print(f"❌ 이미지 저장 실패: {e}")
            return False
    
    def test_health(self):
        """헬스 체크 테스트"""
        print("\n🏥 헬스 체크 테스트...")
        try:
            response = requests.get(f"{self.base_url}/health", timeout=10)
            if response.status_code == 200:
                data = response.json()
                print(f"   ✅ 서버 상태: {data.get('status', 'unknown')}")
                print(f"   📊 가동 시간: {data.get('uptime', 0):.2f}초")
                return True
            else:
                print(f"   ❌ 서버 응답 실패: {response.status_code}")
                return False
        except Exception as e:
            print(f"   ❌ 연결 실패: {e}")
            return False
    
    def test_server_config(self):
        """서버 설정 정보 테스트"""
        print("\n⚙️ 서버 설정 정보 테스트...")
        try:
            response = requests.get(f"{self.base_url}/api/v1/server-config", timeout=10)
            if response.status_code == 200:
                data = response.json()
                print(f"   ✅ 서버 설정 조회 성공")
                print(f"   🖥️ 시스템: {'Jetson Xavier' if data.get('is_jetson') else 'x86_64'}")
                print(f"   🎮 디바이스: {data.get('device', 'unknown')}")
                print(f"   📁 최대 파일 크기: {data.get('max_file_size', 0)}MB")
                print(f"   🎯 사용 가능한 모델: {len(data.get('models', []))}개")
                return True
            else:
                print(f"   ❌ 서버 설정 조회 실패: {response.status_code}")
                return False
        except Exception as e:
            print(f"   ❌ 연결 실패: {e}")
            return False
    
    def test_inpaint(self, image_name="test_image_512.png", mask_name="test_mask_512.png"):
        """인페인팅 API 테스트"""
        print(f"\n🎨 인페인팅 테스트 ({image_name} + {mask_name})...")
        
        # 이미지 파일 경로
        image_path = os.path.join(self.test_images_dir, image_name)
        mask_path = os.path.join(self.test_images_dir, mask_name)
        
        if not os.path.exists(image_path) or not os.path.exists(mask_path):
            print(f"   ❌ 테스트 이미지 파일이 없습니다")
            return False
        
        try:
            # 이미지를 base64로 인코딩
            image_base64 = self.image_to_base64(image_path)
            mask_base64 = self.image_to_base64(mask_path)
            
            if not image_base64 or not mask_base64:
                return False
            
            # 인페인팅 요청
            payload = {
                "image": image_base64,
                "mask": mask_base64,
                "model_name": "simple-lama",
                "sd_seed": 42,
                "prompt": "beautiful landscape",
                "negative_prompt": "ugly, blurry",
                "num_inference_steps": 20,
                "guidance_scale": 7.5,
                "strength": 1.0
            }
            
            start_time = time.time()
            response = requests.post(
                f"{self.base_url}/api/v1/inpaint",
                json=payload,
                timeout=60
            )
            processing_time = time.time() - start_time
            
            if response.status_code == 200:
                data = response.json()
                print(f"   ✅ 인페인팅 성공!")
                print(f"   ⏱️ 처리 시간: {data.get('processing_time', processing_time):.2f}초")
                print(f"   🎲 사용된 시드: {data.get('seed', 'unknown')}")
                
                # 결과 이미지 저장
                if data.get('image'):
                    self.save_base64_image(data['image'], f"inpaint_result_{image_name}")
                
                return True
            else:
                print(f"   ❌ 인페인팅 실패: {response.status_code}")
                try:
                    error_data = response.json()
                    print(f"   📝 에러 메시지: {error_data.get('detail', 'unknown error')}")
                except:
                    print(f"   📝 응답 내용: {response.text}")
                return False
                
        except Exception as e:
            print(f"   ❌ 인페인팅 테스트 실패: {e}")
            return False
    
    def test_remove_bg(self, image_name="test_image_512.png"):
        """배경 제거 API 테스트"""
        print(f"\n🖼️ 배경 제거 테스트 ({image_name})...")
        
        # 이미지 파일 경로
        image_path = os.path.join(self.test_images_dir, image_name)
        
        if not os.path.exists(image_path):
            print(f"   ❌ 테스트 이미지 파일이 없습니다")
            return False
        
        try:
            # 이미지를 base64로 인코딩
            image_base64 = self.image_to_base64(image_path)
            
            if not image_base64:
                return False
            
            # 배경 제거 요청
            payload = {
                "image": image_base64,
                "model_name": "rembg"
            }
            
            start_time = time.time()
            response = requests.post(
                f"{self.base_url}/api/v1/remove_bg",
                json=payload,
                timeout=60
            )
            processing_time = time.time() - start_time
            
            if response.status_code == 200:
                data = response.json()
                print(f"   ✅ 배경 제거 성공!")
                print(f"   ⏱️ 처리 시간: {data.get('processing_time', processing_time):.2f}초")
                
                # 결과 이미지와 마스크 저장
                if data.get('image'):
                    self.save_base64_image(data['image'], f"rembg_result_{image_name}")
                if data.get('mask'):
                    self.save_base64_image(data['mask'], f"rembg_mask_{image_name}")
                
                return True
            else:
                print(f"   ❌ 배경 제거 실패: {response.status_code}")
                try:
                    error_data = response.json()
                    print(f"   📝 에러 메시지: {error_data.get('detail', 'unknown error')}")
                except:
                    print(f"   📝 응답 내용: {response.text}")
                return False
                
        except Exception as e:
            print(f"   ❌ 배경 제거 테스트 실패: {e}")
            return False
    
    def test_plugin_generate_image(self, image_name="test_image_512.png"):
        """플러그인 이미지 생성 테스트"""
        print(f"\n🔌 플러그인 이미지 생성 테스트 ({image_name})...")
        
        # 이미지 파일 경로
        image_path = os.path.join(self.test_images_dir, image_name)
        
        if not os.path.exists(image_path):
            print(f"   ❌ 테스트 이미지 파일이 없습니다")
            return False
        
        try:
            # 이미지를 base64로 인코딩
            image_base64 = self.image_to_base64(image_path)
            
            if not image_base64:
                return False
            
            # 플러그인 요청
            payload = {
                "name": "rembg",
                "image": image_base64,
                "model_name": "rembg"
            }
            
            start_time = time.time()
            response = requests.post(
                f"{self.base_url}/api/v1/run_plugin_gen_image",
                json=payload,
                timeout=60
            )
            processing_time = time.time() - start_time
            
            if response.status_code == 200:
                data = response.json()
                print(f"   ✅ 플러그인 이미지 생성 성공!")
                print(f"   ⏱️ 처리 시간: {data.get('processing_time', processing_time):.2f}초")
                
                # 결과 이미지 저장
                if data.get('image'):
                    self.save_base64_image(data['image'], f"plugin_result_{image_name}")
                
                return True
            else:
                print(f"   ❌ 플러그인 이미지 생성 실패: {response.status_code}")
                try:
                    error_data = response.json()
                    print(f"   📝 에러 메시지: {error_data.get('detail', 'unknown error')}")
                except:
                    print(f"   📝 응답 내용: {response.text}")
                return False
                
        except Exception as e:
            print(f"   ❌ 플러그인 테스트 실패: {e}")
            return False
    
    def test_samplers(self):
        """샘플러 목록 테스트"""
        print("\n🎲 샘플러 목록 테스트...")
        try:
            response = requests.get(f"{self.base_url}/api/v1/samplers", timeout=10)
            if response.status_code == 200:
                samplers = response.json()
                print(f"   ✅ 샘플러 목록 조회 성공")
                print(f"   📊 사용 가능한 샘플러: {len(samplers)}개")
                print(f"   🎯 샘플러들: {', '.join(samplers[:5])}{'...' if len(samplers) > 5 else ''}")
                return True
            else:
                print(f"   ❌ 샘플러 목록 조회 실패: {response.status_code}")
                return False
        except Exception as e:
            print(f"   ❌ 연결 실패: {e}")
            return False
    
    def run_all_tests(self):
        """모든 테스트 실행"""
        print("🧪 전체 API 테스트 시작")
        print("=" * 50)
        
        test_results = []
        
        # 1. 헬스 체크
        test_results.append(("헬스 체크", self.test_health()))
        
        # 2. 서버 설정 정보
        test_results.append(("서버 설정 정보", self.test_server_config()))
        
        # 3. 샘플러 목록
        test_results.append(("샘플러 목록", self.test_samplers()))
        
        # 4. 인페인팅 테스트 (작은 이미지)
        test_results.append(("인페인팅 (256x256)", self.test_inpaint("test_image_256.png", "test_mask_256.png")))
        
        # 5. 인페인팅 테스트 (기본 이미지)
        test_results.append(("인페인팅 (512x512)", self.test_inpaint("test_image_512.png", "test_mask_512.png")))
        
        # 6. 배경 제거 테스트
        test_results.append(("배경 제거", self.test_remove_bg("test_image_512.png")))
        
        # 7. 플러그인 테스트
        test_results.append(("플러그인 이미지 생성", self.test_plugin_generate_image("test_image_512.png")))
        
        # 결과 요약
        print("\n" + "=" * 50)
        print("📊 테스트 결과 요약")
        print("=" * 50)
        
        passed = 0
        total = len(test_results)
        
        for test_name, result in test_results:
            status = "✅ PASS" if result else "❌ FAIL"
            print(f"{test_name:<25} {status}")
            if result:
                passed += 1
        
        print("=" * 50)
        print(f"전체: {total}개, 성공: {passed}개, 실패: {total - passed}개")
        
        if passed == total:
            print("🎉 모든 테스트가 성공했습니다!")
        else:
            print("⚠️ 일부 테스트가 실패했습니다.")
        
        return passed == total

def main():
    """메인 함수"""
    import argparse
    
    parser = argparse.ArgumentParser(description="API 테스트 스크립트")
    parser.add_argument("--url", default="http://localhost:8000", help="테스트할 서버 URL")
    parser.add_argument("--single", help="단일 테스트 실행 (health, config, inpaint, rembg, plugin)")
    
    args = parser.parse_args()
    
    tester = APITester(args.url)
    
    if args.single:
        # 단일 테스트 실행
        if args.single == "health":
            tester.test_health()
        elif args.single == "config":
            tester.test_server_config()
        elif args.single == "inpaint":
            tester.test_inpaint()
        elif args.single == "rembg":
            tester.test_remove_bg()
        elif args.single == "plugin":
            tester.test_plugin_generate_image()
        else:
            print(f"❌ 알 수 없는 테스트: {args.single}")
    else:
        # 전체 테스트 실행
        tester.run_all_tests()

if __name__ == "__main__":
    main()