from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
import os
import sys
import uuid
from pydantic import BaseModel
from typing import List, Dict
import stat
import os
import stat
import uuid
import requests
from pathlib import Path
from typing import List, Dict, Any
import torch
import numpy as np
import soundfile as sf

import logging
from fastapi.responses import JSONResponse

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# Your environment path
ENV_PATH = "/var/www/eduai.edurigo.com/doc_train/edurigo_ai/Puru/ai4bharat/ai4bharat_env"
print("SS")
def activate_environment():
    """Activate the virtual environment"""
    site_packages = os.path.join(ENV_PATH, "lib", "python3.10", "site-packages")
    if site_packages not in sys.path:
        sys.path.insert(0, site_packages)
    os.environ['VIRTUAL_ENV'] = ENV_PATH
    os.environ['PATH'] = os.path.join(ENV_PATH, 'bin') + os.pathsep + os.environ.get('PATH', '')
print(os.environ)
# Activate environment at import time
activate_environment()
print("ssaaa")
print(os.environ)

from transformers import VitsModel, AutoTokenizer
import torch
import scipy 
import soundfile as sf
import numpy as np

class TTSRequest(BaseModel):
    text: str

class TTSItem(BaseModel):
    id: str
    content: str

class TTSBatchRequest(BaseModel):
    data: List[TTSItem]

class TTSBatchRequest2(BaseModel):
    clientId: int
    storigoId: int
    data: List[TTSItem]


app = FastAPI()


@app.post("/generate-tts-batch")
async def generate_tts_batch(request: TTSBatchRequest2):
    """Generate TTS for multiple items with dynamic folder creation"""

    if not request.data:
        raise HTTPException(status_code=400, detail="Data array cannot be empty")
    
    try:
        #base_folder = "convert_mp3"
        base_folder = "/var/www/eduai.edurigo.com/lang_trans_audio/testing/convert_mp3"
        storigo_folder = f"storigo_{request.storigoId}"
        full_folder_path = os.path.join(base_folder, storigo_folder)

        if not os.path.exists(full_folder_path):
            os.makedirs(full_folder_path, exist_ok=True)
            os.chmod(full_folder_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
        
        model = VitsModel.from_pretrained("facebook/mms-tts-hin")
        tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hin")
        results = []
        file_data_for_upload = []

        for item in request.data:

            if not item.content.strip():
                results.append({
                    "id": item.id,
                    "success": False,
                    "error": "Content cannot be empty"
                })
                continue

            try:
                inputs = tokenizer(item.content, return_tensors="pt")
                with torch.no_grad():
                    output = model(**inputs).waveform

                output_np = output.cpu().numpy().astype(np.float32)
                if output_np.ndim == 2 and output_np.shape[0] == 1:
                    output_np = output_np.squeeze(0)

                filename = f"techno_{item.id}_{str(uuid.uuid4())[:8]}.mp3"
                file_path = os.path.join(full_folder_path, filename)

                sf.write(file_path, output_np, samplerate=16000)

                download_url = f"{base_folder}/{storigo_folder}/{filename}"

                file_data_for_upload.append({
                    "id": item.id,
                    "filename": filename,
                    "local_path": file_path,
                    "storigo_id": request.storigoId,
                    "client_id": request.clientId
                    
                })

                results.append({
                    "id": item.id,
                    "success": True,
                    "filename": filename,
                    "download_url": download_url,
                    "content": item.content,
                    "folder": storigo_folder
                })

            except Exception as e:
                logger.error("Error generating TTS for item ID %s: %s", item.id, str(e))
                results.append({
                    "id": item.id,
                    "success": False,
                    "error": f"TTS generation failed: {str(e)}"
                })

        php_response = await send_files_to_php_server(file_data_for_upload)
        clean_audio(php_response, file_data_for_upload)

        response_data = {
        "success": True,
        "message": f"Processed {len(results)} items",
        "folder": storigo_folder,
        "results": results,
        "php_server_response": php_response
        }
        # Log the response for debugging

        # Return a proper JSON response
        return JSONResponse(status_code=200, content=response_data)


    except Exception as e:
        logger.exception("Batch TTS generation failed")
        raise HTTPException(status_code=500, detail=f"Batch TTS generation failed: {str(e)}")

async def send_files_to_php_server(file_data_list: List[Dict]):

    php_server_url = "https://api.edurigo1.com/uploadaudiofile/uploadAudio"
    

    upload_results = []
    
    for file_data in file_data_list:
        try:
            file_path = file_data["local_path"]
            client_id = file_data["client_id"]
            storigo_id = file_data["storigo_id"]
            filename = file_data["filename"]
            resource_id = file_data["id"]

            with open(file_path, 'rb') as f:
                files = {
                    'audioFile': (filename, f, 'audio/mp3')  # match $_FILES['audio_file']
                }
                data = {
                    'clientId': client_id,
                    'storigoId': storigo_id,
                    'fileName': filename,
                    'resourceId': resource_id,
                }

                response = requests.post(php_server_url, files=files, data=data)
                result_json = response.json()

                upload_results.append({
                    "filename": filename,
                    "success": result_json.get("status") == "success",
                    "message": result_json.get("message"),
                    "data": result_json.get("data", {})
                })

        except Exception as e:
            upload_results.append({
                "filename": file_data.get("filename"),
                "success": False,
                "message": str(e)
            })

    return {
        "status": "completed",
        "results": upload_results
    }

def clean_audio(upload_results, file_data_list):
    """Clean up local files after successful uploads"""
    
    # Create filename to file_data mapping
    file_map = {data["filename"]: data for data in file_data_list}
    
    for result in upload_results["results"]:
        if result["success"]:
            filename = result["filename"]
            file_data = file_map.get(filename)
            
            if file_data:
                local_path = file_data["local_path"]
                
                # Delete the local file
                try:
                    if os.path.exists(local_path):
                        os.remove(local_path)
                        print(f"Deleted: {local_path}")
                except Exception as e:
                    print(f"Error deleting {local_path}: {e}")