main.py

import time
import json
import requests
import traceback
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

app = FastAPI(title="Gemini AI Dashboard API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

API_KEYS = [
    "AIzaSyAIo-DJ1iTwoP-jaU0rhhHgojxMdXCDhB8",
    "AIzaSyASgxIkI7xM2QwcNBOUWT7sDqpBRKAOvYs"
]

class AnalysisRequest(BaseModel):
    text: str
    task: str = "dashboard"

@app.get("/")
async def root():
    return {"status": "running", "message": "API is online."}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

def get_best_free_model_and_key():
    for api_key in API_KEYS:
        list_url = f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}"
        try:
            resp = requests.get(list_url, timeout=10)
            if resp.status_code == 200:
                data = resp.json()
                available_models = data.get("models", [])
                
                flash_models = []
                other_models = []
                
                for model in available_models:
                    if "generateContent" in model.get("supportedGenerationMethods", []):
                        name = model["name"].replace("models/", "")
                        if "flash" in name.lower():
                            flash_models.append(name)
                        else:
                            other_models.append(name)
                            
                if flash_models:
                    return flash_models[0], api_key
                elif other_models:
                    return other_models[0], api_key
        except Exception as e:
            print(f"Error fetching models for key {api_key[-4:]}: {e}")
            pass
    return None, None

@app.post("/dashboard")
async def generate_dashboard(request: AnalysisRequest):
    start_time = time.time()
    input_data = request.text[:15000] 
    
    model_name, working_key = get_best_free_model_and_key()
    
    if not model_name:
        raise HTTPException(status_code=500, detail="No free models available for your API keys.")
        
    print(f"Using model: {model_name} with key ending in ...{working_key[-4:]}")
    
    prompt = f"""
    You are a strict, mathematically precise Data Analyst AI. 
    Your job is to analyze the provided raw data and calculate exact metrics.
    
    CRITICAL INSTRUCTIONS:
    1. You MUST actually read the data, count the rows, and sum the numbers. 
    2. DO NOT hallucinate, guess, or use placeholder numbers. 
    3. If you cannot calculate a specific number, do not include it.
    4. All "value" fields in charts and metrics MUST be real numbers calculated directly from the data provided.
    
    The raw data to analyze is:
    ---
    {input_data}
    ---
    
    You MUST return ONLY a valid JSON object with this exact structure. 
    Replace all <placeholders> with actual calculated values from the data.
    
    {{
      "summary": "A factual summary of exactly what the data is and the most obvious mathematical trend.",
      "metrics": [
        {{"label": "<Exact Metric Name 1>", "value": "<Exact Calculated Value>", "change": "<Exact Change if available, else N/A>"}},
        {{"label": "<Exact Metric Name 2>", "value": "<Exact Calculated Value>", "change": "<Exact Change if available, else N/A>"}},
        {{"label": "<Exact Metric Name 3>", "value": "<Exact Calculated Value>", "change": "<Exact Change if available, else N/A>"}},
        {{"label": "<Exact Metric Name 4>", "value": "<Exact Calculated Value>", "change": "<Exact Change if available, else N/A>"}}
      ],
      "charts": [
        {{
          "type": "bar",
          "title": "<Title based on data categories>",
          "data": [{{"name": "<Category 1>", "value": <Exact Sum or Count>}}, {{"name": "<Category 2>", "value": <Exact Sum or Count>}}]
        }},
        {{
          "type": "pie",
          "title": "<Title based on data distribution>",
          "data": [{{"name": "<Segment 1>", "value": <Exact Count>}}, {{"name": "<Segment 2>", "value": <Exact Count>}}]
        }}
      ],
      "insights": ["Factual finding 1 based on the numbers", "Factual finding 2 based on the numbers"],
      "recommendations": ["Actionable recommendation based strictly on the data"],
      "risks": ["Risk identified in the data"],
      "opportunities": ["Opportunity identified in the data"],
      "tables": [
        {{
          "title": "Data Sample Breakdown",
          "headers": ["<Col 1>", "<Col 2>", "<Col 3>"],
          "rows": [
            ["<Actual Row 1 Data>", "<Actual Row 1 Data>", "<Actual Row 1 Data>"],
            ["<Actual Row 2 Data>", "<Actual Row 2 Data>", "<Actual Row 2 Data>"]
          ]
        }}
      ],
      "confidence_score": 90,
      "dataset_info": {{
        "name": "<Inferred Dataset Name>",
        "records": <Exact count of rows/records in the data>, 
        "columns": <Exact count of columns/fields in the data>, 
        "detected_type": "<Type of data>", 
        "upload_time": "0.1s", 
        "processing_time": "1.2s"
      }}
    }}
    """
    
    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {
            "temperature": 0.1,
            "maxOutputTokens": 8000
        }
    }

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={working_key}"
    
    try:
        response = requests.post(url, json=payload, headers={"Content-Type": "application/json"}, timeout=45)
        
        if response.status_code != 200:
            error_msg = f"Gemini API Error {response.status_code}: {response.text}"
            print(error_msg)
            raise HTTPException(status_code=500, detail=error_msg)
            
        gemini_result = response.json()

        if 'candidates' in gemini_result and len(gemini_result['candidates']) > 0:
            if 'content' in gemini_result['candidates'][0] and 'parts' in gemini_result['candidates'][0]['content']:
                raw_text = gemini_result['candidates'][0]['content']['parts'][0]['text']
            else:
                error_msg = f"Gemini blocked the prompt. Full response: {json.dumps(gemini_result)}"
                print(error_msg)
                raise HTTPException(status_code=500, detail=error_msg)
        else:
            error_msg = f"Gemini returned no candidates. Full response: {json.dumps(gemini_result)}"
            print(error_msg)
            raise HTTPException(status_code=500, detail=error_msg)
        
        if raw_text.startswith("```json"):
            raw_text = raw_text[7:]
        if raw_text.endswith("```"):
            raw_text = raw_text[:-3]
            
        # --- ADDED JSON PARSE ERROR HANDLING ---
        try:
            dashboard_data = json.loads(raw_text.strip())
        except json.JSONDecodeError:
            print(f"JSON Parse Error. Raw text from AI: {raw_text}")
            raise HTTPException(status_code=500, detail="The AI returned an empty or invalid JSON response. Please try again.")
        
        processing_time = f"{time.time() - start_time:.1f}s"
        if 'dataset_info' in dashboard_data:
            dashboard_data['dataset_info']['processing_time'] = processing_time
        else:
            dashboard_data['dataset_info'] = {"processing_time": processing_time}
        
        return dashboard_data
        
    except HTTPException:
        raise
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")