import httpx
from pretor.protocol_plugin.model_protocol.modelbase import ModelBase

class OpenAIAdapter(ModelBase):
    def __init__(self, base_url: str, adapter_title: str, api_key: str = "archon-local"):
        self.adapter_title: str = adapter_title
        self.base_url = base_url.rstrip('/')
        if not self.base_url.endswith('/v1'):
            self.base_url += '/v1'
        self.api_key = api_key
        self.model_list = []

    async def get_model(self):
        url = "{}/models".format(self.base_url)
        headers = {
            "Authorization": f"Bearer {self.api_key}"
        }
        async with httpx.AsyncClient() as client:
            response = await client.get(url, headers=headers)
        response.raise_for_status()
        response = response.json()
        self.model_list = [m.get("id", "") for m in response.get("data", [])]
        return self.model_list

    async def post_message(self,model: str, messages: list, stream: bool = False, **kwargs):
        url = f"{self.base_url}/chat/completions"
        headers = {"Authorization": f"Bearer {self.api_key}"}
        payload = {
            "model": model,
            "messages": messages,
            "stream": stream,
            **kwargs
        }
        async with httpx.AsyncClient(timeout=None) as client:
            if not stream:
                response = await client.post(url, headers=headers, json=payload)
                return response.json()
            else:
                return client.stream("POST", url, headers=headers, json=payload)