import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

class AIClient:
    def __init__(self):
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY"),
        )
        # User requested model. Note: If 'openrouter/free' fails, 
        # try a specific free model ID like 'mistralai/mistral-7b-instruct:free'
        self.model = "openrouter/free" 

    def ask(self, system_prompt: str, user_prompt: str) -> str:
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                # Optional headers for OpenRouter ranking
                extra_headers={
                    "HTTP-Referer": "https://localhost", 
                    "X-Title": "AI Code Auditor"
                }
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"[ERROR] AI Request Failed: {str(e)}"