from fastapi import FastAPI, HTTPException, UploadFile, File, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel
import os
import uuid
import subprocess
import logging

# Initialize FastAPI
app = FastAPI(
    title="HF Utility API",
    description="Fast APIs for Embedding and Docx to PDF conversion"
)

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

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize the embedding model globally
try:
    from sentence_transformers import SentenceTransformer
    # Load model. Using a small, fast model for <1s inference.
    embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
    logger.info("SentenceTransformer model loaded successfully.")
except Exception as e:
    logger.error(f"Failed to load embedding model: {e}")
    embedding_model = None

class EmbedRequest(BaseModel):
    text: str

@app.post("/api/embed")
async def embed_text(req: EmbedRequest):
    if not embedding_model:
        raise HTTPException(status_code=500, detail="Embedding model not initialized.")
    
    try:
        # Generate embeddings
        embedding = embedding_model.encode(req.text).tolist()
        return {"embedding": embedding}
    except Exception as e:
        logger.error(f"Embedding error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

def cleanup_files(*file_paths):
    for path in file_paths:
        if path and os.path.exists(path):
            try:
                os.remove(path)
                logger.info(f"Cleaned up {path}")
            except Exception as e:
                logger.error(f"Failed to clean up {path}: {e}")

@app.post("/api/convert-docx")
async def convert_docx(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    if not file.filename.endswith(".docx"):
        raise HTTPException(status_code=400, detail="Only .docx files are supported.")
    
    # Generate unique IDs for the files to avoid collisions
    job_id = str(uuid.uuid4())
    upload_dir = "/tmp/hf_api_uploads"
    # Ensure dir exists
    os.makedirs(upload_dir, exist_ok=True)
    
    # sanitize filename a bit
    safe_filename = file.filename.replace(" ", "_")
    
    input_path = os.path.join(upload_dir, f"{job_id}_{safe_filename}")
    output_filename = f"{job_id}_{safe_filename.replace('.docx', '.pdf')}"
    output_path = os.path.join(upload_dir, output_filename)
    
    try:
        # Save uploaded file
        with open(input_path, "wb") as f:
            content = await file.read()
            f.write(content)
            
        logger.info(f"Saved docx to {input_path}. Starting conversion...")
        
        # Run LibreOffice headless conversion
        command = [
            "libreoffice",
            "--headless",
            "--convert-to", "pdf",
            "--outdir", upload_dir,
            input_path
        ]
        
        # Execute the command
        process = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        if process.returncode != 0:
            logger.error(f"LibreOffice error: {process.stderr}")
            # Clean up the input file on error
            cleanup_files(input_path)
            raise HTTPException(status_code=500, detail="Failed to convert document.")
            
        logger.info(f"Successfully converted to {output_path}")
        
        # Check if output file was created
        if not os.path.exists(output_path):
            cleanup_files(input_path)
            raise HTTPException(status_code=500, detail="Output PDF not found after conversion.")
            
        # Add cleanup to background tasks so they are deleted AFTER the response is sent
        background_tasks.add_task(cleanup_files, input_path, output_path)
            
        return FileResponse(
            path=output_path, 
            filename=safe_filename.replace('.docx', '.pdf'), 
            media_type="application/pdf"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Conversion error: {e}")
        cleanup_files(input_path, output_path)
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/health")
async def health_check():
    return {"status": "healthy", "embedding_model_loaded": embedding_model is not None}

@app.get("/")
async def root():
    return {"message": "HF Utility API is running"}