from datetime import datetime
from typing import AsyncIterator

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, Response

from chatkit.server import ChatKitServer, StreamingResult
from chatkit.types import (
    AssistantMessageContent,
    AssistantMessageItem,
    ThreadItemDoneEvent,
    ThreadMetadata,
    ThreadStreamEvent,
    UserMessageItem,
)

app = FastAPI()

class MyChatKitServer(ChatKitServer[dict]):
    async def respond(
        self,
        thread: ThreadMetadata,
        input_user_message: UserMessageItem | None,
        context: dict,
    ) -> AsyncIterator[ThreadStreamEvent]:
        # Your response logic here
        yield ThreadItemDoneEvent(
            item=AssistantMessageItem(
                thread_id=thread.id,
                id=self.store.generate_item_id("message", thread, context),
                created_at=datetime.now(),
                content=[AssistantMessageContent(text="Hello, world!")],
            ),
        )

server = MyChatKitServer(store=MyChatKitStore())

@app.post("/chatkit")
async def chatkit(request: Request):
    result = await server.process(await request.body(), context={})
    if isinstance(result, StreamingResult):
        return StreamingResponse(result, media_type="text/event-stream")
    return Response(content=result.json, media_type="application/json")

from chatkit.store import Store, NotFoundError
from chatkit.types import ThreadMetadata, ThreadItem, Page

class MyChatKitStore(Store[dict]):
    async def load_thread(self, thread_id: str, context: dict) -> ThreadMetadata:
        # Load thread from your database
        pass
    
    async def save_thread(self, thread: ThreadMetadata, context: dict) -> None:
        # Save thread to your database
        pass
    
    # Implement other required methods...