main.py

import os
import time
import logging
from typing import Optional, Dict, Any

import requests
from fastapi import FastAPI, HTTPException, Header, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

# ---------------------------------------------------------
# Logging Setup
# ---------------------------------------------------------

LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()

logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
)

logger = logging.getLogger("assistant_proxy")

# ---------------------------------------------------------
# Environment config
# ---------------------------------------------------------

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_ASSISTANT_ID = os.getenv("OPENAI_ASSISTANT_ID")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1")  # used when creating runs
COPILOT_PROXY_API_KEY = os.getenv("COPILOT_PROXY_API_KEY")

# In-memory conversation store: conversation_id -> thread_id
conversation_to_thread: Dict[str, str] = {}

# ---------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------

app = FastAPI(
    title="Copilot Assistant Proxy",
    description="FastAPI endpoint that proxies Copilot REST calls to an OpenAI Assistant.",
    version="1.0.0",
)

# ---------------------------------------------------------
# Startup / Shutdown Logging
# ---------------------------------------------------------

@app.on_event("startup")
async def startup_event():
    logger.info("✅ FastAPI service starting up...")
    logger.info(f"OpenAI Assistant ID configured: {bool(OPENAI_ASSISTANT_ID)}")
    logger.info(f"OpenAI API Key present: {bool(OPENAI_API_KEY)}")
    logger.info(f"OpenAI Model: {OPENAI_MODEL}")
    logger.info(f"Copilot Proxy Auth Enabled: {bool(COPILOT_PROXY_API_KEY)}")
    logger.info("Startup complete and server is ready to accept requests ✅")


@app.on_event("shutdown")
async def shutdown_event():
    logger.info("🛑 FastAPI service shutting down gracefully...")


# ---------------------------------------------------------
# Request Logging Middleware
# ---------------------------------------------------------

@app.middleware("http")
async def log_requests(request: Request, call_next):
    logger.info(f"➡️  Incoming request: {request.method} {request.url.path}")

    try:
        response = await call_next(request)
    except HTTPException as e:
        # Let FastAPI handle HTTPExceptions, but log them
        logger.exception(f"⚠️ HTTPException: {e.detail}")
        raise
    except Exception as e:
        logger.exception(f"💥 Unhandled server error: {e}")
        return JSONResponse(
            status_code=500,
            content={
                "error": "Internal server error",
                "details": str(e),
            },
        )

    logger.info(f"⬅️  Response: {request.method} {request.url.path} → {response.status_code}")
    return response


# ---------------------------------------------------------
# Models
# ---------------------------------------------------------

class ChatRequest(BaseModel):
    prompt: str = Field(..., max_length=8000)
    conversation_id: Optional[str] = None
    system_instructions: Optional[str] = Field(None, max_length=4000)
    temperature: Optional[float] = Field(0.7, ge=0.0, le=2.0)
    max_tokens: Optional[int] = Field(512, ge=1, le=4096)
    metadata: Optional[Dict[str, str]] = None
    user_id: Optional[str] = None


class Usage(BaseModel):
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    total_tokens: Optional[int] = None


class ChatResponse(BaseModel):
    reply: str
    conversation_id: Optional[str] = None
    finish_reason: Optional[str] = None
    usage: Optional[Usage] = None
    raw_model_response: Optional[Dict[str, Any]] = None


class ErrorResponse(BaseModel):
    code: str
    message: str
    request_id: Optional[str] = None
    details: Optional[Dict[str, Any]] = None


# ---------------------------------------------------------
# Helper functions for Assistants API (OpenAI platform, v2)
# ---------------------------------------------------------

def _headers() -> Dict[str, str]:
    """
    Standard headers for OpenAI Assistants v2 on platform.openai.com.
    """
    return {
        "Authorization": f"Bearer {OPENAI_API_KEY}",
        "Content-Type": "application/json",
        # Required to use Assistants v2 on the public platform
        "OpenAI-Beta": "assistants=v2",
    }


def _compose_user_content(prompt: str, system_instructions: Optional[str]) -> str:
    """
    We assume the assistant's main instructions live on the assistant itself.
    If system_instructions are provided at call time, prepend them to the prompt.
    """
    if system_instructions:
        return f"System instructions:\n{system_instructions}\n\nUser request:\n{prompt}"
    return prompt


def _create_thread_with_message(content: str) -> str:
    logger.info("🧵 Creating new assistant thread...")
    url = f"{OPENAI_BASE_URL}/threads"
    payload = {
        "messages": [
            {
                "role": "user",
                "content": content,
            }
        ]
    }
    resp = requests.post(url, json=payload, headers=_headers(), timeout=30)

    if resp.status_code >= 400:
        logger.error(f"❌ Thread creation failed: {resp.status_code} {resp.text}")
        raise HTTPException(
            status_code=500,
            detail=ErrorResponse(
                code="openai_create_thread_error",
                message="Failed to create thread with OpenAI Assistants API.",
                details={"status_code": resp.status_code, "body": resp.text},
            ).dict(),
        )

    data = resp.json()
    thread_id = data["id"]
    logger.info(f"✅ Thread created: {thread_id}")
    return thread_id


def _add_message_to_thread(thread_id: str, content: str) -> None:
    logger.info(f"📝 Adding message to thread {thread_id}...")
    url = f"{OPENAI_BASE_URL}/threads/{thread_id}/messages"
    payload = {
        "role": "user",
        "content": content,
    }
    resp = requests.post(url, json=payload, headers=_headers(), timeout=30)

    if resp.status_code >= 400:
        logger.error(f"❌ Error adding message: {resp.status_code} {resp.text}")
        raise HTTPException(
            status_code=500,
            detail=ErrorResponse(
                code="openai_add_message_error",
                message="Failed to add message to thread.",
                details={"status_code": resp.status_code, "body": resp.text},
            ).dict(),
        )

    logger.info("✅ Message added to thread.")


def _create_run(thread_id: str) -> str:
    logger.info(f"🏃 Creating run for thread {thread_id}...")
    url = f"{OPENAI_BASE_URL}/threads/{thread_id}/runs"
    payload = {
        "assistant_id": OPENAI_ASSISTANT_ID,
        # Assistants v2 requires a model in the run payload
        "model": OPENAI_MODEL,
    }
    resp = requests.post(url, json=payload, headers=_headers(), timeout=30)

    if resp.status_code >= 400:
        logger.error(f"❌ Error creating run: {resp.status_code} {resp.text}")
        raise HTTPException(
            status_code=500,
            detail=ErrorResponse(
                code="openai_create_run_error",
                message="Failed to create run for assistant.",
                details={"status_code": resp.status_code, "body": resp.text},
            ).dict(),
        )

    run_id = resp.json()["id"]
    logger.info(f"✅ Run created: {run_id}")
    return run_id


def _wait_for_run(thread_id: str, run_id: str, timeout_seconds: int = 60) -> Dict[str, Any]:
    """
    Poll the run until it completes or times out.
    Returns the final run object.
    """
    logger.info(f"⏳ Waiting for run {run_id} to complete...")
    url = f"{OPENAI_BASE_URL}/threads/{thread_id}/runs/{run_id}"
    start = time.time()

    while True:
        resp = requests.get(url, headers=_headers(), timeout=30)
        if resp.status_code >= 400:
            logger.error(f"❌ Error polling run: {resp.status_code} {resp.text}")
            raise HTTPException(
                status_code=500,
                detail=ErrorResponse(
                    code="openai_poll_run_error",
                    message="Failed while polling run status.",
                    details={"status_code": resp.status_code, "body": resp.text},
                ).dict(),
            )

        run = resp.json()
        status = run.get("status")
        logger.info(f"Run {run_id} status: {status}")

        if status in ("completed", "failed", "cancelled", "expired"):
            return run

        if time.time() - start > timeout_seconds:
            logger.error(f"❌ Run {run_id} timed out")
            raise HTTPException(
                status_code=500,
                detail=ErrorResponse(
                    code="run_timeout",
                    message="Assistant run did not complete in time.",
                    details={"status": status},
                ).dict(),
            )

        time.sleep(1.0)


def _get_latest_assistant_message(thread_id: str) -> str:
    """
    Retrieve messages for a thread and return the latest assistant message text.
    """
    logger.info(f"📥 Fetching messages for thread {thread_id}...")
    url = f"{OPENAI_BASE_URL}/threads/{thread_id}/messages"
    resp = requests.get(url, headers=_headers(), timeout=30)
    if resp.status_code >= 400:
        logger.error(f"❌ Error fetching messages: {resp.status_code} {resp.text}")
        raise HTTPException(
            status_code=500,
            detail=ErrorResponse(
                code="openai_fetch_messages_error",
                message="Failed to fetch messages for thread.",
                details={"status_code": resp.status_code, "body": resp.text},
            ).dict(),
        )

    data = resp.json()
    messages = data.get("data", [])
    # Sort by created_at descending and pick first assistant message
    messages.sort(key=lambda m: m.get("created_at", 0), reverse=True)
    for m in messages:
        if m.get("role") == "assistant":
            parts = m.get("content", [])
            texts = []
            for part in parts:
                if part.get("type") == "text":
                    texts.append(part["text"]["value"])
            result = "\n".join(texts).strip()
            logger.info("✅ Retrieved latest assistant message.")
            return result

    logger.warning("⚠️ No assistant message found in messages.")
    return ""


# ---------------------------------------------------------
# Copilot-facing endpoint
# ---------------------------------------------------------

@app.post(
    "/v1/chat",
    response_model=ChatResponse,
    responses={
        400: {"model": ErrorResponse},
        401: {"model": ErrorResponse},
        429: {"model": ErrorResponse},
        500: {"model": ErrorResponse},
    },
)
async def chat_with_assistant(
    body: ChatRequest,
    authorization: Optional[str] = Header(
        None,
        description="Bearer token from Copilot / plugin configuration.",
    ),
):
    logger.info(f"📨 /v1/chat called (conversation_id={body.conversation_id})")

    # Auth check (optional but recommended)
    if COPILOT_PROXY_API_KEY:
        if not authorization or not authorization.startswith("Bearer "):
            logger.warning("🚫 Missing or invalid Authorization header")
            raise HTTPException(
                status_code=401,
                detail=ErrorResponse(
                    code="unauthorized",
                    message="Missing or invalid Authorization header.",
                ).dict(),
            )
        inbound_token = authorization.split(" ", 1)[1]
        if inbound_token != COPILOT_PROXY_API_KEY:
            logger.warning("🚫 Invalid bearer token")
            raise HTTPException(
                status_code=401,
                detail=ErrorResponse(
                    code="unauthorized",
                    message="Invalid bearer token.",
                ).dict(),
            )

    # Validate env setup
    if not OPENAI_API_KEY or not OPENAI_ASSISTANT_ID:
        logger.error("❌ Missing OpenAI configuration")
        raise HTTPException(
            status_code=500,
            detail=ErrorResponse(
                code="server_misconfigured",
                message="OPENAI_API_KEY or OPENAI_ASSISTANT_ID not configured.",
            ).dict(),
        )

    user_content = _compose_user_content(body.prompt, body.system_instructions)

    # Get or create thread for this conversation
    thread_id: Optional[str] = None
    if body.conversation_id:
        thread_id = conversation_to_thread.get(body.conversation_id)

    if not thread_id:
        thread_id = _create_thread_with_message(user_content)
        if body.conversation_id:
            conversation_to_thread[body.conversation_id] = thread_id
    else:
        _add_message_to_thread(thread_id, user_content)

    # Create run
    run_id = _create_run(thread_id)

    # Wait for completion
    run = _wait_for_run(thread_id, run_id)

    run_status = run.get("status")
    finish_reason = "stop" if run_status == "completed" else "error"

    if run_status != "completed":
        logger.error(f"❌ Run ended with non-completed status: {run_status}")
        raise HTTPException(
            status_code=500,
            detail=ErrorResponse(
                code="assistant_run_error",
                message=f"Run ended with status: {run_status}",
                details=run,
            ).dict(),
        )

    # Get latest assistant message
    reply_text = _get_latest_assistant_message(thread_id)
    if not reply_text:
        reply_text = "[No assistant reply was found in the thread messages.]"

    # Token usage – Assistants API might not return it; keep empty for now
    usage_obj = Usage(
        prompt_tokens=None,
        completion_tokens=None,
        total_tokens=None,
    )

    logger.info("✅ /v1/chat completed successfully")

    return ChatResponse(
        reply=reply_text,
        conversation_id=body.conversation_id,
        finish_reason=finish_reason,
        usage=usage_obj,
        raw_model_response=run,  # can be removed in production
    )