models.py

import time
import uuid
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Literal, Union

# --- Models Moved from main.py ---

# Input Models (OpenAI-like)
class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str

class ChatCompletionRequest(BaseModel):
    messages: List[ChatMessage]
    model: str = "notion-proxy" # Model name can be passed, but we map to Notion's model
    stream: bool = False
    # Add other potential OpenAI params if needed, though they might not map directly
    # max_tokens: Optional[int] = None
    # temperature: Optional[float] = None
    # space_id and thread_id are now handled globally via environment variables
    notion_model: str = "anthropic-opus-4" # Default Notion model, can be overridden


# Notion Models
class NotionTranscriptConfigValue(BaseModel):
    type: str = "markdown-chat"
    model: str # e.g., "anthropic-opus-4"

class NotionTranscriptItem(BaseModel):
    type: Literal["config", "user", "markdown-chat"]
    value: Union[List[List[str]], str, NotionTranscriptConfigValue]

class NotionDebugOverrides(BaseModel):
    cachedInferences: Dict = Field(default_factory=dict)
    annotationInferences: Dict = Field(default_factory=dict)
    emitInferences: bool = False

class NotionRequestBody(BaseModel):
    traceId: str = Field(default_factory=lambda: str(uuid.uuid4()))
    spaceId: str
    transcript: List[NotionTranscriptItem]
    # threadId is removed, createThread will be set to true
    createThread: bool = True
    debugOverrides: NotionDebugOverrides = Field(default_factory=NotionDebugOverrides)
    generateTitle: bool = False
    saveAllThreadOperations: bool = True


# Output Models (OpenAI SSE)
class ChoiceDelta(BaseModel):
    content: Optional[str] = None

class Choice(BaseModel):
    index: int = 0
    delta: ChoiceDelta
    finish_reason: Optional[Literal["stop", "length"]] = None

class ChatCompletionChunk(BaseModel):
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4()}")
    object: str = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = "notion-proxy" # Or could reflect the underlying Notion model
    choices: List[Choice]


# --- Models for /v1/models Endpoint ---

class Model(BaseModel):
    id: str
    object: str = "model"
    created: int = Field(default_factory=lambda: int(time.time()))
    owned_by: str = "notion" # Or specify based on actual model origin if needed

class ModelList(BaseModel):
    object: str = "list"
    data: List[Model]