src/app.py

import os
import json
from fastapi import FastAPI
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from typing import Dict, List
from src.doper import UniversalMaterialsOptimizer
from dotenv import load_dotenv

app = FastAPI(
    title="AURELIUS Battery Optimizer",
    description="Generative Design API for Co-Doped Solid State Batteries with Streaming Support",
    version="4.0"
)

# Load environment variables from .env file
load_dotenv()

# Get API Key from Hugging Face Secrets (Environment Variable)
MP_API_KEY = os.getenv("MP_API_KEY") or os.getenv("MP_API")

# --- SERVE THE UI ---
# This makes the "static" folder accessible
app.mount("/static", StaticFiles(directory="static"), name="static")

# --- NEW INPUT MODEL ---
class MultiDopantRequest(BaseModel):
    host_formula: str = Field(..., example="Li3PS4")
    host_site_element: str = Field(..., example="S")
    
    # Now we accept a list of RECIPES to test
    # Example: Check "Just Cl" AND "Cl + Br mixed"
    recipes: List[Dict[str, float]] = Field(..., example=[
        {"Cl": 0.2},              # Pure Cl doping
        {"Cl": 0.1, "Br": 0.1},   # Co-doping (High Entropy)
        {"I": 0.05, "F": 0.05}    # Exotic mix
    ])

@app.get("/")
async def read_index():
    # When user hits the homepage, send the HTML file
    return FileResponse('static/index.html')

# --- THE GENERATOR FUNCTION ---
# This function doesn't run all at once. It pauses and resumes.
async def run_simulation_stream(job: MultiDopantRequest):
    
    # 1. Initialize Engine (Once)
    optimizer = UniversalMaterialsOptimizer(
        host_formula=job.host_formula,
        host_element_symbol=job.host_site_element,
        api_key=MP_API_KEY
    )

    # 2. First, stream the metadata header
    # This tells the client "I am ready, here is the context"
    metadata = {
        "type": "meta",
        "host_system": job.host_formula,
        "base_properties": optimizer.props
    }
    yield json.dumps(metadata) + "\n"

    # 3. Stream results one by one
    for recipe in job.recipes:
        try:
            result = optimizer.analyze_mixture(recipe)
            # Add a type tag so client knows this is a data row
            result["type"] = "data"
            
            # Send this chunk immediately!
            yield json.dumps(result) + "\n"
            
        except Exception as e:
            error_msg = {"type": "error", "recipe": recipe, "msg": str(e)}
            yield json.dumps(error_msg) + "\n"

@app.post("/simulate_stream")
def simulate_mixture_stream(job: MultiDopantRequest):
    """
    Returns a StreamingResponse. 
    The client will receive line-by-line JSON as calculations finish.
    Ideally, use this endpoint for large numbers of recipes.
    For small numbers of recipes, use /simulate_mix instead.
    """
    return StreamingResponse(
        run_simulation_stream(job), 
        media_type="application/x-ndjson"
    )

@app.post("/simulate_mix")
def simulate_mixture(job: MultiDopantRequest):
    """
    Non-streaming endpoint for backward compatibility.
    Returns all results at once after all calculations complete.
    Use this endpoint if there are relatively few recipes to test.
    For large numbers of recipes, use /simulate_stream instead.
    """
    optimizer = UniversalMaterialsOptimizer(
        host_formula=job.host_formula,
        host_element_symbol=job.host_site_element,
        api_key=MP_API_KEY
    )

    results = []
    for recipe in job.recipes:
        try:
            res = optimizer.analyze_mixture(recipe)
            results.append(res)
        except Exception as e:
            results.append({"recipe": recipe, "error": str(e)})

    return {
        "host_system": job.host_formula,
        "base_properties": optimizer.props,
        "results": results
    }