from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from typing import List
import uvicorn
from hindi_translator import Translator

app = FastAPI(title="Translation API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

translator = Translator(model="gemma3:12b")

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

@app.post("/translate")
async def translate(request: dict):
    try:
        # Extract all fields from request
        client_id = request.get('client_id')
        target_language = request.get('target_language', 'English')
        slides_data = request.get('slides_data', [])

        if not slides_data:
            return {
                "success": False,
                "message": "No data provided",
                "client_id": client_id,
                "target_language": target_language,
                "slides_data": []
            }
        
        # Collect all text that needs translation
        sentences = []
        sentence_mapping = []  # Track which sentence belongs to which item and field
        
        for i, item in enumerate(slides_data):
            if isinstance(item, dict):
                # Add name field if present
                if item.get('name'):
                    sentences.append(item['name'])
                    sentence_mapping.append(('name', i))
                
                # Add backName field if present
                if item.get('backName'):
                    sentences.append(item['backName'])
                    sentence_mapping.append(('backName', i))
                
                # Add questionOption from options if present
                if item.get('options'):
                    for j, option in enumerate(item['options']):
                        if option.get('questionOption'):
                            sentences.append(option['questionOption'])
                            sentence_mapping.append(('questionOption', i, j))
        
        if not sentences:
            return {
                "success": False,
                "message": "No valid sentences provided",
                "client_id": client_id,
                "target_language": target_language,
                "slides_data": []
            }
        
        result = translator.translate_array(sentences, target_language)
        print(f"translate_array returned: {type(result)} - {result}")
        
        if isinstance(result, tuple) and len(result) == 2:
            translated, token_usage = result
        else:
            # Handle unexpected return format
            if isinstance(result, list):
                translated = result
                token_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
            else:
                raise Exception(f"Unexpected return type from translate_array: {type(result)}")
        
        # Create deep copy of original data to preserve structure
        import copy
        result_slides_data = copy.deepcopy(slides_data)

        # Convert contentId to string to match desired format
        for item in result_slides_data:
            if 'contentId' in item:
                item['contentId'] = str(item['contentId'])

        # Map translations back to their original locations
        for translation_idx, (field_type, *indices) in enumerate(sentence_mapping):
            if translation_idx < len(translated):
                if field_type == 'name':
                    item_idx = indices[0]
                    result_slides_data[item_idx]['name'] = translated[translation_idx]
                
                elif field_type == 'backName':
                    item_idx = indices[0]
                    result_slides_data[item_idx]['backName'] = translated[translation_idx]
                
                elif field_type == 'questionOption':
                    item_idx, option_idx = indices[0], indices[1]
                    result_slides_data[item_idx]['options'][option_idx]['questionOption'] = translated[translation_idx]
        
        return {
            "success": True,
            "message": "Translation completed successfully",
            "client_id": client_id,
            "target_language": target_language,
            "slides_data": result_slides_data,
            "token_usage": token_usage
        }
        
    except Exception as e:
        return {
            "success": False,
            "message": str(e),
            "client_id": request.get('client_id'),
            "target_language": request.get('target_language', 'English'),
            "slides_data": []
        }

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