main.py

import os
import uuid
import json
import time
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 models import (
    ChatMessage, ChatCompletionRequest, NotionTranscriptConfigValue,
    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."""
    transcript = [
        NotionTranscriptItem(
            type="config",
            value=NotionTranscriptConfigValue(model=request_data.notion_model)
        )
    ]
    for message in request_data.messages:
        # Map 'assistant' role to 'markdown-chat', all others to 'user'
        if message.role == "assistant":
            # Notion uses "markdown-chat" for assistant replies in the transcript history
            transcript.append(NotionTranscriptItem(type="markdown-chat", value=message.content))
        else:  # Handles 'user', 'system', etc.
            content = message.content
            if isinstance(content, str):
                # Handle string content: Append one item, using default [[""]] for empty strings
                notion_value = [[content]] if content else [[""]]
                transcript.append(NotionTranscriptItem(type="user", value=notion_value))
            elif isinstance(content, list):
                # Handle list content: Append a SEPARATE item for each valid text part
                found_text_part = False
                for part in content:
                    # Check if part is a dict with type="text" and non-empty text
                    if isinstance(part, dict) and part.get("type") == "text":
                        text_content = part.get("text")
                        if isinstance(text_content, str) and text_content:
                            # Create and append a SEPARATE item for this text part
                            transcript.append(NotionTranscriptItem(type="user", value=[[text_content]]))
                            found_text_part = True
                # If the list was empty or had no valid text parts, append a default empty item to maintain behavior
                if not found_text_part:
                    print(f'Error: no valid input found: {message}')
                    transcript.append(NotionTranscriptItem(type="user", value=[[""]]))
            else:
                # Handle unexpected content types (e.g., None, int) by appending a default empty item
                transcript.append(NotionTranscriptItem(type="user", value=[[""]]))
                print(f'Error: no valid input found: {message}')

    # 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,zh-CN;q=0.8,zh;q=0.7,zh-TW;q=0.6,ja;q=0.5',
        'content-type': 'application/json',
        'notion-audit-log-platform': 'web',
        'notion-client-version': '23.13.0.3604', # 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
            async with client.stream("POST", NOTION_API_URL, json=notion_request_body.dict(), 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)