routers/powerpoint.py

from fastapi import APIRouter, Depends, HTTPException, File, UploadFile, Form
from fastapi.responses import FileResponse
from pydantic import BaseModel
from typing import Annotated, List, Literal, Optional, Union, Dict
from google import genai
from google.genai import types
from auth import verify_token
import os
import tempfile
from pptx import Presentation
from pptx.util import Inches, Pt
from pptx.enum.text import PP_ALIGN, MSO_ANCHOR
from pptx.enum.shapes import MSO_CONNECTOR
from pptx.chart.data import CategoryChartData
from pptx.enum.chart import XL_CHART_TYPE
from pptx.dml.color import RGBColor
from enum import Enum
from datetime import datetime, timedelta
from collections import deque
from threading import Lock
import json
from abc import ABC, abstractmethod

router = APIRouter(prefix="/powerpoint", tags=["powerpoint"])

# Get Gemini client
gemini_key = os.environ.get('GEMINI_KEY', '')
client = genai.Client(api_key=gemini_key)

class APIResponse(BaseModel):
    success: bool
    data: Optional[Union[Dict, List[Dict], str]] = None
    error: Optional[str] = None

class ChartDataPoint(BaseModel):
    label: str
    value: float

class MultimodalModel(str, Enum):
    gemini_25_flash = "gemini-2.5-flash"
    gemini_25_flash_lite = "gemini-2.5-flash-lite"
    gemini_20_flash = "gemini-2.0-flash"
    gemini_20_flash_lite = "gemini-2.0-flash-lite"
    gemini_25_pro = "gemini-2.5-pro"

class TemplateType(str, Enum):
    single_graph = "single_graph"
    two_column_graph = "two_column_graph"

# Base class for all presentation data models
class BasePresentationData(BaseModel, ABC):
    title: str

# Single graph template (original)
class SingleGraphPresentationData(BasePresentationData):
    main_text: str
    graph_title: str
    chart_type: Literal["bar", "column", "line", "pie", "area"]
    chart_data: List[ChartDataPoint]
    side_text: str

# Two column graph template (new)
class TwoColumnGraphPresentationData(BasePresentationData):
    left_text: str
    left_graph_title: str
    left_chart_type: Literal["bar", "column", "line", "pie", "area"]
    left_chart_data: List[ChartDataPoint]
    left_side_text: str
    right_text: str
    right_graph_title: str
    right_chart_type: Literal["bar", "column", "line", "pie", "area"]
    right_chart_data: List[ChartDataPoint]
    right_side_text: str

MULTIMODAL_MODEL_LIMITS = {
    "gemini-2.5-flash": {
        "rpm": 10,
        "daily": 250
    },
    "gemini-2.5-flash-lite": {
        "rpm": 30,
        "daily": 1000,
    },
    "gemini-2.0-flash": {
        "rpm": 15,
        "daily": 200
    },
    "gemini-2.0-flash-lite": {
        "rpm": 30,
        "daily": 200,
    },
}

# Track usage per model
model_usage = {
    model: {
        "daily_count": 0,
        "last_daily_reset": datetime.utcnow(),
        "rpm_timestamps": deque(),  # Tracks timestamps of requests in the past 60 seconds
    }
    for model in MULTIMODAL_MODEL_LIMITS
}

usage_lock = Lock()  # Thread-safe tracking

def get_available_model():
    now = datetime.utcnow()

    with usage_lock:
        for model, limits in MULTIMODAL_MODEL_LIMITS.items():
            usage = model_usage[model]

            # --- Daily Limit Reset ---
            if now - usage["last_daily_reset"] > timedelta(days=1):
                usage["daily_count"] = 0
                usage["last_daily_reset"] = now

            # --- RPM Cleanup ---
            rpm_window = timedelta(seconds=60)
            while usage["rpm_timestamps"] and now - usage["rpm_timestamps"][0] > rpm_window:
                usage["rpm_timestamps"].popleft()

            # --- Limit Checks ---
            if usage["daily_count"] < limits["daily"] and len(usage["rpm_timestamps"]) < limits["rpm"]:
                usage["daily_count"] += 1
                usage["rpm_timestamps"].append(now)
                return model

    raise HTTPException(
        status_code=429,
        detail="All Gemini multimodal models have reached their rate limits (daily or RPM)."
    )

def create_chart_from_data(slide, chart_data, chart_type, x, y, cx, cy):
    """Create a chart in the slide based on the provided data and type"""
    
    # Chart type mapping
    chart_type_mapping = {
        "bar": XL_CHART_TYPE.BAR_CLUSTERED,
        "column": XL_CHART_TYPE.COLUMN_CLUSTERED,
        "line": XL_CHART_TYPE.LINE,
        "pie": XL_CHART_TYPE.PIE,
        "area": XL_CHART_TYPE.AREA
    }
    
    # Create chart data
    chart_data_obj = CategoryChartData()
    chart_data_obj.categories = [item.label for item in chart_data]
    chart_data_obj.add_series('Series 1', [item.value for item in chart_data])
    
    # Add chart to slide
    chart = slide.shapes.add_chart(
        chart_type_mapping.get(chart_type, XL_CHART_TYPE.COLUMN_CLUSTERED),
        x, y, cx, cy, chart_data_obj
    ).chart
    
    chart.chart_title.text_frame.text = ""
    
    return chart

# Abstract base class for template generators
class TemplateGenerator(ABC):
    @abstractmethod
    def generate_pptx(self, presentation_data: BasePresentationData) -> str:
        pass
    
    @abstractmethod
    def get_prompt(self) -> str:
        pass
    
    @abstractmethod
    def get_response_schema(self):
        pass

class SingleGraphTemplateGenerator(TemplateGenerator):
    def get_prompt(self) -> str:
        return """
        Analyze this image and create a professional presentation slide based on what you see. Don't modify or invent text.
        
        Please provide:
        1. A title from the image
        2. Secondary title from the image (main_text)
        3. A graph title that describes any data visualization or key metric
        4. Chart data with meaningful labels and values (create realistic data if needed based on the image)
        5. Choose the most appropriate chart type (bar, column, line, pie, or area)
        6. Side text with bullet points •, dash - and \\n
        
        Make sure the content is professional, data-driven, and suitable for a business presentation.
        The chart data should have 3-7 data points with meaningful labels and realistic values.
        """
    
    def get_response_schema(self):
        return SingleGraphPresentationData
    
    def generate_pptx(self, presentation_data: SingleGraphPresentationData) -> str:
        """Generate PPTX file from single graph presentation data"""
        
        prs = Presentation()
        prs.slide_width = Inches(13.33)  # standard 16:9
        prs.slide_height = Inches(7.5)
        slide = prs.slides.add_slide(prs.slide_layouts[6])
        
        # Shape 1: Title
        title_box = slide.shapes.add_textbox(Inches(0.39), Inches(0.24), Inches(12.53), Inches(0.9))
        title_tf = title_box.text_frame
        title_tf.word_wrap = True
        title_tf.text = presentation_data.title
        title_tf.paragraphs[0].font.size = Pt(24)
        title_tf.paragraphs[0].font.name = "Verdana"
        
        # Shape 2: Sources
        src_box = slide.shapes.add_textbox(Inches(0.38), Inches(6.95), Inches(12.53), Inches(0.19))
        src_tf = src_box.text_frame
        src_tf.text = "Sources: TBD manually"
        src_tf.paragraphs[0].font.size = Pt(10)
        
        # Shape 4: Ligne horizontale sous le titre
        line1 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(0.39), Inches(1.23),
            Inches(12.92), Inches(1.23)
        )
        line1.line.width = Pt(0.5)
        line1.shadow.inherit = False
        line1.line.fill.solid()
        line1.line.fill.fore_color.rgb = RGBColor(1, 1, 1)
        
        # Shape 7: [Text] gauche
        text_box = slide.shapes.add_textbox(Inches(0.37), Inches(1.66), Inches(9.61), Inches(0.35))
        text_tf = text_box.text_frame
        text_tf.word_wrap = True
        text_tf.text = presentation_data.main_text
        text_tf.paragraphs[0].font.size = Pt(16)
        text_tf.paragraphs[0].font.bold = True
        text_tf.paragraphs[0].font.name = "Verdana"
        
        # Shape 6: Ligne horizontale sous [Text]
        line2 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(0.37), Inches(2.09),
            Inches(9.99), Inches(2.09)
        )
        line2.line.width = Pt(6)
        line2.shadow.inherit = False
        line2.line.fill.solid()
        line2.line.fill.fore_color.rgb = RGBColor(0, 32, 96)
        
        # Shape 5: Graph title
        graph_title = slide.shapes.add_textbox(Inches(0.37), Inches(2.16), Inches(9.23), Inches(0.48))
        gtf = graph_title.text_frame
        gtf.word_wrap = True
        gtf.text = presentation_data.graph_title
        gtf.paragraphs[0].font.size = Pt(14)
        gtf.paragraphs[0].font.name = "Verdana"
        
        # Shape 10: Create actual chart
        try:
            create_chart_from_data(slide, presentation_data.chart_data, 
                                  presentation_data.chart_type, 
                                  Inches(0.39), Inches(2.96), Inches(9.23), Inches(3.64))
        except Exception as e:
            # Fallback to text box if chart creation fails
            graph_box = slide.shapes.add_textbox(Inches(0.39), Inches(2.96), Inches(9.23), Inches(3.64))
            graph_tf = graph_box.text_frame
            graph_tf.text = f"Chart ({presentation_data.chart_type}): {presentation_data.graph_title}"
            graph_tf.paragraphs[0].alignment = PP_ALIGN.CENTER
            graph_tf.paragraphs[0].font.italic = True
        
        # Shape 8: Ligne verticale séparatrice
        line3 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(10.0), Inches(2.23),
            Inches(10.0), Inches(6.6)
        )
        line3.line.width = Pt(1.5)
        line3.shadow.inherit = False
        line3.line.fill.solid()
        line3.line.fill.fore_color.rgb = RGBColor(0, 32, 96)
        
        # Shape 11: Texte colonne droite
        side_box = slide.shapes.add_textbox(Inches(10.19), Inches(2.2), Inches(2.72), Inches(4.39))
        fill = side_box.fill
        fill.solid()
        fill.fore_color.rgb = RGBColor(247, 248, 250)
        side_tf = side_box.text_frame
        side_tf.word_wrap = True
        side_tf.text = presentation_data.side_text
        for paragraph in side_tf.paragraphs:
            paragraph.font.size = Pt(12)
            paragraph.font.name = "Verdana"
        side_tf.vertical_anchor = MSO_ANCHOR.MIDDLE
        
        # Save to temporary file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.pptx')
        prs.save(temp_file.name)
        
        return temp_file.name

class TwoColumnGraphTemplateGenerator(TemplateGenerator):
    def get_prompt(self) -> str:
        return """
        Analyze this image and create a professional presentation slide with two columns based on what you see. Don't modify or invent text.
        
        Please provide:
        1. A title from the image
        2. Left column text (main insight or finding)
        3. Left graph title that describes the first data visualization
        4. Left chart data with meaningful labels and values (3-5 data points)
        5. Choose appropriate chart type for left chart (bar, column, line, pie, or area)
        6. Left side text with bullet points •, dash - and \\n
        7. Right column text (secondary insight or comparison)
        8. Right graph title that describes the second data visualization
        9. Right chart data with meaningful labels and values (3-5 data points)
        10. Choose appropriate chart type for right chart (bar, column, line, pie, or area)
        11. Right side text with bullet points •, dash - and \\n
        
        Make sure the content is professional, data-driven, and suitable for a business presentation.
        The two columns should complement each other and tell a coherent story.
        """
    
    def get_response_schema(self):
        return TwoColumnGraphPresentationData
    
    def generate_pptx(self, presentation_data: TwoColumnGraphPresentationData) -> str:
        """Generate PPTX file from two column graph presentation data"""
        
        prs = Presentation()
        prs.slide_width = Inches(13.33)  # 16:9 widescreen
        prs.slide_height = Inches(7.5)
        slide = prs.slides.add_slide(prs.slide_layouts[6])  # blank layout
        
        # === Shape 1 : Title ===
        title_box = slide.shapes.add_textbox(Inches(0.39), Inches(0.24), Inches(12.53), Inches(0.9))
        title_tf = title_box.text_frame
        title_tf.text = presentation_data.title
        title_tf.word_wrap = True
        title_tf.paragraphs[0].font.name = "Verdana"
        title_tf.paragraphs[0].font.size = Pt(22)
        title_tf.vertical_anchor = MSO_ANCHOR.MIDDLE
        
        # === Shape 2 : Sources ===
        src_box = slide.shapes.add_textbox(Inches(0.38), Inches(6.91), Inches(12.53), Inches(0.23))
        src_tf = src_box.text_frame
        src_tf.text = "Sources: TBD manually"
        src_tf.paragraphs[0].font.size = Pt(10)
        src_tf.paragraphs[0].font.name = "Verdana"
        
        # === Shape 3 : Placeholder (page num / legend) ===
        ph_box = slide.shapes.add_textbox(Inches(10.15), Inches(7.04), Inches(3.0), Inches(0.4))
        ph_tf = ph_box.text_frame
        ph_tf.text = "1"
        ph_tf.paragraphs[0].alignment = PP_ALIGN.RIGHT
        ph_tf.paragraphs[0].font.size = Pt(10)
        
        # Shape 4: Ligne horizontale sous le titre
        line1 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(0.39), Inches(1.23),
            Inches(12.92), Inches(1.23)
        )
        line1.line.width = Pt(0.5)
        line1.shadow.inherit = False
        line1.line.fill.solid()
        line1.line.fill.fore_color.rgb = RGBColor(1, 1, 1)
        
        # === Shape 5 : Ligne horizontale gauche ===
        line2 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(0.37), Inches(2.09),
            Inches(6.3), Inches(2.09)
        )
        line2.line.width = Pt(6)
        line2.shadow.inherit = False
        line2.line.fill.solid()
        line2.line.fill.fore_color.rgb = RGBColor(0, 32, 96)
        
        # === Shape 6 : Graph title gauche ===
        gt_left = slide.shapes.add_textbox(Inches(0.38), Inches(2.21), Inches(5.9), Inches(0.39))
        gtl_tf = gt_left.text_frame
        gtl_tf.word_wrap = True
        gtl_tf.text = presentation_data.left_graph_title
        gtl_tf.paragraphs[0].font.size = Pt(12)
        gtl_tf.paragraphs[0].font.name = "Verdana"
        
        # === Shape 7 : Ligne horizontale droite ===
        line3 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(7.07), Inches(2.09),
            Inches(12.97), Inches(2.09)
        )
        line3.line.width = Pt(6)
        line3.shadow.inherit = False
        line3.line.fill.solid()
        line3.line.fill.fore_color.rgb = RGBColor(0, 32, 96)
        
        # === Shape 8 : Graph title droite ===
        gt_right = slide.shapes.add_textbox(Inches(7.07), Inches(2.21), Inches(5.9), Inches(0.39))
        gtr_tf = gt_right.text_frame
        gtr_tf.word_wrap = True
        gtr_tf.text = presentation_data.right_graph_title
        gtr_tf.paragraphs[0].font.size = Pt(12)
        gtr_tf.paragraphs[0].font.name = "Verdana"
        
        # === Shape 9 : Bloc texte gauche ===
        text_left = slide.shapes.add_textbox(Inches(4.47), Inches(2.85), Inches(1.8), Inches(3.6))
        tl_tf_fill = text_left.fill
        tl_tf_fill.solid()
        tl_tf_fill.fore_color.rgb = RGBColor(247, 248, 250)
        tl_tf = text_left.text_frame
        tl_tf.word_wrap = True
        tl_tf.text = presentation_data.left_side_text
        for paragraph in tl_tf.paragraphs:
            paragraph.font.size = Pt(12)
            paragraph.font.name = "Verdana"
        tl_tf.vertical_anchor = MSO_ANCHOR.MIDDLE
        
        # === Shape 10 : Bloc texte droite ===
        text_right = slide.shapes.add_textbox(Inches(11.13), Inches(2.85), Inches(1.8), Inches(3.6))
        tr_tf_fill = text_right.fill
        tr_tf_fill.solid()
        tr_tf_fill.fore_color.rgb = RGBColor(247, 248, 250)
        tr_tf = text_right.text_frame
        tr_tf.word_wrap = True
        tr_tf.text = presentation_data.right_side_text
        for paragraph in tr_tf.paragraphs:
            paragraph.font.size = Pt(12)
            paragraph.font.name = "Verdana"
        tr_tf.vertical_anchor = MSO_ANCHOR.MIDDLE
        
        # === Shape 11 : Ligne verticale séparation ===
        line4 = slide.shapes.add_connector(
            MSO_CONNECTOR.STRAIGHT,
            Inches(6.69), Inches(2.23),
            Inches(6.69), Inches(6.32)
        )
        line4.line.width = Pt(0.25)
        line4.shadow.inherit = False
        line4.line.fill.solid()
        line4.line.fill.fore_color.rgb = RGBColor(217, 217, 217)
        
        # === Shape 12 : Texte haut gauche ===
        txt12 = slide.shapes.add_textbox(Inches(0.38), Inches(1.52), Inches(5.9), Inches(0.44))
        txt12_tf = txt12.text_frame
        txt12_tf.word_wrap = True
        txt12_tf.text = presentation_data.left_text
        txt12_tf.paragraphs[0].font.size = Pt(14)
        txt12_tf.paragraphs[0].font.name = "Verdana"
        txt12_tf.paragraphs[0].font.bold = True
        txt12_tf.paragraphs[0].font.color.rgb = RGBColor(0, 32, 96)
        
        # === Shape 13 : Texte haut droite ===
        txt13 = slide.shapes.add_textbox(Inches(7.07), Inches(1.52), Inches(5.9), Inches(0.44))
        txt13_tf = txt13.text_frame
        txt13_tf.word_wrap = True
        txt13_tf.text = presentation_data.right_text
        txt13_tf.paragraphs[0].font.size = Pt(14)
        txt13_tf.paragraphs[0].font.name = "Verdana"
        txt13_tf.paragraphs[0].font.bold = True
        txt13_tf.paragraphs[0].font.color.rgb = RGBColor(0, 32, 96)
        
        # === Shape 14 : Graph gauche ===
        try:
            create_chart_from_data(slide, presentation_data.left_chart_data, 
                                  presentation_data.left_chart_type,
                                  Inches(0.38), Inches(2.85), Inches(4.1), Inches(3.6))
        except Exception as e:
            graph_left = slide.shapes.add_shape(1, Inches(0.38), Inches(2.85), Inches(4.1), Inches(3.6))
            gl_tf = graph_left.text_frame
            gl_tf.text = f"Graph ({presentation_data.left_chart_type})"
            gl_tf.paragraphs[0].alignment = PP_ALIGN.CENTER
            gl_tf.paragraphs[0].font.italic = True
        
        # === Shape 15 : Graph droite ===
        try:
            create_chart_from_data(slide, presentation_data.right_chart_data, 
                                  presentation_data.right_chart_type,
                                  Inches(7.07), Inches(2.85), Inches(4.06), Inches(3.6))
        except Exception as e:
            graph_right = slide.shapes.add_shape(1, Inches(7.07), Inches(2.85), Inches(4.06), Inches(3.6))
            gr_tf = graph_right.text_frame
            gr_tf.text = f"Graph ({presentation_data.right_chart_type})"
            gr_tf.paragraphs[0].alignment = PP_ALIGN.CENTER
            gr_tf.paragraphs[0].font.italic = True
        
        # Save to temporary file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.pptx')
        prs.save(temp_file.name)
        
        return temp_file.name

# Template registry
TEMPLATE_GENERATORS = {
    TemplateType.single_graph: SingleGraphTemplateGenerator(),
    TemplateType.two_column_graph: TwoColumnGraphTemplateGenerator()
}

@router.post("/image-to-pptx")
async def image_to_pptx(
    token: Annotated[str, Depends(verify_token)],
    file: UploadFile = File(..., description="Image file to convert to PowerPoint"),
    template: TemplateType = Form(..., description="Template type to use for the presentation"),
    model: MultimodalModel = Form(..., description="Gemini model (Gemini 2.5 Flash recommended"),
):
    """
    Analyze an image with Gemini and generate a PowerPoint presentation based on the analysis using the selected template
    """
    try:
        # Get available model
        #model = get_available_model()
        
        # Get template generator
        template_generator = TEMPLATE_GENERATORS.get(template)
        if not template_generator:
            raise HTTPException(status_code=400, detail=f"Unsupported template: {template}")
        
        # Read uploaded file
        file_content = await file.read()
        
        # Send to Gemini
        response = client.models.generate_content(
            model=model,
            contents=[
                types.Part.from_bytes(
                    data=file_content,
                    mime_type=file.content_type,
                ),
                template_generator.get_prompt()
            ],
            config={
                'response_mime_type': 'application/json',
                'response_schema': template_generator.get_response_schema(),
            }
        )
        # Display input/output tokens
        print(response.usage_metadata)
        
        # Parse the response
        presentation_data = template_generator.get_response_schema().model_validate(json.loads(response.text))
        
        # Generate PPTX file
        pptx_path = template_generator.generate_pptx(presentation_data)
        
        # Return file as download
        return FileResponse(
            path=pptx_path,
            filename=f"presentation_{template.value}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pptx",
            media_type="application/vnd.openxmlformats-officedocument.presentationml.presentation"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        return APIResponse(success=False, error=f"Failed to generate presentation: {str(e)}")

@router.get("/templates")
async def get_available_templates(token: Annotated[str, Depends(verify_token)]):
    """Get list of available presentation templates"""
    templates = []
    for template_type in TemplateType:
        templates.append({
            "name": template_type.value,
            "description": {
                TemplateType.single_graph: "Single graph with sidebar - ideal for focusing on one key insight",
                TemplateType.two_column_graph: "Two column layout with graphs and sidebars - perfect for comparisons"
            }.get(template_type, "")
        })
    
    return APIResponse(success=True, data=templates)

@router.get("/model-stats")
async def get_model_stats(token: Annotated[str, Depends(verify_token)]):
    """Get current model usage statistics"""
    now = datetime.utcnow()
    model_stats = []

    with usage_lock:
        for model, limits in MULTIMODAL_MODEL_LIMITS.items():
            usage = model_usage[model]

            # Reset daily count if 24h passed
            if now - usage["last_daily_reset"] > timedelta(days=1):
                usage["daily_count"] = 0
                usage["last_daily_reset"] = now

            # Clean up old RPM timestamps
            rpm_window = timedelta(seconds=60)
            while usage["rpm_timestamps"] and now - usage["rpm_timestamps"][0] > rpm_window:
                usage["rpm_timestamps"].popleft()

            model_stats.append({
                "model": model,
                "daily_count": usage["daily_count"],
                "daily_limit": limits["daily"],
                "rpm_count": len(usage["rpm_timestamps"]),
                "rpm_limit": limits["rpm"],
                "last_daily_reset": usage["last_daily_reset"].isoformat(),
            })

    return APIResponse(success=True, data=model_stats)