api/routes/multimodal.py

"""
Multimodal Translation API Routes
Handles image, document, and website translation
"""

import asyncio
import logging
import base64
import tempfile
import os
from typing import Optional, Dict, Any
from datetime import datetime

from fastapi import APIRouter, HTTPException, Depends, UploadFile, File, Form
from pydantic import BaseModel, Field
from utils.auth import verify_platform_token

from services.multimodal_translation_service import multimodal_translator
from services.translation_service import unified_translation_service as real_translation_service
from services.cache_service import cache_service

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/multimodal", tags=["multimodal"])


class ImageTranslationRequest(BaseModel):
    """Request model for image translation"""
    image_base64: str = Field(..., description="Base64 encoded image data")
    mime_type: str = Field(..., description="MIME type of the image")
    source_language: Optional[str] = Field(None, description="Source language code")
    target_language: str = Field(..., description="Target language code")
    enhance_image: bool = Field(default=True, description="Enhance image for better OCR")
    annotate_image: bool = Field(default=False, description="Create annotated image with translated text")


class DocumentTranslationRequest(BaseModel):
    """Request model for document translation"""
    document_base64: str = Field(..., description="Base64 encoded document data")
    mime_type: str = Field(..., description="MIME type of the document")
    file_name: str = Field(..., description="Original file name")
    source_language: Optional[str] = Field(None, description="Source language code")
    target_language: str = Field(..., description="Target language code")
    preserve_format: bool = Field(default=True, description="Preserve document formatting")


class WebsiteTranslationRequest(BaseModel):
    """Request model for website translation"""
    url: str = Field(..., description="Website URL to translate")
    source_language: Optional[str] = Field(None, description="Source language code")
    target_language: str = Field(..., description="Target language code")
    extract_images: bool = Field(default=False, description="Extract and translate images")
    max_pages: int = Field(default=1, min=1, max=10, description="Maximum pages to process")


class ImageTranslationResponse(BaseModel):
    """Response model for image translation"""
    original_text: str
    translated_text: str
    detected_language: str
    target_language: str
    ocr_confidence: float
    translation_confidence: float
    quality_score: float
    model: str
    image_regions: list
    processing_time: float
    annotated_image: Optional[str] = None


class DocumentTranslationResponse(BaseModel):
    """Response model for document translation"""
    original_text: str
    translated_text: str
    detected_language: str
    target_language: str
    confidence: float
    quality_score: float
    model: str
    pages: int
    processing_time: float
    metadata: dict


class WebsiteTranslationResponse(BaseModel):
    """Response model for website translation"""
    original_text: str
    translated_text: str
    detected_language: str
    target_language: str
    confidence: float
    quality_score: float
    model: str
    title: str
    url: str
    images: list
    pages: int
    processing_time: float
    metadata: dict


@router.post("/image", response_model=ImageTranslationResponse)
async def translate_image(
    request: ImageTranslationRequest,
):
    """Translate text from image using OCR"""
    start_time = datetime.now()
    
    try:
        logger.info(f"Image translation request: {request.mime_type}, target: {request.target_language}")
        
        # Decode base64 image
        try:
            image_data = base64.b64decode(request.image_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 image data: {str(e)}")
        
        # Save to temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{request.mime_type.split('/')[-1]}") as tmp_file:
            tmp_file.write(image_data)
            tmp_path = tmp_file.name
        
        try:
            # Translate image
            result = await multimodal_translator.translate_image(
                image_path=tmp_path,
                target_lang=request.target_language,
                source_lang=request.source_language,
                enhance_image=request.enhance_image,
                annotate_image=request.annotate_image
            )
            
            # Convert annotated image to base64 if available
            annotated_image_b64 = None
            if result.annotated_image is not None and request.annotate_image:
                import cv2
                _, buffer = cv2.imencode('.jpg', result.annotated_image)
                annotated_image_b64 = base64.b64encode(buffer).decode('utf-8')
            
            processing_time = (datetime.now() - start_time).total_seconds()
            
            return ImageTranslationResponse(
                original_text=result.original_text,
                translated_text=result.translated_text,
                detected_language=result.detected_language,
                target_language=result.target_language,
                ocr_confidence=result.ocr_confidence,
                translation_confidence=result.translation_confidence,
                quality_score=min(result.ocr_confidence, result.translation_confidence),
                model="multimodal-ocr",
                image_regions=result.image_regions,
                processing_time=processing_time,
                annotated_image=annotated_image_b64
            )
            
        finally:
            # Clean up temporary file
            try:
                os.unlink(tmp_path)
            except Exception:
                pass
                
    except Exception as e:
        logger.error(f"Image translation failed: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Image translation failed: {str(e)}")


@router.post("/document", response_model=DocumentTranslationResponse)
async def translate_document(
    request: DocumentTranslationRequest,
):
    """Translate document content"""
    start_time = datetime.now()
    
    try:
        logger.info(f"Document translation request: {request.file_name}, target: {request.target_language}")
        
        # Decode base64 document
        try:
            document_data = base64.b64decode(request.document_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 document data: {str(e)}")
        
        # Save to temporary file
        file_extension = os.path.splitext(request.file_name)[1] or '.txt'
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
            tmp_file.write(document_data)
            tmp_path = tmp_file.name
        
        try:
            # Extract text from document
            from utils.file_parser import file_parser
            text_content, metadata = file_parser.parse_file(tmp_path)
            
            if not text_content.strip():
                raise HTTPException(status_code=400, detail="No text content found in document")
            
            # Detect language if needed
            source_lang = request.source_language
            if not source_lang:
                detection_result = await real_translation_service.detect_language(text_content[:1000])
                source_lang = detection_result['language']
            
            # Translate text
            if len(text_content) > 500:
                # Use long text translation for large documents
                translation_result = await real_translation_service.translate_long_text(
                    text=text_content,
                    source_lang=source_lang,
                    target_lang=request.target_language
                )
            else:
                # Use regular translation for small documents
                translation_result = await real_translation_service.translate(
                    text=text_content,
                    source_lang=source_lang,
                    target_lang=request.target_language
                )
            
            processing_time = (datetime.now() - start_time).total_seconds()
            
            return DocumentTranslationResponse(
                original_text=text_content,
                translated_text=translation_result.translated_text,
                detected_language=source_lang,
                target_language=request.target_language,
                confidence=translation_result.confidence_score,
                quality_score=translation_result.quality_score or 0.95,
                model=translation_result.model_used or "document-translator",
                pages=metadata.get('pages', 1),
                processing_time=processing_time,
                metadata=metadata
            )
            
        finally:
            # Clean up temporary file
            try:
                os.unlink(tmp_path)
            except Exception:
                pass
                
    except Exception as e:
        logger.error(f"Document translation failed: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Document translation failed: {str(e)}")


@router.post("/website", response_model=WebsiteTranslationResponse)
async def translate_website(
    request: WebsiteTranslationRequest,
    _: str = Depends(verify_platform_token)
):
    """Translate website content"""
    start_time = datetime.now()
    
    try:
        logger.info(f"Website translation request: {request.url}, target: {request.target_language}")
        
        # Web scraping and content extraction
        import requests
        from bs4 import BeautifulSoup
        
        # Fetch website content
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        
        response = requests.get(request.url, headers=headers, timeout=30)
        response.raise_for_status()
        
        # Parse HTML content
        soup = BeautifulSoup(response.content, 'html.parser')
        
        # Extract title
        title = soup.find('title')
        title_text = title.get_text().strip() if title else "Untitled"
        
        # Extract main text content
        # Remove script and style elements
        for script in soup(["script", "style"]):
            script.decompose()
        
        # Get text content
        text_content = soup.get_text()
        
        # Clean up text
        lines = (line.strip() for line in text_content.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        text_content = ' '.join(chunk for chunk in chunks if chunk)
        
        if not text_content.strip():
            raise HTTPException(status_code=400, detail="No text content found on website")
        
        # Detect language if needed
        source_lang = request.source_language
        if not source_lang:
            detection_result = await real_translation_service.detect_language(text_content[:1000])
            source_lang = detection_result['language']
        
        # Translate text
        if len(text_content) > 500:
            # Use long text translation for large websites
            translation_result = await real_translation_service.translate_long_text(
                text=text_content,
                source_lang=source_lang,
                target_lang=request.target_language
            )
        else:
            # Use regular translation for small websites
            translation_result = await real_translation_service.translate(
                text=text_content,
                source_lang=source_lang,
                target_lang=request.target_language
            )
        
        # Extract images if requested
        images = []
        if request.extract_images:
            img_tags = soup.find_all('img')
            for img in img_tags[:10]:  # Limit to 10 images
                src = img.get('src')
                if src:
                    # Convert relative URLs to absolute
                    if src.startswith('//'):
                        src = 'https:' + src
                    elif src.startswith('/'):
                        from urllib.parse import urljoin
                        src = urljoin(request.url, src)
                    images.append({
                        'src': src,
                        'alt': img.get('alt', ''),
                        'title': img.get('title', '')
                    })
        
        processing_time = (datetime.now() - start_time).total_seconds()
        
        return WebsiteTranslationResponse(
            original_text=text_content,
            translated_text=translation_result.translated_text,
            detected_language=source_lang,
            target_language=request.target_language,
            confidence=translation_result.confidence_score,
            quality_score=translation_result.quality_score or 0.95,
            model=translation_result.model_used or "website-translator",
            title=title_text,
            url=request.url,
            images=images,
            pages=1,  # Single page for now
            processing_time=processing_time,
            metadata={
                'user_agent': headers['User-Agent'],
                'status_code': response.status_code,
                'content_type': response.headers.get('content-type', ''),
                'content_length': len(response.content)
            }
        )
        
    except requests.RequestException as e:
        logger.error(f"Website request failed: {str(e)}")
        raise HTTPException(status_code=400, detail=f"Failed to fetch website: {str(e)}")
    except Exception as e:
        logger.error(f"Website translation failed: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Website translation failed: {str(e)}")


@router.get("/supported-types")
async def get_supported_types(_: str = Depends(verify_platform_token)):
    """Get supported file types for multimodal translation"""
    return {
        "success": True,
        "data": {
            "images": {
                "extensions": [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff", ".svg"],
                "mime_types": ["image/jpeg", "image/jpg", "image/png", "image/gif", "image/bmp", "image/webp", "image/tiff", "image/svg+xml"],
                "max_size": "10MB",
                "features": ["OCR", "Text Detection", "Context Analysis", "Quality Scoring"]
            },
            "documents": {
                "extensions": [".pdf", ".docx", ".doc", ".txt", ".rtf"],
                "mime_types": ["application/pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/msword", "text/plain", "application/rtf"],
                "max_size": "10MB",
                "features": ["Text Extraction", "Format Preservation", "Batch Processing", "Quality Validation"]
            },
            "websites": {
                "features": ["Web Scraping", "Content Extraction", "Image Translation", "Multi-page Support"],
                "max_pages": 10,
                "supported_domains": "All public websites"
            }
        }
    }


@router.get("/health")
async def multimodal_health_check(_: str = Depends(verify_platform_token)):
    """Health check for multimodal translation services"""
    try:
        # Check OCR engines
        ocr_engines = multimodal_translator.ocr_engines.keys()
        
        return {
            "success": True,
            "data": {
                "status": "healthy",
                "services": {
                    "image_translation": len(ocr_engines) > 0,
                    "document_translation": True,
                    "website_translation": True,
                    "ocr_engines": list(ocr_engines),
                },
                "timestamp": datetime.now().isoformat(),
            }
        }
    except Exception as e:
        logger.error(f"Multimodal health check failed: {str(e)}")
        return {
            "success": False,
            "error": "Health check failed",
            "details": str(e),
        }