main.py

"""
VNC-Perplexity API Server
FastAPI server with OpenAI and Anthropic compatible endpoints
"""
import asyncio
import logging
import signal
import sys
import time
import uuid
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Dict, List, Optional

import uvicorn
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from core.config import config
from core.vnc_manager import VNCManager
from core.response_processor import ResponseProcessor
from core.api_models import (
    OpenAIChatRequest, OpenAIChatResponse, OpenAIChoice, OpenAIMessage, OpenAIUsage,
    AnthropicRequest, AnthropicResponse,
    ModelListResponse, ModelInfo,
    HealthResponse, StatusResponse, VNCStatus,
    ErrorResponse, ErrorDetail
)


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(f'logs/server_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


# Global managers
vnc_manager = None
response_processor = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifespan"""
    global vnc_manager, response_processor
    
    # Startup
    logger.info("🚀 Starting VNC-Perplexity API Server...")
    
    try:
        # Initialize VNC manager
        vnc_manager = VNCManager()
        if not await vnc_manager.initialize():
            logger.error("❌ Failed to initialize VNC manager")
        else:
            logger.info("✅ VNC manager initialized")
        
        # Initialize response processor
        response_processor = ResponseProcessor()
        logger.info("✅ Response processor initialized")
        
        # Store in app state
        app.state.vnc_manager = vnc_manager
        app.state.response_processor = response_processor
        
        logger.info("✅ Server startup completed successfully")
        
    except Exception as e:
        logger.error(f"❌ Startup error: {e}")
    
    yield
    
    # Shutdown
    logger.info("🛑 Shutting down VNC-Perplexity API Server...")
    
    try:
        # Cleanup VNC environment
        if vnc_manager:
            await vnc_manager.cleanup()
        
        # Cleanup response processor
        if response_processor:
            await response_processor._cleanup_old_responses()
        
        logger.info("✅ Shutdown cleanup completed")
    except Exception as e:
        logger.error(f"❌ Shutdown cleanup failed: {e}")


# Create FastAPI app with lifespan
app = FastAPI(
    title="VNC-Perplexity API Server",
    description="OpenAI and Anthropic compatible API using VNC browser automation with Perplexity.ai",
    version="1.0.0",
    lifespan=lifespan
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global state
server_start_time = datetime.now()
request_counter = 0


# Middleware for request logging
@app.middleware("http")
async def log_requests(request: Request, call_next):
    global request_counter
    request_counter += 1
    
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    
    logger.info(f"📝 {request.method} {request.url.path} - {response.status_code} - {process_time:.3f}s")
    return response


# Error handlers
@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    """Handle HTTP exceptions"""
    error_response = ErrorResponse(
        error=exc.detail,
        code=str(exc.status_code)
    )
    return JSONResponse(
        status_code=exc.status_code,
        content=error_response.model_dump()
    )


@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """Handle general exceptions"""
    logger.error(f"❌ Unhandled error: {exc}")
    error_response = ErrorResponse(
        error="An internal server error occurred",
        detail=str(exc),
        code="500"
    )
    return JSONResponse(
        status_code=500,
        content=error_response.model_dump()
    )


# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    global vnc_manager
    
    vnc_status = await vnc_manager.get_status() if vnc_manager else {}
    
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "vnc_running": vnc_status.get("vnc_running", False),
        "workers_active": vnc_status.get("active_workers", 0),
        "uptime_seconds": int((datetime.now() - server_start_time).total_seconds())
    }


# Models endpoint (OpenAI compatible)
@app.get("/v1/models", response_model=ModelListResponse)
async def list_models():
    """List available models (OpenAI compatible)"""
    models = [
        ModelInfo(
            id="perplexity-pro",
            object="model",
            created=int(server_start_time.timestamp()),
            owned_by="perplexity"
        ),
        ModelInfo(
            id="perplexity-fast",
            object="model", 
            created=int(server_start_time.timestamp()),
            owned_by="perplexity"
        )
    ]
    
    return ModelListResponse(
        object="list",
        data=models
    )


# Status endpoint
@app.get("/status")
async def get_status():
    """Get server status"""
    global vnc_manager, response_processor
    
    # Get VNC status
    vnc_status = await vnc_manager.get_status() if vnc_manager else {}
    
    # Get response processor stats
    processor_stats = await response_processor.get_stats() if response_processor else {}
    
    return {
        "server_time": datetime.now().isoformat(),
        "uptime_seconds": int((datetime.now() - server_start_time).total_seconds()),
        "requests_processed": request_counter,
        "vnc": vnc_status,
        "response_processor": processor_stats
    }


# OpenAI Chat Completions endpoint
@app.post("/v1/chat/completions", response_model=OpenAIChatResponse)
async def create_chat_completion(request: OpenAIChatRequest):
    """Create chat completion (OpenAI compatible)"""
    global vnc_manager, response_processor
    
    if not vnc_manager or not response_processor:
        raise HTTPException(
            status_code=503,
            detail="VNC manager or response processor not initialized"
        )
    
    try:
        # Extract the user message
        user_messages = [msg for msg in request.messages if msg.role == "user"]
        if not user_messages:
            raise HTTPException(status_code=400, detail="No user message found")
        
        query = user_messages[-1].content
        
        # Process with VNC using configured mode
        if config.use_download_mode:
            result = await vnc_manager.process_query_with_download(query)
        else:
            result = await vnc_manager.process_query(query)
            
        if result.status != "success":
            # Try to get response from saved MD files as fallback
            logger.info("🔄 Trying to get response from saved MD files...")
            latest_response = await response_processor.get_latest_response_from_files()
            if latest_response:
                logger.info("✅ Found recent response in MD files")
                if config.use_pattern_extraction:
                    processed_response = await response_processor.process_response_with_extraction(latest_response)
                else:
                    processed_response = await response_processor.process_response(latest_response)
            else:
                raise HTTPException(status_code=500, detail=f"VNC processing failed: {result.error}")
        else:
            # Process response with configured method
            if config.use_pattern_extraction:
                processed_response = await response_processor.process_response_with_extraction(result.response)
            else:
                processed_response = await response_processor.process_response(result.response)
        
        # Create OpenAI compatible response
        choice = OpenAIChoice(
            index=0,
            message=OpenAIMessage(
                role="assistant",
                content=processed_response
            ),
            finish_reason="stop"
        )
        
        usage = OpenAIUsage(
            prompt_tokens=len(query.split()),
            completion_tokens=len(processed_response.split()),
            total_tokens=len(query.split()) + len(processed_response.split())
        )
        
        return OpenAIChatResponse(
            id=f"chatcmpl-{uuid.uuid4().hex[:12]}",
            object="chat.completion",
            created=int(time.time()),
            model=request.model,
            choices=[choice],
            usage=usage
        )
        
    except Exception as e:
        logger.error(f"❌ Chat completion error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# Anthropic Messages endpoint
@app.post("/v1/messages", response_model=AnthropicResponse)
async def create_message(request: AnthropicRequest):
    """Create message (Anthropic compatible)"""
    global vnc_manager, response_processor
    
    if not vnc_manager or not response_processor:
        raise HTTPException(
            status_code=503,
            detail="VNC manager or response processor not initialized"
        )
    
    try:
        # Extract the user message
        user_messages = [msg for msg in request.messages if msg.role == "user"]
        if not user_messages:
            raise HTTPException(status_code=400, detail="No user message found")
        
        query = user_messages[-1].content
        
        # Process with VNC using configured mode
        if config.use_download_mode:
            result = await vnc_manager.process_query_with_download(query)
        else:
            result = await vnc_manager.process_query(query)
            
        if result.status != "success":
            # Try to get response from saved MD files as fallback
            logger.info("🔄 Trying to get response from saved MD files...")
            latest_response = await response_processor.get_latest_response_from_files()
            if latest_response:
                logger.info("✅ Found recent response in MD files")
                if config.use_pattern_extraction:
                    processed_response = await response_processor.process_response_with_extraction(latest_response)
                else:
                    processed_response = await response_processor.process_response(latest_response)
            else:
                raise HTTPException(status_code=500, detail=f"VNC processing failed: {result.error}")
        else:
            # Process response with configured method
            if config.use_pattern_extraction:
                processed_response = await response_processor.process_response_with_extraction(result.response)
            else:
                processed_response = await response_processor.process_response(result.response)
        
        # Create Anthropic compatible response
        return AnthropicResponse(
            id=f"msg_{uuid.uuid4().hex[:12]}",
            type="message",
            role="assistant",
            content=[{"type": "text", "text": processed_response}],
            model=request.model,
            stop_reason="end_turn",
            stop_sequence=None,
            usage={
                "input_tokens": len(query.split()),
                "output_tokens": len(processed_response.split())
            }
        )
        
    except Exception as e:
        logger.error(f"❌ Message creation error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# Root endpoint
@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "VNC-Perplexity API Server",
        "version": "1.0.0",
        "status": "running",
        "uptime_seconds": int((datetime.now() - server_start_time).total_seconds()),
        "endpoints": {
            "health": "/health",
            "models": "/v1/models", 
            "status": "/status",
            "openai_chat": "/v1/chat/completions",
            "anthropic_messages": "/v1/messages"
        },
        "docs": "/docs"
    }


if __name__ == "__main__":
    logger.info("🚀 Starting VNC-Perplexity API Server...")
    
    def signal_handler(signum, frame):
        """Handle shutdown signals"""
        logger.info(f"🛑 Received signal {signum}, shutting down...")
        sys.exit(0)
    
    # Register signal handlers
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        uvicorn.run(
            "main:app",
            host=config.host,
            port=config.port,
            reload=False,  # Set to True for development
            log_level="info"
        )
    except KeyboardInterrupt:
        logger.info("🛑 Server interrupted by user")
    except Exception as e:
        logger.error(f"❌ Server error: {e}")
    finally:
        logger.info("✅ Server shutdown complete")