memories/backend/main.py

from fastapi import FastAPI, UploadFile, File, Depends, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from sqlalchemy.orm import Session
from geoalchemy2.elements import WKTElement
import os
import datetime
from datetime import datetime as dt
import json
import redis
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib import colors
from io import BytesIO

from data_importer import GoogleTakeoutImporter
from analysis import LocationAnalyzer
from models import SessionLocal, engine, Base, Location, Visit, Trip, create_db_tables

app = FastAPI()

# Redis connection
redis_client = redis.from_url(os.getenv("REDIS_URL", "redis://localhost:6379"))

# Dependency to get the DB session
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

@app.on_event("startup")
async def startup_event():
    create_db_tables()

@app.post("/upload")
async def upload_file(file: UploadFile = File(...), db: Session = Depends(get_db)):
    if not file.filename:
        raise HTTPException(status_code=400, detail="No selected file")

    file_location = f"/tmp/{file.filename}"
    try:
        with open(file_location, "wb+") as file_object:
            file_object.write(await file.read())

        importer = GoogleTakeoutImporter()
        raw_locations = importer.import_locations(file_location)

        new_locations = []
        for loc_data in raw_locations:
            if all(k in loc_data for k in ['latitude', 'longitude', 'timestamp']):
                if isinstance(loc_data['timestamp'], str):
                    timestamp = dt.fromisoformat(loc_data['timestamp'].replace('Z', '+00:00'))
                else:
                    timestamp = loc_data['timestamp']
                
                point = WKTElement(f'POINT({loc_data["longitude"]} {loc_data["latitude"]})', srid=4326)
                new_location = Location(
                    latitude=loc_data['latitude'],
                    longitude=loc_data['longitude'],
                    timestamp=timestamp,
                    geom=point
                )
                new_locations.append(new_location)

        analyzer = LocationAnalyzer()
        visits = analyzer.detect_visits(new_locations)
        trips = analyzer.detect_trips(new_locations, visits)

        db.add_all(new_locations)
        db.add_all(visits)
        db.add_all(trips)
        db.commit()

        # Invalidate cache
        redis_client.flushdb()

        return JSONResponse(content={
            'message': f'Successfully imported {len(new_locations)} locations, detected {len(visits)} visits and {len(trips)} trips.'
        }, status_code=200)
    except Exception as e:
        if 'db' in locals():
            db.rollback()
        raise HTTPException(status_code=500, detail=f"Failed to process file: {str(e)}")
    finally:
        if os.path.exists(file_location):
            os.remove(file_location)

@app.get("/data/locations")
def get_locations(db: Session = Depends(get_db)):
    cache_key = "locations"
    cached = redis_client.get(cache_key)
    if cached:
        return JSONResponse(content=json.loads(cached))
    
    locations = db.query(Location).all()
    data = {"locations": [{"id": l.id, "lat": l.latitude, "lon": l.longitude, "time": l.timestamp.isoformat()} for l in locations]}
    
    redis_client.setex(cache_key, 300, json.dumps(data))  # Cache for 5 min
    return data

@app.get("/data/visits")
def get_visits(db: Session = Depends(get_db)):
    cache_key = "visits"
    cached = redis_client.get(cache_key)
    if cached:
        return JSONResponse(content=json.loads(cached))
    
    visits = db.query(Visit).all()
    data = {"visits": [{"id": v.id, "start": v.start_time.isoformat(), "end": v.end_time.isoformat(), 
                       "centroid": str(v.centroid_geom) if v.centroid_geom else None} for v in visits]}
    
    redis_client.setex(cache_key, 300, json.dumps(data))
    return data

@app.get("/data/trips")
def get_trips(db: Session = Depends(get_db)):
    cache_key = "trips"
    cached = redis_client.get(cache_key)
    if cached:
        return JSONResponse(content=json.loads(cached))
    
    trips = db.query(Trip).all()
    data = {"trips": [{"id": t.id, "start": t.start_time.isoformat(), "end": t.end_time.isoformat(), 
                      "mode": t.transport_mode, "path": str(t.path_geom) if t.path_geom else None} for t in trips]}
    
    redis_client.setex(cache_key, 300, json.dumps(data))
    return data

@app.get("/export/report")
def export_report(db: Session = Depends(get_db)):
    # Fetch data
    visits = db.query(Visit).count()
    trips = db.query(Trip).count()
    locations = db.query(Location).count()

    buffer = BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter)
    styles = getSampleStyleSheet()
    story = []

    # Title
    title = Paragraph("Travel Memories Report", styles['Title'])
    story.append(title)
    story.append(Spacer(1, 12))

    # Stats
    stats_data = [
        ['Metric', 'Count'],
        ['Total Locations', locations],
        ['Detected Visits', visits],
        ['Detected Trips', trips]
    ]
    stats_table = Table(stats_data)
    stats_table.setStyle([
        ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
        ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, 0), 14),
        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
        ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
        ('GRID', (0, 0), (-1, -1), 1, colors.black)
    ])
    story.append(stats_table)

    # Build PDF
    doc.build(story)
    buffer.seek(0)

    return StreamingResponse(buffer, media_type="application/pdf", headers={"Content-Disposition": "attachment; filename=travel-report.pdf"})