OCR/backend/main.py

import io
import datetime
from typing import Optional, List
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from pypdf import PdfReader
import pytesseract
from PIL import Image
from pdf2image import convert_from_bytes
import pdfplumber
from sqlalchemy.orm import Session



# 2. OCR Module
def extract_text_from_pdf(file_bytes: bytes) -> str:
    # Try pdfplumber first for layout preservation
    try:
        with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
            text = ""
            for page in pdf.pages:
                page_text = page.extract_text(layout=True)
                if page_text:
                    text += page_text + "\n"
            if text.strip():
                return text.strip()
    except Exception as e:
        print(f"pdfplumber error: {e}")

    # Fallback to pypdf
    try:
        reader = PdfReader(io.BytesIO(file_bytes))
        text = ""
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
        return text.strip()
    except Exception as e:
        print(f"Error reading PDF: {e}")
        return ""

# Internal modules
from database import get_db, Email
from scheduler import start_scheduler, stop_scheduler
from mail_service import fetch_and_store_emails

# Lifespan for Scheduler
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    start_scheduler()
    yield
    # Shutdown
    stop_scheduler()

app = FastAPI(lifespan=lifespan)

# CORS configuration
origins = [
    "http://localhost",
    "http://localhost:3000",
    "http://localhost:5173", # Vite default
    "http://localhost:4200", # Angular default
    "http://127.0.0.1:4200",
]

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

# Models
class LoginRequest(BaseModel):
    username: str
    password: str

class LoginResponse(BaseModel):
    token: str
    message: str

@app.post("/api/login", response_model=LoginResponse)
def login(request: LoginRequest):
    if request.username == "admin" and request.password == "admin":
        return LoginResponse(token="fake-super-secret-token", message="Success")
    raise HTTPException(status_code=401, detail="Invalid username or password")

class NERResponse(BaseModel):
    text: str



def extract_text_from_image(file_bytes: bytes) -> str:
    try:
        image = Image.open(io.BytesIO(file_bytes))
        text = pytesseract.image_to_string(image)
        return text.strip()
    except Exception as e:
        print(f"Error reading Image: {e}")
        return ""

@app.post("/api/ocr/extract", response_model=NERResponse)
async def extract_text(file: UploadFile = File(...)):
    content = await file.read()
    filename = file.filename.lower()
    
    extracted_text = ""
    
    if filename.endswith(".pdf"):
        # Try text extraction first
        extracted_text = extract_text_from_pdf(content)
        
        # If text is empty, it might be a scanned PDF. 
        if not extracted_text.strip():
             try:
                 images = convert_from_bytes(content)
                 for i, image in enumerate(images):
                     page_text = pytesseract.image_to_string(image)
                     extracted_text += f"\n--- Page {i+1} ---\n{page_text}"
             except Exception as e:
                 extracted_text = f"Error processing Scanned PDF: {str(e)}\n\n(Hint: Ensure 'poppler' is installed on your system. Run 'brew install poppler')"

    elif filename.endswith((".png", ".jpg", ".jpeg", ".tiff", ".bmp")):
        extracted_text = extract_text_from_image(content)
    else:
        raise HTTPException(status_code=400, detail="Unsupported file type")
    
    return NERResponse(text=extracted_text)

import zipfile
import mimetypes
from fastapi.responses import StreamingResponse

class EmailDTO(BaseModel):
    id: int
    subject: Optional[str]
    sender: Optional[str]
    received_date: Optional[datetime.datetime]
    is_read: bool
    has_attachments: bool = False

class AttachmentDTO(BaseModel):
    id: int
    filename: str
    content_type: str

class EmailListResponse(BaseModel):
    items: List[EmailDTO]
    total: int

class EmailDetailDTO(EmailDTO):
    body: Optional[str]
    attachments: List[AttachmentDTO]

@app.get("/api/emails", response_model=EmailListResponse)
def get_emails(skip: int = 0, limit: int = 50, db: Session = Depends(get_db)):
    total = db.query(Email).count()
    emails = db.query(Email).order_by(Email.received_date.desc()).offset(skip).limit(limit).all()
    
    items = [
        EmailDTO(
            id=e.id, 
            subject=e.subject, 
            sender=e.sender, 
            received_date=e.received_date, 
            is_read=e.is_read,
            has_attachments=bool(e.attachments) # Check if list is not empty
        ) for e in emails
    ]
    return EmailListResponse(items=items, total=total)

# ...

@app.get("/api/emails/{email_id}", response_model=EmailDetailDTO)
def get_email_details(email_id: int, db: Session = Depends(get_db)):
    email = db.query(Email).filter(Email.id == email_id).first()
    if not email:
        raise HTTPException(status_code=404, detail="Email not found")
    
    # Mark as read if opened
    if not email.is_read:
        email.is_read = True
        db.commit()
    
    # Safely handle missing content_type and filename
    attachments = []
    for a in email.attachments:
        fname = a.filename
        if not fname or fname == "unknown":
            ext = mimetypes.guess_extension(a.content_type or "") or ".bin"
            fname = f"attachment_{a.id}{ext}"
            
        attachments.append(
            AttachmentDTO(
                id=a.id, 
                filename=fname, 
                content_type=a.content_type or "application/octet-stream"
            )
        )
        
    return EmailDetailDTO(
        id=email.id,
        subject=email.subject,
        sender=email.sender,
        received_date=email.received_date,
        is_read=email.is_read,
        has_attachments=bool(email.attachments),
        body=email.body,
        attachments=attachments
    )

@app.get("/api/emails/{email_id}/download-all")
def download_all_attachments(email_id: int, db: Session = Depends(get_db)):
    email = db.query(Email).filter(Email.id == email_id).first()
    if not email:
        raise HTTPException(status_code=404, detail="Email not found")
        
    if not email.attachments:
         raise HTTPException(status_code=400, detail="No attachments found")

    # Create a zip file in memory
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
        for att in email.attachments:
            # Use filename from DB, ensure unique names if duplicates exist (simple append here)
            zip_file.writestr(att.filename, att.file_content)
    
    zip_buffer.seek(0)
    
    return StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={"Content-Disposition": f"attachment; filename=email-{email_id}-attachments.zip"}
    )

@app.get("/api/attachments/{attachment_id}")
def download_attachment(attachment_id: int, view: bool = False, db: Session = Depends(get_db)):
    attachment = db.query(DBAttachment).filter(DBAttachment.id == attachment_id).first()
    if not attachment:
        raise HTTPException(status_code=404, detail="Attachment not found")
    
    disposition = "inline" if view else "attachment"
    
    return StreamingResponse(
        io.BytesIO(attachment.file_content), 
        media_type=attachment.content_type,
        headers={"Content-Disposition": f"{disposition}; filename={attachment.filename}"}
    )

@app.post("/api/emails/sync")
def sync_emails():
    result = fetch_and_store_emails()
    return result

@app.get("/")
def read_root():
    return {"message": "OCR Backend API is running"}