import asyncio
from datetime import UTC, datetime
from openai import APIError, AsyncOpenAI

from models import (
    Message,
    SendMessageResponse,
    Session,
    SessionHistoryResponse,
    TokenUsage,
)
from storage import JsonSessionStorage


class ChatProviderError(Exception):
    pass


class ChatService:
    def __init__(
        self,
        storage: JsonSessionStorage,
        client: AsyncOpenAI,
        model: str,
    ) -> None:
        self.storage = storage
        self.client = client
        self.model = model
        self._locks: dict[str, asyncio.Lock] = {}

    async def generate_response(
        self, session_id: str, content: str
    ) -> SendMessageResponse:
        lock = self._locks.setdefault(session_id, asyncio.Lock())
        async with lock:
            user_created_at = datetime.now(UTC)
            session = await self.storage.read(session_id)
            api_messages = self._build_api_messages(session, content)

            try:
                completion = await self.client.chat.completions.create(
                    model=self.model,
                    messages=api_messages,  # type: ignore[arg-type]
                    extra_body={"thinking": {"type": "disabled"}},
                )
                assistant_content = completion.choices[0].message.content
            except (APIError, IndexError, AttributeError) as exc:
                raise ChatProviderError("Upstream chat request failed") from exc

            if assistant_content is None:
                raise ChatProviderError("Upstream chat provider returned an empty response")

            assistant_message = await self._save_exchange(
                session, content, user_created_at, assistant_content
            )
            usage = self._parse_usage(completion.usage)

            return SendMessageResponse(
                session_id=session_id,
                message=assistant_message,
                usage=usage,
            )

    async def get_session_history(self, session_id: str) -> SessionHistoryResponse:
        lock = self._locks.setdefault(session_id, asyncio.Lock())
        async with lock:
            session = await self.storage.read(session_id)
            return SessionHistoryResponse(
                session_id=session.session_id,
                system_prompt=session.system_prompt,
                created_at=session.created_at,
                updated_at=session.updated_at,
                messages=session.messages,
            )

    @staticmethod
    def _build_api_messages(
        session: Session, content: str
    ) -> list[dict[str, str]]:
        messages = [{"role": "system", "content": session.system_prompt}]
        messages.extend(
            {"role": message.role, "content": message.content}
            for message in session.messages
        )
        messages.append({"role": "user", "content": content})
        return messages

    async def _save_exchange(
        self,
        session: Session,
        user_content: str,
        user_created_at: datetime,
        assistant_content: str,
    ) -> Message:
        assistant_message = Message(
            role="assistant",
            content=assistant_content,
            created_at=datetime.now(UTC),
        )
        session.messages.extend(
            [
                Message(
                    role="user",
                    content=user_content,
                    created_at=user_created_at,
                ),
                assistant_message,
            ]
        )
        session.updated_at = assistant_message.created_at
        await self.storage.write(session)
        return assistant_message

    @staticmethod
    def _parse_usage(usage: object | None) -> TokenUsage | None:
        if usage is None:
            return None
        return TokenUsage(
            prompt_tokens=usage.prompt_tokens,  # type: ignore[attr-defined]
            completion_tokens=usage.completion_tokens,  # type: ignore[attr-defined]
            total_tokens=usage.total_tokens,  # type: ignore[attr-defined]
        )