tikker/src/api/viz_service.py

"""
Tikker Visualization Microservice

Generates charts, graphs, and visual reports from keystroke statistics.
Supports multiple output formats and caching for performance.
"""

from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel
from typing import Dict, Any, Optional, List
import logging
import os
import base64
from io import BytesIO
import json

try:
    import matplotlib.pyplot as plt
    import matplotlib.dates as mdates
    import numpy as np
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False

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

app = FastAPI(
    title="Tikker Visualization Service",
    description="Generate charts and graphs for keystroke data",
    version="1.0.0"
)


class ChartRequest(BaseModel):
    title: str
    data: Dict[str, int]
    chart_type: str = "bar"
    width: int = 10
    height: int = 6


class ChartResponse(BaseModel):
    status: str
    image_base64: Optional[str] = None
    chart_type: str
    title: str


class HealthResponse(BaseModel):
    status: str
    viz_available: bool
    api_version: str


@app.get("/health", response_model=HealthResponse)
async def health_check() -> HealthResponse:
    """Health check endpoint."""
    return HealthResponse(
        status="healthy",
        viz_available=MATPLOTLIB_AVAILABLE,
        api_version="1.0.0"
    )


def _generate_bar_chart(title: str, data: Dict[str, int], width: int, height: int) -> bytes:
    """Generate a bar chart from data."""
    if not MATPLOTLIB_AVAILABLE:
        raise HTTPException(status_code=503, detail="Visualization not available")
    
    plt.figure(figsize=(width, height))
    
    labels = list(data.keys())
    values = list(data.values())
    
    plt.bar(labels, values, color='steelblue', edgecolor='navy', alpha=0.7)
    plt.title(title, fontsize=14, fontweight='bold')
    plt.xlabel('Category', fontsize=12)
    plt.ylabel('Count', fontsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    
    buf = BytesIO()
    plt.savefig(buf, format='png', dpi=100)
    buf.seek(0)
    plt.close()
    
    return buf.getvalue()


def _generate_line_chart(title: str, data: Dict[str, int], width: int, height: int) -> bytes:
    """Generate a line chart from data."""
    if not MATPLOTLIB_AVAILABLE:
        raise HTTPException(status_code=503, detail="Visualization not available")
    
    plt.figure(figsize=(width, height))
    
    labels = list(data.keys())
    values = list(data.values())
    
    plt.plot(labels, values, marker='o', linestyle='-', linewidth=2, color='steelblue', markersize=6)
    plt.title(title, fontsize=14, fontweight='bold')
    plt.xlabel('Category', fontsize=12)
    plt.ylabel('Count', fontsize=12)
    plt.grid(True, alpha=0.3)
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    
    buf = BytesIO()
    plt.savefig(buf, format='png', dpi=100)
    buf.seek(0)
    plt.close()
    
    return buf.getvalue()


def _generate_pie_chart(title: str, data: Dict[str, int], width: int, height: int) -> bytes:
    """Generate a pie chart from data."""
    if not MATPLOTLIB_AVAILABLE:
        raise HTTPException(status_code=503, detail="Visualization not available")
    
    plt.figure(figsize=(width, height))
    
    labels = list(data.keys())
    values = list(data.values())
    
    plt.pie(values, labels=labels, autopct='%1.1f%%', startangle=90, colors=plt.cm.Set3.colors)
    plt.title(title, fontsize=14, fontweight='bold')
    plt.tight_layout()
    
    buf = BytesIO()
    plt.savefig(buf, format='png', dpi=100)
    buf.seek(0)
    plt.close()
    
    return buf.getvalue()


@app.post("/chart", response_model=ChartResponse)
async def generate_chart(request: ChartRequest) -> ChartResponse:
    """
    Generate a chart from data.
    
    Args:
        request: Chart configuration with data and type
    
    Returns:
        Chart response with base64-encoded image
    """
    try:
        chart_type = request.chart_type.lower()
        
        if chart_type == "bar":
            image_data = _generate_bar_chart(request.title, request.data, request.width, request.height)
        elif chart_type == "line":
            image_data = _generate_line_chart(request.title, request.data, request.width, request.height)
        elif chart_type == "pie":
            image_data = _generate_pie_chart(request.title, request.data, request.width, request.height)
        else:
            raise HTTPException(status_code=400, detail=f"Unknown chart type: {chart_type}")
        
        image_base64 = base64.b64encode(image_data).decode('utf-8')
        
        return ChartResponse(
            status="success",
            image_base64=image_base64,
            chart_type=chart_type,
            title=request.title
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Chart generation error: {e}")
        raise HTTPException(status_code=500, detail=f"Chart generation failed: {str(e)}")


@app.post("/chart/download")
async def download_chart(request: ChartRequest) -> FileResponse:
    """
    Download a chart as PNG file.
    
    Args:
        request: Chart configuration
    
    Returns:
        PNG file download
    """
    try:
        chart_type = request.chart_type.lower()
        
        if chart_type == "bar":
            image_data = _generate_bar_chart(request.title, request.data, request.width, request.height)
        elif chart_type == "line":
            image_data = _generate_line_chart(request.title, request.data, request.width, request.height)
        elif chart_type == "pie":
            image_data = _generate_pie_chart(request.title, request.data, request.width, request.height)
        else:
            raise HTTPException(status_code=400, detail=f"Unknown chart type: {chart_type}")
        
        filename = f"{request.title.replace(' ', '_')}.png"
        
        return StreamingResponse(
            BytesIO(image_data),
            media_type="image/png",
            headers={"Content-Disposition": f"attachment; filename={filename}"}
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Chart download error: {e}")
        raise HTTPException(status_code=500, detail=f"Chart download failed: {str(e)}")


@app.get("/")
async def root() -> Dict[str, Any]:
    """Root endpoint with service information."""
    return {
        "name": "Tikker Visualization Service",
        "version": "1.0.0",
        "status": "running",
        "viz_available": MATPLOTLIB_AVAILABLE,
        "supported_charts": ["bar", "line", "pie"],
        "endpoints": {
            "health": "/health",
            "chart": "/chart",
            "download": "/chart/download"
        }
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8002)