main.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.responses import HTMLResponse
import requests
import os


# ----------------------------
# 1. Configuration
# ----------------------------

# Remove hardcoded API key and use an environment variable
HF_API_KEY = os.getenv("HF_API_KEY")

if not HF_API_KEY:
    raise RuntimeError("Hugging Face API key is not set. Please set the HF_API_KEY environment variable.")

HF_MODEL_NAME = "gpt2"  # A reliable text generation model available on HF Inference API

# ----------------------------
# 2. FastAPI App Initialization
# ----------------------------

app = FastAPI(
    title="AI Code Review Service",
    description="An API to get AI-powered code reviews for pull request diffs.",
    version="1.0.0",
)

# ----------------------------
# 3. No Local Model Loading (Using HF API Instead)
# ----------------------------

@app.on_event("startup")
async def startup_event():
    """
    On server startup, validate HF API key.
    """
    print("Server starting up...")
    print(f"Using Hugging Face API with model: {HF_MODEL_NAME}")
    if not HF_API_KEY:
        print("WARNING: HF_API_KEY not set!")
    else:
        print("HF_API_KEY is configured.")

# ----------------------------
# 4. API Request/Response Models
# ----------------------------

class ReviewRequest(BaseModel):
    diff: str

class ReviewComment(BaseModel):
    file_path: str
    line_number: int
    comment_text: str

class ReviewResponse(BaseModel):
    comments: list[ReviewComment]

# ----------------------------
# 5. The AI Review Logic
# ----------------------------

def run_ai_inference(diff: str) -> str:
    """
    Sends the code diff to Hugging Face Inference API to get the review.
    """
    # Better prompt for meaningful completions
    prompt = f"""Code review feedback:

{diff[:200]}

Feedback: This code could be improved by"""

    headers = {
        "Authorization": f"Bearer {HF_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": 32,
            "temperature": 0.7,
            "top_p": 0.9
        }
    }

    try:
        response = requests.post(
            f"https://api-inference.huggingface.co/models/{HF_MODEL_NAME}",
            headers=headers,
            json=payload,
            timeout=30
        )

        if response.status_code != 200:
            print(f"HF API Error: {response.status_code} - {response.text}")
            return "Consider adding proper documentation and error handling."

        response_data = response.json()
        print(f"HF API Response: {response_data}")
        
        if isinstance(response_data, list) and len(response_data) > 0:
            generated_text = response_data[0].get("generated_text", "")
            # Extract only the new generated part (after our prompt)
            if generated_text.startswith(prompt):
                response_text = generated_text[len(prompt):].strip()
            else:
                response_text = generated_text.strip()
        else:
            response_text = "Unable to generate a meaningful review."

    except Exception as e:
        print(f"HF API Exception: {e}")
        return "Consider adding proper documentation and error handling."

    # Clean up the response
    response_text = response_text.strip()

    # Handle different completion patterns
    if response_text.startswith("error handling"):
        review = "Consider adding error handling and input validation."
    elif response_text.startswith("documentation"):
        review = "Consider adding documentation and type hints."
    elif response_text.startswith("input validation"):
        review = "Consider adding input validation and error checks."
    elif response_text.startswith("type hints"):
        review = "Consider adding type hints and documentation."
    else:
        # Extract meaningful content
        lines = [line.strip() for line in response_text.split('\n') if line.strip()]
        if lines and len(lines[0]) > 3:
            first_line = lines[0]
            # Clean up common artifacts
            if first_line.startswith('#'):
                first_line = first_line[1:].strip()
            if len(first_line) > 10:
                review = f"Consider adding {first_line.lower()}."
            else:
                review = "Consider adding proper documentation and error handling."
        else:
            review = "Consider adding proper documentation and error handling."

    return review

def parse_ai_response(response_text: str) -> list[ReviewComment]:
    """
    Parses the raw text from the AI to extract the JSON array.
    """
    # For codegen-350M-mono, just wrap the review in a single comment
    return [ReviewComment(
        file_path="code_reviewed.py",
        line_number=1,
        comment_text=response_text.strip()
    )]

# ----------------------------
# 6. The API Endpoint
# ----------------------------

@app.post("/review", response_model=ReviewResponse)
async def get_code_review(request: ReviewRequest):
    if not request.diff:
        raise HTTPException(status_code=400, detail="Diff content cannot be empty.")

    import time
    start_time = time.time()
    print(f"Starting review request at {start_time}")

    try:
        print("Running AI inference...")
        ai_response_text = run_ai_inference(request.diff)
        print(f"AI inference completed in {time.time() - start_time:.2f} seconds")
        
        print("Parsing AI response...")
        parsed_comments = parse_ai_response(ai_response_text)
        print(f"Total processing time: {time.time() - start_time:.2f} seconds")
        
        return ReviewResponse(comments=parsed_comments)

    except Exception as e:
        print(f"An unexpected error occurred after {time.time() - start_time:.2f} seconds: {e}")
        raise HTTPException(status_code=500, detail="An internal error occurred while processing the review.")

# ----------------------------
# 7. Health Check Endpoint
# ----------------------------
@app.get("/", response_class=HTMLResponse)
def root_html():
    """Return HTML for browser viewing."""
    return """
    <!DOCTYPE html>
    <html>
    <head>
        <title>AI Code Review Service</title>
        <style>
            body { font-family: Arial, sans-serif; margin: 40px; }
            .status { color: green; font-weight: bold; }
            .endpoint { background: #f4f4f4; padding: 10px; margin: 10px 0; border-radius: 5px; }
        </style>
    </head>
    <body>
        <h1>AI Code Review Service</h1>
        <p class="status">✅ Service is running with AI model</p>
        
        <h2>Available Endpoints:</h2>
        <div class="endpoint"><strong>GET /health</strong> - Health check</div>
        <div class="endpoint"><strong>POST /review</strong> - Submit code diff for review</div>
        <div class="endpoint"><strong>GET /docs</strong> - Interactive API documentation</div>
        
        <h2>Quick Test:</h2>
        <p><a href="/health">Test Health Endpoint</a></p>
        <p><a href="/docs">View API Documentation</a></p>
        
        <h2>Status:</h2>
        <ul>
            <li>Mode: Hugging Face API</li>
            <li>AI Model: GPT-2</li>
            <li>Response Time: ~2-5 seconds</li>
        </ul>
    </body>
    </html>
    """

@app.get("/health")
async def health_check():
    return {"status": "ok", "api_configured": HF_API_KEY is not None}