import requests
import base64
import json
from pathlib import Path

class TTSClient:
    def __init__(self, base_url="http://localhost:7860"):
        self.base_url = base_url.rstrip("/")
    
    def health_check(self):
        """APIの健康状態をチェック"""
        try:
            response = requests.get(f"{self.base_url}/health")
            return response.status_code == 200, response.json()
        except Exception as e:
            return False, str(e)
    
    def synthesize_to_file(self, text, output_path="output.wav", sample_rate=22050):
        """テキストを音声に変換してファイルに保存"""
        try:
            response = requests.post(
                f"{self.base_url}/synthesize",
                json={"text": text, "sample_rate": sample_rate}
            )
            
            if response.status_code == 200:
                with open(output_path, "wb") as f:
                    f.write(response.content)
                return True, f"Audio saved to {output_path}"
            else:
                return False, f"Error: {response.status_code} - {response.text}"
                
        except Exception as e:
            return False, str(e)
    
    def synthesize_to_base64(self, text, sample_rate=22050):
        """テキストを音声に変換してBase64形式で取得"""
        try:
            response = requests.post(
                f"{self.base_url}/synthesize_base64",
                json={"text": text, "sample_rate": sample_rate}
            )
            
            if response.status_code == 200:
                return True, response.json()
            else:
                return False, f"Error: {response.status_code} - {response.text}"
                
        except Exception as e:
            return False, str(e)
    
    def save_base64_audio(self, audio_base64, output_path="output_from_base64.wav"):
        """Base64形式の音声データをファイルに保存"""
        try:
            audio_data = base64.b64decode(audio_base64)
            with open(output_path, "wb") as f:
                f.write(audio_data)
            return True, f"Audio saved to {output_path}"
        except Exception as e:
            return False, str(e)

def main():
    # クライアントを初期化（ローカル開発用）
    client = TTSClient("http://localhost:7860")
    
    # Hugging Face Spacesデプロイ用（URLを適切に変更してください）
    # client = TTSClient("https://your-space-name.hf.space")
    
    print("=== SpeechBrain TTS API Client Test ===\n")
    
    # 健康状態チェック
    print("1. Health Check:")
    is_healthy, health_result = client.health_check()
    print(f"   Status: {'✓ Healthy' if is_healthy else '✗ Unhealthy'}")
    print(f"   Response: {health_result}\n")
    
    if not is_healthy:
        print("API is not available. Please check if the server is running.")
        return
    
    # テスト用テキスト
    test_texts = [
        "Hello, this is a test of the SpeechBrain TTS API.",
        "The quick brown fox jumps over the lazy dog.",
        "Welcome to Hugging Face Spaces!"
    ]
    
    # テスト1: WAVファイル直接保存
    print("2. Testing direct WAV file synthesis:")
    for i, text in enumerate(test_texts):
        print(f"   Testing: '{text}'")
        success, result = client.synthesize_to_file(
            text, 
            f"test_output_{i+1}.wav",
            sample_rate=22050
        )
        print(f"   Result: {'✓ Success' if success else '✗ Failed'} - {result}")
    print()
    
    # テスト2: Base64形式での取得
    print("3. Testing Base64 synthesis:")
    for i, text in enumerate(test_texts):
        print(f"   Testing: '{text}'")
        success, result = client.synthesize_to_base64(text, sample_rate=22050)
        
        if success:
            print(f"   ✓ Success - Audio length: {len(result['audio_base64'])} chars")
            
            # Base64からファイルに保存
            save_success, save_result = client.save_base64_audio(
                result['audio_base64'],
                f"test_base64_{i+1}.wav"
            )
            print(f"   Save result: {'✓ Success' if save_success else '✗ Failed'} - {save_result}")
        else:
            print(f"   ✗ Failed - {result}")
    print()
    
    # テスト3: エラーハンドリング
    print("4. Testing error handling:")
    
    # 空のテキスト
    print("   Testing empty text:")
    success, result = client.synthesize_to_file("", "empty_test.wav")
    print(f"   Result: {'✓ Success' if success else '✗ Expected failure'} - {result}")
    
    # 長すぎるテキスト
    print("   Testing too long text:")
    long_text = "This is a very long text. " * 50  # 500文字以上
    success, result = client.synthesize_to_file(long_text, "long_test.wav")
    print(f"   Result: {'✓ Success' if success else '✗ Expected failure'} - {result}")
    
    print("\n=== Test Complete ===")
    print("Check the generated audio files:")
    for i in range(len(test_texts)):
        print(f"   - test_output_{i+1}.wav")
        print(f"   - test_base64_{i+1}.wav")

if __name__ == "__main__":
    main()