Atomizer/atomizer-dashboard/backend/api/routes/optimization.py

"""
Optimization API endpoints
Handles study status, history retrieval, and control operations
"""

from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from fastapi.responses import JSONResponse, FileResponse
from pathlib import Path
from typing import List, Dict, Optional
import json
import sys
import sqlite3
import shutil
from datetime import datetime

# Add project root to path
sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))

router = APIRouter()

# Base studies directory
STUDIES_DIR = Path(__file__).parent.parent.parent.parent.parent / "studies"

@router.get("/studies")
async def list_studies():
    """List all available optimization studies"""
    try:
        studies = []
        if not STUDIES_DIR.exists():
            return {"studies": []}

        for study_dir in STUDIES_DIR.iterdir():
            if not study_dir.is_dir():
                continue

            # Look for optimization config (check multiple locations)
            config_file = study_dir / "optimization_config.json"
            if not config_file.exists():
                config_file = study_dir / "1_setup" / "optimization_config.json"
            if not config_file.exists():
                continue

            # Load config
            with open(config_file) as f:
                config = json.load(f)

            # Check if results directory exists
            results_dir = study_dir / "2_results"

            # Check for Optuna database (Protocol 10) or JSON history (other protocols)
            study_db = results_dir / "study.db"
            history_file = results_dir / "optimization_history_incremental.json"

            status = "not_started"
            trial_count = 0
            best_value = None

            # Protocol 10: Read from Optuna SQLite database
            if study_db.exists():
                try:
                    conn = sqlite3.connect(str(study_db))
                    cursor = conn.cursor()

                    # Get trial count and status
                    cursor.execute("SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'")
                    trial_count = cursor.fetchone()[0]

                    # Get best trial (for single-objective, or first objective for multi-objective)
                    if trial_count > 0:
                        cursor.execute("""
                            SELECT value FROM trial_values
                            WHERE trial_id IN (
                                SELECT trial_id FROM trials WHERE state = 'COMPLETE'
                            )
                            ORDER BY value ASC
                            LIMIT 1
                        """)
                        result = cursor.fetchone()
                        if result:
                            best_value = result[0]

                    conn.close()

                    # Determine status
                    total_trials = config.get('optimization_settings', {}).get('n_trials', 50)
                    if trial_count >= total_trials:
                        status = "completed"
                    else:
                        status = "running"  # Simplified - would need process check

                except Exception as e:
                    print(f"Warning: Failed to read Optuna database for {study_dir.name}: {e}")
                    status = "error"

            # Legacy: Read from JSON history
            elif history_file.exists():
                with open(history_file) as f:
                    history = json.load(f)
                    trial_count = len(history)
                    if history:
                        # Find best trial
                        best_trial = min(history, key=lambda x: x['objective'])
                        best_value = best_trial['objective']

                        # Determine status
                        total_trials = config.get('trials', {}).get('n_trials', 50)
                        if trial_count >= total_trials:
                            status = "completed"
                        else:
                            status = "running"  # Simplified - would need process check

            # Get total trials from config (supports both formats)
            total_trials = (
                config.get('optimization_settings', {}).get('n_trials') or
                config.get('trials', {}).get('n_trials', 50)
            )

            studies.append({
                "id": study_dir.name,
                "name": study_dir.name.replace("_", " ").title(),
                "status": status,
                "progress": {
                    "current": trial_count,
                    "total": total_trials
                },
                "best_value": best_value,
                "target": config.get('target', {}).get('value'),
                "path": str(study_dir)
            })

        return {"studies": studies}

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to list studies: {str(e)}")

@router.get("/studies/{study_id}/status")
async def get_study_status(study_id: str):
    """Get detailed status of a specific study"""
    try:
        study_dir = STUDIES_DIR / study_id
        if not study_dir.exists():
            raise HTTPException(status_code=404, detail=f"Study {study_id} not found")

        # Load config (check multiple locations)
        config_file = study_dir / "optimization_config.json"
        if not config_file.exists():
            config_file = study_dir / "1_setup" / "optimization_config.json"

        with open(config_file) as f:
            config = json.load(f)

        # Check for results
        results_dir = study_dir / "2_results"
        study_db = results_dir / "study.db"
        history_file = results_dir / "optimization_history_incremental.json"

        # Protocol 10: Read from Optuna database
        if study_db.exists():
            conn = sqlite3.connect(str(study_db))
            cursor = conn.cursor()

            # Get trial counts by state
            cursor.execute("SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'")
            trial_count = cursor.fetchone()[0]

            cursor.execute("SELECT COUNT(*) FROM trials WHERE state = 'PRUNED'")
            pruned_count = cursor.fetchone()[0]

            # Get best trial (first objective for multi-objective)
            best_trial = None
            if trial_count > 0:
                cursor.execute("""
                    SELECT t.trial_id, t.number, tv.value
                    FROM trials t
                    JOIN trial_values tv ON t.trial_id = tv.trial_id
                    WHERE t.state = 'COMPLETE'
                    ORDER BY tv.value ASC
                    LIMIT 1
                """)
                result = cursor.fetchone()
                if result:
                    trial_id, trial_number, best_value = result

                    # Get parameters for this trial
                    cursor.execute("""
                        SELECT param_name, param_value
                        FROM trial_params
                        WHERE trial_id = ?
                    """, (trial_id,))
                    params = {row[0]: row[1] for row in cursor.fetchall()}

                    best_trial = {
                        "trial_number": trial_number,
                        "objective": best_value,
                        "design_variables": params,
                        "results": {"first_frequency": best_value}
                    }

            conn.close()

            total_trials = config.get('optimization_settings', {}).get('n_trials', 50)
            status = "completed" if trial_count >= total_trials else "running"

            return {
                "study_id": study_id,
                "status": status,
                "progress": {
                    "current": trial_count,
                    "total": total_trials,
                    "percentage": (trial_count / total_trials * 100) if total_trials > 0 else 0
                },
                "best_trial": best_trial,
                "pruned_trials": pruned_count,
                "config": config
            }

        # Legacy: Read from JSON history
        if not history_file.exists():
            return {
                "study_id": study_id,
                "status": "not_started",
                "progress": {"current": 0, "total": config.get('trials', {}).get('n_trials', 50)},
                "config": config
            }

        with open(history_file) as f:
            history = json.load(f)

        trial_count = len(history)
        total_trials = config.get('trials', {}).get('n_trials', 50)

        # Find best trial
        best_trial = None
        if history:
            best_trial = min(history, key=lambda x: x['objective'])

        # Check for pruning data
        pruning_file = results_dir / "pruning_history.json"
        pruned_count = 0
        if pruning_file.exists():
            with open(pruning_file) as f:
                pruning_history = json.load(f)
                pruned_count = len(pruning_history)

        status = "completed" if trial_count >= total_trials else "running"

        return {
            "study_id": study_id,
            "status": status,
            "progress": {
                "current": trial_count,
                "total": total_trials,
                "percentage": (trial_count / total_trials * 100) if total_trials > 0 else 0
            },
            "best_trial": best_trial,
            "pruned_trials": pruned_count,
            "config": config
        }

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get study status: {str(e)}")

@router.get("/studies/{study_id}/history")
async def get_optimization_history(study_id: str, limit: Optional[int] = None):
    """Get optimization history (all trials)"""
    try:
        study_dir = STUDIES_DIR / study_id
        results_dir = study_dir / "2_results"
        study_db = results_dir / "study.db"
        history_file = results_dir / "optimization_history_incremental.json"

        # Protocol 10: Read from Optuna database
        if study_db.exists():
            conn = sqlite3.connect(str(study_db))
            cursor = conn.cursor()

            # Get all completed trials
            cursor.execute("""
                SELECT trial_id, number, datetime_start, datetime_complete
                FROM trials
                WHERE state = 'COMPLETE'
                ORDER BY number DESC
            """ + (f" LIMIT {limit}" if limit else ""))

            trial_rows = cursor.fetchall()

            trials = []
            for trial_id, trial_num, start_time, end_time in trial_rows:
                # Get objectives for this trial
                cursor.execute("""
                    SELECT value
                    FROM trial_values
                    WHERE trial_id = ?
                    ORDER BY objective
                """, (trial_id,))
                values = [row[0] for row in cursor.fetchall()]

                # Get parameters for this trial
                cursor.execute("""
                    SELECT param_name, param_value
                    FROM trial_params
                    WHERE trial_id = ?
                """, (trial_id,))
                params = {}
                for param_name, param_value in cursor.fetchall():
                    try:
                        params[param_name] = float(param_value) if param_value is not None else None
                    except (ValueError, TypeError):
                        params[param_name] = param_value

                # Get user attributes (extracted results: mass, frequency, stress, displacement, etc.)
                cursor.execute("""
                    SELECT key, value_json
                    FROM trial_user_attributes
                    WHERE trial_id = ?
                """, (trial_id,))
                user_attrs = {}
                for key, value_json in cursor.fetchall():
                    try:
                        user_attrs[key] = json.loads(value_json)
                    except (ValueError, TypeError):
                        user_attrs[key] = value_json

                # Extract relevant metrics for results (mass, frequency, stress, displacement, etc.)
                results = {}
                if "mass" in user_attrs:
                    results["mass"] = user_attrs["mass"]
                if "frequency" in user_attrs:
                    results["frequency"] = user_attrs["frequency"]
                if "max_stress" in user_attrs:
                    results["max_stress"] = user_attrs["max_stress"]
                if "max_displacement" in user_attrs:
                    results["max_displacement"] = user_attrs["max_displacement"]
                # Fallback to first frequency from objectives if available
                if not results and len(values) > 0:
                    results["first_frequency"] = values[0]

                # CRITICAL: Extract design_vars from user_attrs if stored there
                # The optimization code does: trial.set_user_attr("design_vars", design_vars)
                design_vars_from_attrs = user_attrs.get("design_vars", {})

                # Merge with params (prefer user_attrs design_vars if available)
                final_design_vars = {**params, **design_vars_from_attrs} if design_vars_from_attrs else params

                trials.append({
                    "trial_number": trial_num,
                    "objective": values[0] if len(values) > 0 else None,  # Primary objective
                    "objectives": values if len(values) > 1 else None,  # All objectives for multi-objective
                    "design_variables": final_design_vars,  # Use merged design vars
                    "results": results,
                    "user_attrs": user_attrs,  # Include all user attributes
                    "start_time": start_time,
                    "end_time": end_time
                })

            conn.close()
            return {"trials": trials}

        # Legacy: Read from JSON history
        if not history_file.exists():
            return {"trials": []}

        with open(history_file) as f:
            history = json.load(f)

        # Apply limit if specified
        if limit:
            history = history[-limit:]

        return {"trials": history}

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get history: {str(e)}")

@router.get("/studies/{study_id}/pruning")
async def get_pruning_history(study_id: str):
    """Get pruning diagnostics"""
    try:
        study_dir = STUDIES_DIR / study_id
        pruning_file = study_dir / "2_results" / "pruning_history.json"

        if not pruning_file.exists():
            return {"pruned_trials": []}

        with open(pruning_file) as f:
            pruning_history = json.load(f)

        return {"pruned_trials": pruning_history}

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get pruning history: {str(e)}")

def _infer_objective_unit(objective: Dict) -> str:
    """Infer unit from objective name and description"""
    name = objective.get("name", "").lower()
    desc = objective.get("description", "").lower()

    # Common unit patterns
    if "frequency" in name or "hz" in desc:
        return "Hz"
    elif "stiffness" in name or "n/mm" in desc:
        return "N/mm"
    elif "mass" in name or "kg" in desc:
        return "kg"
    elif "stress" in name or "mpa" in desc or "pa" in desc:
        return "MPa"
    elif "displacement" in name or "mm" in desc:
        return "mm"
    elif "force" in name or "newton" in desc:
        return "N"
    elif "%" in desc or "percent" in desc:
        return "%"

    # Check if unit is explicitly mentioned in description (e.g., "(N/mm)")
    import re
    unit_match = re.search(r'\(([^)]+)\)', desc)
    if unit_match:
        return unit_match.group(1)

    return ""  # No unit found

@router.get("/studies/{study_id}/metadata")
async def get_study_metadata(study_id: str):
    """Read optimization_config.json for objectives, design vars, units (Protocol 13)"""
    try:
        study_dir = STUDIES_DIR / study_id
        if not study_dir.exists():
            raise HTTPException(status_code=404, detail=f"Study {study_id} not found")

        # Load config (check multiple locations)
        config_file = study_dir / "optimization_config.json"
        if not config_file.exists():
            config_file = study_dir / "1_setup" / "optimization_config.json"

        if not config_file.exists():
            raise HTTPException(status_code=404, detail=f"Config file not found for study {study_id}")

        with open(config_file) as f:
            config = json.load(f)

        # Enhance objectives with inferred units if not present
        objectives = config.get("objectives", [])
        for obj in objectives:
            if "unit" not in obj or not obj["unit"]:
                obj["unit"] = _infer_objective_unit(obj)

        return {
            "objectives": objectives,
            "design_variables": config.get("design_variables", []),
            "constraints": config.get("constraints", []),
            "study_name": config.get("study_name", study_id),
            "description": config.get("description", "")
        }

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get study metadata: {str(e)}")

@router.get("/studies/{study_id}/optimizer-state")
async def get_optimizer_state(study_id: str):
    """Read realtime optimizer state from intelligent_optimizer/ (Protocol 13)"""
    try:
        study_dir = STUDIES_DIR / study_id
        results_dir = study_dir / "2_results"
        state_file = results_dir / "intelligent_optimizer" / "optimizer_state.json"

        if not state_file.exists():
            return {"available": False}

        with open(state_file) as f:
            state = json.load(f)

        return {"available": True, **state}

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get optimizer state: {str(e)}")

@router.get("/studies/{study_id}/pareto-front")
async def get_pareto_front(study_id: str):
    """Get Pareto-optimal solutions for multi-objective studies (Protocol 13)"""
    try:
        study_dir = STUDIES_DIR / study_id
        results_dir = study_dir / "2_results"
        study_db = results_dir / "study.db"

        if not study_db.exists():
            return {"is_multi_objective": False, "pareto_front": []}

        # Import optuna here to avoid loading it for all endpoints
        import optuna

        storage = optuna.storages.RDBStorage(f"sqlite:///{study_db}")
        study = optuna.load_study(study_name=study_id, storage=storage)

        # Check if multi-objective
        if len(study.directions) == 1:
            return {"is_multi_objective": False, "pareto_front": []}

        # Get Pareto front
        pareto_trials = study.best_trials

        return {
            "is_multi_objective": True,
            "pareto_front": [
                {
                    "trial_number": t.number,
                    "values": t.values,
                    "params": t.params,
                    "user_attrs": dict(t.user_attrs),
                    "constraint_satisfied": t.user_attrs.get("constraint_satisfied", True)
                }
                for t in pareto_trials
            ]
        }

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get Pareto front: {str(e)}")

@router.post("/studies")
async def create_study(
    config: str = Form(...),
    prt_file: Optional[UploadFile] = File(None),
    sim_file: Optional[UploadFile] = File(None),
    fem_file: Optional[UploadFile] = File(None)
):
    """
    Create a new optimization study
    Accepts:
    - config: JSON string with study configuration
    - prt_file: NX part file (optional if using existing study)
    - sim_file: NX simulation file (optional)
    - fem_file: NX FEM file (optional)
    """
    try:
        # Parse config
        config_data = json.loads(config)
        study_name = config_data.get("name") # Changed from study_name to name to match frontend

        if not study_name:
            raise HTTPException(status_code=400, detail="name is required in config")

        # Create study directory structure
        study_dir = STUDIES_DIR / study_name
        if study_dir.exists():
            raise HTTPException(status_code=400, detail=f"Study {study_name} already exists")

        setup_dir = study_dir / "1_setup"
        model_dir = setup_dir / "model"
        results_dir = study_dir / "2_results"

        setup_dir.mkdir(parents=True, exist_ok=True)
        model_dir.mkdir(parents=True, exist_ok=True)
        results_dir.mkdir(parents=True, exist_ok=True)

        # Save config file
        config_file = setup_dir / "optimization_config.json"
        with open(config_file, 'w') as f:
            json.dump(config_data, f, indent=2)

        # Save uploaded files
        files_saved = {}
        if prt_file:
            prt_path = model_dir / prt_file.filename
            with open(prt_path, 'wb') as f:
                content = await prt_file.read()
                f.write(content)
            files_saved['prt_file'] = str(prt_path)

        if sim_file:
            sim_path = model_dir / sim_file.filename
            with open(sim_path, 'wb') as f:
                content = await sim_file.read()
                f.write(content)
            files_saved['sim_file'] = str(sim_path)

        if fem_file:
            fem_path = model_dir / fem_file.filename
            with open(fem_path, 'wb') as f:
                content = await fem_file.read()
                f.write(content)
            files_saved['fem_file'] = str(fem_path)

        return JSONResponse(
            status_code=201,
            content={
                "status": "created",
                "study_id": study_name,
                "study_path": str(study_dir),
                "config_path": str(config_file),
                "files_saved": files_saved,
                "message": f"Study {study_name} created successfully. Ready to run optimization."
            }
        )

    except json.JSONDecodeError as e:
        raise HTTPException(status_code=400, detail=f"Invalid JSON in config: {str(e)}")
    except Exception as e:
        # Clean up on error
        if 'study_dir' in locals() and study_dir.exists():
            shutil.rmtree(study_dir)
        raise HTTPException(status_code=500, detail=f"Failed to create study: {str(e)}")

@router.post("/studies/{study_id}/convert-mesh")
async def convert_study_mesh(study_id: str):
    """
    Convert study mesh to GLTF for 3D visualization
    Creates a web-viewable 3D model with FEA results as vertex colors
    """
    try:
        study_dir = STUDIES_DIR / study_id
        if not study_dir.exists():
            raise HTTPException(status_code=404, detail=f"Study {study_id} not found")

        # Import mesh converter
        sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
        from optimization_engine.mesh_converter import convert_study_mesh

        # Convert mesh
        output_path = convert_study_mesh(study_dir)

        if output_path and output_path.exists():
            return {
                "status": "success",
                "gltf_path": str(output_path),
                "gltf_url": f"/api/optimization/studies/{study_id}/mesh/model.gltf",
                "metadata_url": f"/api/optimization/studies/{study_id}/mesh/model.json",
                "message": "Mesh converted successfully"
            }
        else:
            raise HTTPException(status_code=500, detail="Mesh conversion failed")

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to convert mesh: {str(e)}")

@router.get("/studies/{study_id}/mesh/{filename}")
async def get_mesh_file(study_id: str, filename: str):
    """
    Serve GLTF mesh files and metadata
    Supports .gltf, .bin, and .json files
    """
    try:
        # Validate filename to prevent directory traversal
        if '..' in filename or '/' in filename or '\\' in filename:
            raise HTTPException(status_code=400, detail="Invalid filename")

        study_dir = STUDIES_DIR / study_id
        visualization_dir = study_dir / "3_visualization"

        file_path = visualization_dir / filename

        if not file_path.exists():
            raise HTTPException(status_code=404, detail=f"File {filename} not found")

        # Determine content type
        suffix = file_path.suffix.lower()
        content_types = {
            '.gltf': 'model/gltf+json',
            '.bin': 'application/octet-stream',
            '.json': 'application/json',
            '.glb': 'model/gltf-binary'
        }

        content_type = content_types.get(suffix, 'application/octet-stream')

        return FileResponse(
            path=str(file_path),
            media_type=content_type,
            filename=filename
        )

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"File not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to serve mesh file: {str(e)}")

@router.get("/studies/{study_id}/optuna-url")
async def get_optuna_dashboard_url(study_id: str):
    """
    Get the Optuna dashboard URL for a specific study.
    Returns the URL to access the study in Optuna dashboard.

    The Optuna dashboard should be started with a relative path from the Atomizer root:
    sqlite:///studies/{study_id}/2_results/study.db
    """
    try:
        study_dir = STUDIES_DIR / study_id
        if not study_dir.exists():
            raise HTTPException(status_code=404, detail=f"Study {study_id} not found")

        results_dir = study_dir / "2_results"
        study_db = results_dir / "study.db"

        if not study_db.exists():
            raise HTTPException(status_code=404, detail=f"No Optuna database found for study {study_id}")

        # Get the study name from the database (may differ from folder name)
        import optuna
        storage = optuna.storages.RDBStorage(f"sqlite:///{study_db}")
        studies = storage.get_all_studies()

        if not studies:
            raise HTTPException(status_code=404, detail=f"No Optuna study found in database for {study_id}")

        # Use the actual study name from the database
        optuna_study_name = studies[0].study_name

        # Return URL info for the frontend
        # The dashboard should be running on port 8081 with the correct database
        return {
            "study_id": study_id,
            "optuna_study_name": optuna_study_name,
            "database_path": f"studies/{study_id}/2_results/study.db",
            "dashboard_url": f"http://localhost:8081/dashboard/studies/{studies[0]._study_id}",
            "dashboard_base": "http://localhost:8081",
            "note": "Optuna dashboard must be started with: sqlite:///studies/{study_id}/2_results/study.db"
        }

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to get Optuna URL: {str(e)}")

@router.post("/studies/{study_id}/generate-report")
async def generate_report(
    study_id: str,
    format: str = "markdown",
    include_llm_summary: bool = False
):
    """
    Generate an optimization report in the specified format

    Args:
        study_id: Study identifier
        format: Report format ('markdown', 'html', or 'pdf')
        include_llm_summary: Whether to include LLM-generated executive summary

    Returns:
        Information about the generated report including download URL
    """
    try:
        study_dir = STUDIES_DIR / study_id
        if not study_dir.exists():
            raise HTTPException(status_code=404, detail=f"Study {study_id} not found")

        # Validate format
        valid_formats = ['markdown', 'md', 'html', 'pdf']
        if format.lower() not in valid_formats:
            raise HTTPException(status_code=400, detail=f"Invalid format. Must be one of: {', '.join(valid_formats)}")

        # Import report generator
        sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
        from optimization_engine.report_generator import generate_study_report

        # Generate report
        output_path = generate_study_report(
            study_dir=study_dir,
            output_format=format.lower(),
            include_llm_summary=include_llm_summary
        )

        if output_path and output_path.exists():
            # Get relative path for URL
            rel_path = output_path.relative_to(study_dir)

            return {
                "status": "success",
                "format": format,
                "file_path": str(output_path),
                "download_url": f"/api/optimization/studies/{study_id}/reports/{output_path.name}",
                "file_size": output_path.stat().st_size,
                "message": f"Report generated successfully in {format} format"
            }
        else:
            raise HTTPException(status_code=500, detail="Report generation failed")

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Study {study_id} not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to generate report: {str(e)}")

@router.get("/studies/{study_id}/reports/{filename}")
async def download_report(study_id: str, filename: str):
    """
    Download a generated report file

    Args:
        study_id: Study identifier
        filename: Report filename

    Returns:
        Report file for download
    """
    try:
        # Validate filename to prevent directory traversal
        if '..' in filename or '/' in filename or '\\' in filename:
            raise HTTPException(status_code=400, detail="Invalid filename")

        study_dir = STUDIES_DIR / study_id
        results_dir = study_dir / "2_results"

        file_path = results_dir / filename

        if not file_path.exists():
            raise HTTPException(status_code=404, detail=f"Report file {filename} not found")

        # Determine content type
        suffix = file_path.suffix.lower()
        content_types = {
            '.md': 'text/markdown',
            '.html': 'text/html',
            '.pdf': 'application/pdf',
            '.json': 'application/json'
        }

        content_type = content_types.get(suffix, 'application/octet-stream')

        return FileResponse(
            path=str(file_path),
            media_type=content_type,
            filename=filename,
            headers={"Content-Disposition": f"attachment; filename={filename}"}
        )

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Report file not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to download report: {str(e)}")