from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from transformers import pipeline
import traceback

app = FastAPI()

# Enable CORS to allow POST requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


app.mount("/assets", StaticFiles(directory="."), name="static")

# Serve the HTML page at the root
@app.get("/", response_class=HTMLResponse)
async def serve_html():
    with open("index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

# Load pre-trained models from Hugging Face Hub
try:
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
    translators = {
        "fr": pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr"),
        "de": pipeline("translation", model="Helsinki-NLP/opus-mt-en-de"),
        "es": pipeline("translation", model="Helsinki-NLP/opus-mt-en-es"),
        "ar": pipeline("translation", model="Helsinki-NLP/opus-mt-en-ar")
    }
except Exception as e:
    print(f"Error loading models: {str(e)}")
    raise e

# Summarization endpoint
@app.post("/summarize")
async def summarize(request: Request):
    try:
        data = await request.json()
        text = data["text"]
        input_length = len(summarizer.tokenizer(text)["input_ids"])
        max_length = min(200, max(70, int(input_length * 0.7)))  
        min_length = max(30, int(input_length * 0.3))
        # Generate summary
        summary = summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)
        return {"summary": summary[0]["summary_text"]}
    except Exception as e:
        print(f"Error in summarize: {str(e)}")
        print(traceback.format_exc())
        return {"error": f"Failed to summarize: {str(e)}"}

# Translation endpoint
@app.post("/translate")
async def translate(request: Request):
    try:
        data = await request.json()
        text = data["text"]
        lang = data["lang"]
        if lang not in translators:
            return {"error": "Language not supported"}
        
        # Perform translation
        result = translators[lang](text)
        print(f"Translation result for {lang}: {result}")  
        
        # Check if result is a list and has at least one item
        if not isinstance(result, list) or len(result) == 0:
            return {"error": "Translation failed: empty or invalid result"}
        
        translation = result[0].get("translation_text")
        if translation is None:
            return {"error": "Translation failed: 'translation_text' not found in result"}
        
        return {"translation": translation}
    except Exception as e:
        print(f"Error in translate: {str(e)}")
        print(traceback.format_exc())
        return {"error": f"Failed to translate: {str(e)}"}

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