import requests
import json
import base64
import os
from pydub import AudioSegment
import io
import argparse

def text_to_speech_ollama(text, model_name="mistral-tts", output_path="output.wav"):
    """
    Convertește text în audio folosind un model TTS prin Ollama
    
    Args:
        text: Textul în limba română pentru a fi convertit în audio
        model_name: Numele modelului de folosit (trebuie să fie un model TTS disponibil în Ollama)
        output_path: Calea unde va fi salvat fișierul audio
    """
    # URL-ul API-ului Ollama (presupunând că rulează local)
    url = "http://localhost:11434/api/generate"
    
    # Construirea solicitării pentru API
    payload = {
        "model": model_name,
        "prompt": text,
        "stream": False,
        "options": {
            "language": "ro",  # Specificare limba română
            "temperature": 0.7,
            "format": "wav"
        }
    }
    
    print(f"Se generează audio pentru textul: {text}")
    
    try:
        response = requests.post(url, json=payload)
        
        if response.status_code == 200:
            # Presupunem că răspunsul conține audio codificat base64
            response_data = response.json()
            
            if "audio" in response_data:
                # Decodificare audio din base64
                audio_data = base64.b64decode(response_data["audio"])
                
                # Salvare audio în fișier
                with open(output_path, "wb") as audio_file:
                    audio_file.write(audio_data)
                print(f"Audio generat cu succes și salvat la: {output_path}")
                return True
            else:
                print("Răspunsul nu conține date audio.")
                print(f"Răspuns primit: {response_data}")
                return False
        else:
            print(f"Eroare la generarea audio: {response.status_code}")
            print(f"Răspuns: {response.text}")
            return False
            
    except Exception as e:
        print(f"Eroare la procesarea cererii: {str(e)}")
        return False

def load_model_if_needed(model_name):
    """
    Verifică dacă modelul este disponibil local și îl încarcă dacă este necesar
    """
    url = "http://localhost:11434/api/tags"
    
    try:
        response = requests.get(url)
        models = response.json().get("models", [])
        model_names = [model["name"] for model in models]
        
        if model_name not in model_names:
            print(f"Modelul {model_name} nu este încărcat. Se încarcă acum...")
            pull_url = "http://localhost:11434/api/pull"
            pull_payload = {"name": model_name}
            pull_response = requests.post(pull_url, json=pull_payload, stream=True)
            
            for line in pull_response.iter_lines():
                if line:
                    print(json.loads(line.decode("utf-8")).get("status", ""))
            
            print(f"Modelul {model_name} a fost încărcat cu succes!")
        else:
            print(f"Modelul {model_name} este deja disponibil.")
    
    except Exception as e:
        print(f"Eroare la verificarea modelelor disponibile: {str(e)}")

def main():
    # Parsare argumente linie de comandă
    parser = argparse.ArgumentParser(description="Convertire text în audio în limba română folosind Ollama")
    parser.add_argument("--text", type=str, help="Textul de transformat în audio")
    parser.add_argument("--file", type=str, help="Fișier text de transformat în audio")
    parser.add_argument("--model", type=str, default="mistral-tts", help="Model TTS de folosit")
    parser.add_argument("--output", type=str, default="output.wav", help="Calea fișierului audio de output")
    
    args = parser.parse_args()
    
    if not args.text and not args.file:
        print("Eroare: Trebuie să specificați fie un text, fie un fișier text de transformat.")
        return
    
    # Verifică și încarcă modelul dacă e necesar
    load_model_if_needed(args.model)
    
    # Procesează textul
    if args.file:
        try:
            with open(args.file, "r", encoding="utf-8") as f:
                text = f.read()
            text_to_speech_ollama(text, args.model, args.output)
        except Exception as e:
            print(f"Eroare la citirea fișierului: {str(e)}")
    else:
        text_to_speech_ollama(args.text, args.model, args.output)

if __name__ == "__main__":
    main()