from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, List, Any, Optional
import os
from groq import Groq

app = FastAPI(title="Architecture Evaluator LLM")

# Enable CORS for all origins (adjust in production)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize Groq client - FREE and FAST!
client = Groq(
    api_key=os.getenv("GROQ_API_KEY", "your-groq-api-key-here")
)

class EvaluationRequest(BaseModel):
    question: str
    architecture: Dict[str, Any]

class HeuristicScores(BaseModel):
    DURABILITY: float
    AVAILABILITY: float
    ENERGY_EFFICIENCY: float
    CONSISTENCY: float
    MAINTAINABILITY: float
    LATENCY: float
    COST: float
    SECURITY: float
    THROUGHPUT: float
    SCALABILITY: float

class EvaluationResponse(BaseModel):
    heuristic_scores: HeuristicScores
    suggestion: str

def create_evaluation_prompt(question: str, architecture: Dict[str, Any]) -> str:
    """Create a detailed prompt for the LLM to evaluate the architecture."""
    
    prompt = f"""You are an expert system architect. Analyze the following system design architecture and provide an evaluation.

**Question/Context:**
{question}

**Architecture JSON:**
{architecture}

**Your Task:**
1. Analyze the overall architecture based on the components and their connections (links)
2. Calculate aggregate heuristic scores (0-10 scale) for the entire architecture considering:
   - Individual component scores
   - How components interact through links
   - Overall system design patterns
   - Redundancy and fault tolerance
   - Data flow efficiency

3. Provide constructive suggestions on:
   - What's good about this architecture
   - What could be improved
   - Specific recommendations
   - Alternative approaches if applicable

**Important:** You must respond with ONLY a valid JSON object in this exact format:
{{
  "heuristic_scores": {{
    "DURABILITY": <float between 0-10>,
    "AVAILABILITY": <float between 0-10>,
    "ENERGY_EFFICIENCY": <float between 0-10>,
    "CONSISTENCY": <float between 0-10>,
    "MAINTAINABILITY": <float between 0-10>,
    "LATENCY": <float between 0-10>,
    "COST": <float between 0-10>,
    "SECURITY": <float between 0-10>,
    "THROUGHPUT": <float between 0-10>,
    "SCALABILITY": <float between 0-10>
  }},
  "suggestion": "<your detailed analysis and suggestions as a string which should be in bulleted point>"
}}

Do not include any text outside the JSON object. Ensure all scores are numeric values between 0 and 10."""
    
    return prompt

@app.get("/")
async def root():
    return {
        "message": "Architecture Evaluator LLM API",
        "status": "online",
        "endpoints": {
            "/evaluate": "POST - Evaluate a system architecture",
            "/health": "GET - Health check"
        }
    }

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

@app.post("/evaluate", response_model=EvaluationResponse)
async def evaluate_architecture(request: EvaluationRequest):
    """
    Evaluate a system design architecture using LLM.
    
    Args:
        request: Contains the question and architecture JSON
        
    Returns:
        EvaluationResponse with heuristic scores and suggestions
    """
    try:
        # Create the prompt
        prompt = create_evaluation_prompt(request.question, request.architecture)
        
        # Call Groq LLM - Using best reasoning model (deepseek-r1-distill-llama-70b)
        # Alternative models: llama-3.3-70b-versatile, llama-3.1-70b-versatile
        response = client.chat.completions.create(
            model=os.getenv("MODEL_NAME", "llama-3.3-70b-versatile"),
            messages=[
                {
                    "role": "system",
                    "content": "You are an expert system architect who provides detailed analysis and returns responses"
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=0.7,
            response_format={"type": "json_object"}  # Enforce JSON response
        )
        
        # Parse the response
        import json
        result = json.loads(response.choices[0].message.content)
        
        # Validate and return
        return EvaluationResponse(**result)
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error evaluating architecture: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)  # Port 7860 is standard for HuggingFace Spaces