main.py

import os
import uuid
import json
import time
import random
import httpx
from fastapi import FastAPI, Request, HTTPException, Depends, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import StreamingResponse
from dotenv import load_dotenv
import secrets # Added for secure comparison
from datetime import datetime, timedelta, timezone # Explicit datetime imports
from zoneinfo import ZoneInfo # For timezone handling
from models import (
    ChatMessage, ChatCompletionRequest, NotionTranscriptConfigValue,
    NotionTranscriptContextValue, NotionTranscriptItem, NotionDebugOverrides,
    NotionRequestBody, ChoiceDelta, Choice, ChatCompletionChunk, Model, ModelList
)

# Load environment variables from .env file
load_dotenv()

# --- Configuration ---
NOTION_API_URL = "https://www.notion.so/api/v3/runInferenceTranscript"
# IMPORTANT: Load the Notion cookie securely from environment variables
NOTION_COOKIE = os.getenv("NOTION_COOKIE")

NOTION_SPACE_ID = os.getenv("NOTION_SPACE_ID")
if not NOTION_COOKIE:
    print("Error: NOTION_COOKIE environment variable not set.")
    # Consider raising HTTPException or exiting in a real app
if not NOTION_SPACE_ID:
    print("Warning: NOTION_SPACE_ID environment variable not set. Using a default UUID.")
    # Using a default might not be ideal, depends on Notion's behavior
    # Consider raising an error instead: raise ValueError("NOTION_SPACE_ID not set")
    NOTION_SPACE_ID = str(uuid.uuid4()) # Default or raise error

# --- Authentication ---
EXPECTED_TOKEN = os.getenv("PROXY_AUTH_TOKEN", "default_token") # Default token
security = HTTPBearer()

def authenticate(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """Compares provided token with the expected token."""
    correct_token = secrets.compare_digest(credentials.credentials, EXPECTED_TOKEN)
    if not correct_token:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication credentials",
            # WWW-Authenticate header removed for Bearer
        )
    return True # Indicate successful authentication

# --- FastAPI App ---
app = FastAPI()

# --- Helper Functions ---

def build_notion_request(request_data: ChatCompletionRequest) -> NotionRequestBody:
    """Transforms OpenAI-style messages to Notion transcript format, adding userId and createdAt."""
    
    # --- Timestamp and User ID Logic ---
    user_id = os.getenv("NOTION_ACTIVE_USER_HEADER")
    # Get all non-assistant messages to assign timestamps
    non_assistant_messages = [msg for msg in request_data.messages if msg.role != "assistant"]
    num_non_assistant_messages = len(non_assistant_messages)
    message_timestamps = {} # Store timestamps keyed by message id

    if num_non_assistant_messages > 0:
        # Get current time specifically in Pacific Time (America/Los_Angeles)
        pacific_tz = ZoneInfo("America/Los_Angeles")
        now_pacific = datetime.now(timezone.utc).astimezone(pacific_tz)

        # Assign timestamp to the last non-assistant message
        last_msg_id = non_assistant_messages[-1].id
        message_timestamps[last_msg_id] = now_pacific

        # Calculate timestamps for previous non-assistant messages (random intervals earlier)
        current_timestamp = now_pacific
        for i in range(num_non_assistant_messages - 2, -1, -1): # Iterate backwards from second-to-last
            current_timestamp -= timedelta(minutes=random.randint(3, 20)) # Use random interval (3-20 mins)
            message_timestamps[non_assistant_messages[i].id] = current_timestamp

    # --- Build Transcript ---
    # Get current time in Pacific timezone for context
    pacific_tz = ZoneInfo("America/Los_Angeles")
    now_pacific = datetime.now(timezone.utc).astimezone(pacific_tz)
    # Format timestamp exactly as YYYY-MM-DDTHH:MM:SS.fff-HH:MM
    dt_str = now_pacific.strftime("%Y-%m-%dT%H:%M:%S")
    ms = f"{now_pacific.microsecond // 1000:03d}" # Ensure 3 digits for milliseconds
    tz_str = now_pacific.strftime("%z") # Gets +HHMM or -HHMM
    formatted_tz = f"{tz_str[:-2]}:{tz_str[-2:]}" # Insert colon
    current_datetime_iso = f"{dt_str}.{ms}{formatted_tz}"
    
    # Generate random text for userName and spaceName
    random_words = ["Project", "Workspace", "Team", "Studio", "Lab", "Hub", "Zone", "Space"]
    user_name = f"User{random.randint(100, 999)}"
    space_name = f"{random.choice(random_words)} {random.randint(1, 99)}"
    
    transcript = [
        NotionTranscriptItem(
            type="config",
            value=NotionTranscriptConfigValue(model=request_data.notion_model)
        ),
        NotionTranscriptItem(
            type="context",
            value=NotionTranscriptContextValue(
                userId=user_id or "",  # Use the user_id from env or empty string
                spaceId=NOTION_SPACE_ID,
                surface="home_module",
                timezone="America/Los_Angeles",
                userName=user_name,
                spaceName=space_name,
                spaceViewId=str(uuid.uuid4()),  # Random UUID for spaceViewId
                currentDatetime=current_datetime_iso
            )
        ),
        NotionTranscriptItem(
            type="agent-integration"
            # No value field needed for agent-integration
        )
    ]

    for message in request_data.messages:
        if message.role == "assistant":
            # Assistant messages get type="markdown-chat" and a traceId
            transcript.append(NotionTranscriptItem(
                type="markdown-chat",
                value=message.content,
                traceId=str(uuid.uuid4()) # Generate unique traceId for assistant message
            ))
        else: # Treat all other roles (user, system, etc.) as "user" type
            created_at_dt = message_timestamps.get(message.id) # Use the unified timestamp dict
            created_at_iso = None
            if created_at_dt:
                # Format timestamp exactly as YYYY-MM-DDTHH:MM:SS.fff-HH:MM
                dt_str = created_at_dt.strftime("%Y-%m-%dT%H:%M:%S")
                ms = f"{created_at_dt.microsecond // 1000:03d}" # Ensure 3 digits for milliseconds
                tz_str = created_at_dt.strftime("%z") # Gets +HHMM or -HHMM
                formatted_tz = f"{tz_str[:-2]}:{tz_str[-2:]}" # Insert colon
                created_at_iso = f"{dt_str}.{ms}{formatted_tz}"

            content = message.content
            # Ensure content is treated as a string for user/system messages
            if isinstance(content, list):
                 # Attempt to extract text from list format, default to empty string
                 text_content = ""
                 for part in content:
                     if isinstance(part, dict) and part.get("type") == "text":
                         text_part = part.get("text")
                         if isinstance(text_part, str):
                             text_content += text_part # Concatenate text parts if needed
                 content = text_content if text_content else "" # Use extracted text or empty string
            elif not isinstance(content, str):
                 content = "" # Default to empty string if not list or string

            # Format value as expected by Notion for user type: [[content_string]]
            notion_value = [[content]] if content else [[""]]

            transcript.append(NotionTranscriptItem(
                type="user", # Set type to "user" for non-assistant roles
                value=notion_value,
                userId=user_id, # Assign userId
                createdAt=created_at_iso # Assign timestamp
                # No traceId for user/system messages
            ))

    # Use globally configured spaceId, set createThread=True
    return NotionRequestBody(
        spaceId=NOTION_SPACE_ID, # From environment variable
        transcript=transcript,
        createThread=True,       # Always create a new thread
        # Generate a new traceId for each request
        traceId=str(uuid.uuid4()),
        # Explicitly set debugOverrides, generateTitle, and saveAllThreadOperations
        debugOverrides=NotionDebugOverrides(
            cachedInferences={},
            annotationInferences={},
            emitInferences=False
        ),
        generateTitle=False,
        saveAllThreadOperations=False
    )

async def stream_notion_response(notion_request_body: NotionRequestBody):
    """Streams the request to Notion and yields OpenAI-compatible SSE chunks."""
    headers = {
        'accept': 'application/x-ndjson',
        'accept-language': 'en-US,en;q=0.9',
        'content-type': 'application/json',
        'notion-audit-log-platform': 'web',
        'notion-client-version': '23.13.0.3668', # Consider making this configurable
        'origin': 'https://www.notion.so',
        'priority': 'u=1, i',
        # Referer might be optional or need adjustment. Removing threadId part.
        'referer': 'https://www.notion.so/chat',
        'sec-ch-ua': '"Chromium";v="136", "Google Chrome";v="136", "Not.A/Brand";v="99"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Windows"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-origin',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36',
        'cookie': NOTION_COOKIE, # Loaded from .env
        'x-notion-space-id': NOTION_SPACE_ID # Added space ID header
    }

    # Conditionally add the active user header
    notion_active_user = os.getenv("NOTION_ACTIVE_USER_HEADER")
    if notion_active_user: # Checks for None and empty string implicitly
        headers['x-notion-active-user-header'] = notion_active_user

    chunk_id = f"chatcmpl-{uuid.uuid4()}"
    created_time = int(time.time())

    try:
        async with httpx.AsyncClient(timeout=None) as client: # No timeout for streaming
            # Explicitly serialize using .json() to respect Pydantic Config (like json_encoders for UUID)
            request_body_json = notion_request_body.json()
            async with client.stream("POST", NOTION_API_URL, content=request_body_json, headers=headers) as response:
                if response.status_code != 200:
                    error_content = await response.aread()
                    print(f"Error from Notion API: {response.status_code}")
                    print(f"Response: {error_content.decode()}")
                    # Yield an error message in SSE format? Or just raise exception?
                    # For now, raise internal server error in the endpoint
                    raise HTTPException(status_code=response.status_code, detail=f"Notion API Error: {error_content.decode()}")

                async for line in response.aiter_lines():
                    if not line.strip():
                        continue
                    try:
                        data = json.loads(line)
                        # Check if it's the type of message containing text chunks
                        if data.get("type") == "markdown-chat" and isinstance(data.get("value"), str):
                            content_chunk = data["value"]
                            if content_chunk: # Only send if there's content
                                chunk = ChatCompletionChunk(
                                    id=chunk_id,
                                    created=created_time,
                                    choices=[Choice(delta=ChoiceDelta(content=content_chunk))]
                                )
                                yield f"data: {chunk.json()}\n\n"
                        # Add logic here to detect the end of the stream if Notion has a specific marker
                        # For now, we assume markdown-chat stops when the main content is done.
                        # If we see a recordMap, it's definitely past the text stream.
                        elif "recordMap" in data:
                             print("Detected recordMap, stopping stream.")
                             break # Stop processing after recordMap

                    except json.JSONDecodeError:
                        print(f"Warning: Could not decode JSON line: {line}")
                    except Exception as e:
                        print(f"Error processing line: {line} - {e}")
                        # Decide if we should continue or stop

        # Send the final chunk indicating stop
        final_chunk = ChatCompletionChunk(
            id=chunk_id,
            created=created_time,
            choices=[Choice(delta=ChoiceDelta(), finish_reason="stop")]
        )
        yield f"data: {final_chunk.json()}\n\n"
        yield "data: [DONE]\n\n"

    except httpx.RequestError as e:
        print(f"HTTPX Request Error: {e}")
        # Yield an error message or handle in the endpoint
        # For now, let the endpoint handle it
        raise HTTPException(status_code=500, detail=f"Error connecting to Notion API: {e}")
    except Exception as e:
        print(f"Unexpected error during streaming: {e}")
        # Yield an error message or handle in the endpoint
        raise HTTPException(status_code=500, detail=f"Internal server error during streaming: {e}")


# --- API Endpoint ---

@app.get("/v1/models", response_model=ModelList)
async def list_models(authenticated: bool = Depends(authenticate)):
    """
    Endpoint to list available Notion models, mimicking OpenAI's /v1/models.
    """
    available_models = [
        "openai-gpt-4.1",
        "anthropic-opus-4",
        "anthropic-sonnet-4"
    ]
    model_list = [
        Model(id=model_id, owned_by="notion") # created uses default_factory
        for model_id in available_models
    ]
    return ModelList(data=model_list)
@app.post("/v1/chat/completions")
async def chat_completions(request_data: ChatCompletionRequest, request: Request, authenticated: bool = Depends(authenticate)):
    """
    Endpoint to mimic OpenAI's chat completions, proxying to Notion.
    """
    if not NOTION_COOKIE:
         raise HTTPException(status_code=500, detail="Server configuration error: Notion cookie not set.")

    notion_request_body = build_notion_request(request_data)

    if request_data.stream:
        return StreamingResponse(
            stream_notion_response(notion_request_body),
            media_type="text/event-stream"
        )
    else:
        # --- Non-Streaming Logic (Optional - Collects stream internally) ---
        # Note: The primary goal is streaming, but a non-streaming version
        # might be useful for testing or simpler clients.
        # This requires collecting all chunks from the async generator.
        full_response_content = ""
        final_finish_reason = None
        chunk_id = f"chatcmpl-{uuid.uuid4()}" # Generate ID for the non-streamed response
        created_time = int(time.time())

        try:
            async for line in stream_notion_response(notion_request_body):
                 if line.startswith("data: ") and "[DONE]" not in line:
                    try:
                        data_json = line[len("data: "):].strip()
                        if data_json:
                            chunk_data = json.loads(data_json)
                            if chunk_data.get("choices"):
                                delta = chunk_data["choices"][0].get("delta", {})
                                content = delta.get("content")
                                if content:
                                    full_response_content += content
                                finish_reason = chunk_data["choices"][0].get("finish_reason")
                                if finish_reason:
                                    final_finish_reason = finish_reason
                    except json.JSONDecodeError:
                        print(f"Warning: Could not decode JSON line in non-streaming mode: {line}")

            # Construct the final OpenAI-compatible non-streaming response
            return {
                "id": chunk_id,
                "object": "chat.completion",
                "created": created_time,
                "model": request_data.model, # Return the model requested by the client
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": full_response_content,
                        },
                        "finish_reason": final_finish_reason or "stop", # Default to stop if not explicitly set
                    }
                ],
                "usage": { # Note: Token usage is not available from Notion
                    "prompt_tokens": None,
                    "completion_tokens": None,
                    "total_tokens": None,
                },
            }
        except HTTPException as e:
             # Re-raise HTTP exceptions from the streaming function
             raise e
        except Exception as e:
            print(f"Error during non-streaming processing: {e}")
            raise HTTPException(status_code=500, detail="Internal server error processing Notion response")


# --- Uvicorn Runner ---
# Allows running with `python main.py` for simple testing,
# but `uvicorn main:app --reload` is recommended for development.
if __name__ == "__main__":
    import uvicorn
    print("Starting server. Access at http://127.0.0.1:7860")
    print("Ensure NOTION_COOKIE is set in your .env file or environment.")
    uvicorn.run(app, host="127.0.0.1", port=7860)