app.py

import os
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
from transformers import pipeline
from PIL import Image
import io
import fitz  # PyMuPDF
from docx import Document
import pandas as pd
import logging
from typing import Optional

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

app = FastAPI(title="AI Web Services")
app.mount("/static", StaticFiles(directory="static"), name="static")

# Initialize models (Spaces will cache these)
try:
    logger.info("Loading AI models...")
    image_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    text_pipeline = pipeline("text2text-generation", model="t5-small")
    logger.info("Models loaded successfully")
except Exception as e:
    logger.error(f"Model loading Failed: {e}")
    raise RuntimeError("Failed to initialize AI models")

@app.get("/", response_class=HTMLResponse)
async def home():
    """Serve the frontend interface"""
    try:
        with open("static/index.html") as f:
            return f.read()
    except Exception as e:
        logger.error(f"Failed to load frontend: {e}")
        raise HTTPException(500, "Frontend loading failed")

@app.post("/api/summarize")
async def summarize(
    file: Optional[UploadFile] = File(None),
    text: Optional[str] = Form(None)
):
    """
    Summarize text or document
    Accepts: PDF, DOCX or raw text
    Returns: {'summary': str}
    """
    try:
        if file:
            text = await extract_text(file)
        if not text:
            raise HTTPException(400, "No text provided")
        
        result = text_pipeline(f"summarize: {text}", max_length=150)
        return JSONResponse({"summary": result[0]['generated_text']})
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Summarization error: {e}")
        raise HTTPException(500, "Summarization failed")

@app.post("/api/caption")
async def caption_image(file: UploadFile = File(...)):
    """
    Generate caption for image
    Accepts: JPEG, PNG
    Returns: {'caption': str}
    """
    try:
        image = Image.open(io.BytesIO(await file.read()))
        result = image_pipeline(image)
        return JSONResponse({"caption": result[0]['generated_text']})
    except Exception as e:
        logger.error(f"Captioning error: {e}")
        raise HTTPException(500, "Image captioning failed")

@app.post("/api/answer")
async def answer_question(
    file: Optional[UploadFile] = File(None),
    text: Optional[str] = Form(None),
    question: str = Form(...)
):
    """
    Answer questions about text/document
    Accepts: PDF, DOCX or raw text + question
    Returns: {'answer': str}
    """
    try:
        if file:
            text = await extract_text(file)
        if not text:
            raise HTTPException(400, "No text provided")
        
        result = text_pipeline(f"question: {question} context: {text}")
        return JSONResponse({"answer": result[0]['generated_text']})
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"QA error: {e}")
        raise HTTPException(500, "Question answering failed")

@app.post("/api/visualize")
async def generate_visualization(
    file: UploadFile = File(...),
    chart_type: str = Form("bar")
):
    """
    Generate visualization code for Excel data
    Accepts: XLSX, CSV
    Returns: {'code': str, 'columns': list}
    """
    try:
        df = pd.read_excel(io.BytesIO(await file.read()))
        
        if chart_type.lower() == "bar":
            code = f"""import matplotlib.pyplot as plt
plt.bar(df['{df.columns[0]}'], df['{df.columns[1]}'])
plt.title('Bar Chart')
plt.show()"""
        else:
            code = f"""import seaborn as sns
sns.pairplot(df)
plt.title('Data Distribution')
plt.show()"""
            
        return JSONResponse({
            "code": code,
            "columns": list(df.columns)
        })
    except Exception as e:
        logger.error(f"Visualization error: {e}")
        raise HTTPException(500, "Visualization code generation failed")

async def extract_text(file: UploadFile) -> str:
    """Extract text from PDF or DOCX files"""
    try:
        content = await file.read()
        
        if file.filename.endswith(".pdf"):
            with fitz.open(stream=content, filetype="pdf") as doc:
                return " ".join(page.get_text() for page in doc)
        elif file.filename.endswith(".docx"):
            doc = Document(io.BytesIO(content))
            return "\n".join(p.text for p in doc.paragraphs)
        else:
            raise ValueError("Unsupported file format")
    except Exception as e:
        logger.error(f"Text extraction failed: {e}")
        raise HTTPException(400, f"Could not extract text: {e}")

# Health check endpoint
@app.get("/health")
async def health_check():
    return JSONResponse({"status": "healthy", "models": "loaded"})
    if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)