lambda_function.py

import boto3
import sqlite3
from PIL import Image
import io
import generator
import json
import base64
import logging
import os
import time
from boto3.dynamodb.conditions import Key

logger = logging.getLogger()
logger.setLevel(logging.INFO)

CORS_HEADERS = {
    'Access-Control-Allow-Origin': '*',
    'Access-Control-Allow-Headers': 'Content-Type',
    'Access-Control-Allow-Methods': 'POST,OPTIONS'
}

DB_CACHE_TTL = 300  # 5 minutes


def make_response(status_code, body):
    return {
        'statusCode': status_code,
        'headers': CORS_HEADERS,
        'body': json.dumps(body)
    }

def get_random_question(db_path, topic=None, subject_id=None):
    """Get a random question from the database."""
    try:
        with sqlite3.connect(db_path) as connection:
            cursor = connection.cursor()

            table_name = "QuestionsAdd" if subject_id == "0606" else "Questions"

            # Check table exists
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
            if not cursor.fetchone():
                logger.error(f"{table_name} table not found in database")
                return None

            # Build WHERE clause
            where_clause = []
            params = []

            if topic:
                cursor.execute(f"SELECT DISTINCT Topic FROM {table_name}")
                available_topics = [row[0] for row in cursor.fetchall()]

                matching_topic = next(
                    (db_topic for db_topic in available_topics
                     if db_topic and topic and db_topic.lower() == topic.lower()),
                    None
                )

                if matching_topic:
                    where_clause.append("Topic = ?")
                    params.append(matching_topic)
                else:
                    logger.warning(f"No matching topic found for '{topic}' in {table_name}")

            where_sql = " AND ".join(where_clause) if where_clause else "1=1"

            # Get a random question
            question_sql = f"""
                SELECT Image, Topic, Marks, Year, Month, Variant, Paper
                FROM {table_name}
                WHERE {where_sql}
                ORDER BY RANDOM()
                LIMIT 1
            """
            cursor.execute(question_sql, params)

            result = cursor.fetchone()
            if result:
                image_base64 = base64.b64encode(result[0]).decode('utf-8') if result[0] else None
                return {
                    'image': image_base64,
                    'topic': result[1],
                    'marks': result[2],
                    'year': result[3],
                    'month': result[4],
                    'variant': result[5],
                    'paper': result[6]
                }
            logger.error(f"No questions found in {table_name} for topic: {topic or 'all'}")
            return None

    except sqlite3.Error as e:
        logger.error(f"Database error: {str(e)}")
        return None
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        return None

def lambda_handler(event, context):
    try:
        s3 = boto3.client('s3')
        ddb = boto3.client('dynamodb')
        counters_table = os.environ.get('COUNTERS_TABLE', 'AppCounters')

        def increment_papers_total():
            ddb.update_item(
                TableName=counters_table,
                Key={'pk': {'S': 'papers_total'}},
                UpdateExpression='ADD #c :inc',
                ExpressionAttributeNames={'#c': 'count'},
                ExpressionAttributeValues={':inc': {'N': '1'}},
            )

        def get_papers_total():
            resp = ddb.get_item(
                TableName=counters_table,
                Key={'pk': {'S': 'papers_total'}},
                ConsistentRead=True
            )
            item = resp.get('Item')
            if not item or 'count' not in item:
                return 0
            return int(item['count']['N'])

        # Cache DB download — reuse if file exists and is recent
        s3_bucket = 'shahu-exam'
        s3_key = 'examcopy.db'
        local_file_path = '/tmp/my-database.db'

        should_download = True
        if os.path.exists(local_file_path):
            file_age = time.time() - os.path.getmtime(local_file_path)
            if file_age < DB_CACHE_TTL:
                should_download = False

        if should_download:
            s3.download_file(s3_bucket, s3_key, local_file_path)
            logger.info("Downloaded fresh database from S3")

        # Random question request
        if event.get('action') == 'random_question':
            question_data = get_random_question(
                local_file_path,
                event.get('topic'),
                event.get('subjectId')
            )
            if question_data:
                return make_response(200, question_data)
            return make_response(404, {'error': 'No questions found'})

        # Stats request
        if event.get('action') == 'stats':
            total = get_papers_total()
            return make_response(200, {'totalPapersGenerated': total})

        # Paper generation
        logger.info("Processing paper generation request")
        raw_paper = generator.exam(event['topic'], event['mark'], event['subjectId'], event['info'])
        if raw_paper == "Error":
            logger.error("Failed to generate paper")
            return make_response(500, {'error': 'Failed to generate paper'})

        first_page = raw_paper[0]
        pdf_filename = "/tmp/shahu.pdf"
        first_page.save(pdf_filename, "PDF", resolution=100.0, save_all=True, append_images=raw_paper[1:])
        object_name = "new-exam.pdf"
        s3.upload_file(pdf_filename, s3_bucket, object_name)

        # Clean up temp PDF
        try:
            os.remove(pdf_filename)
        except OSError:
            pass

        increment_papers_total()

        url = s3.generate_presigned_url(
            'get_object',
            Params={'Bucket': s3_bucket, 'Key': object_name},
            ExpiresIn=3600
        )

        logger.info("Successfully generated paper and presigned URL")
        return make_response(200, {'url': url})

    except Exception as e:
        logger.error(f"Error in lambda_handler: {str(e)}")
        return make_response(500, {'error': str(e)})