from fastapi import APIRouter
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'
)
logger = logging.getLogger(__name__)

def get_router():
    router = APIRouter()
    
    # Import sub-modules from the routes directory
    from .routes.auth import auth
    from .routes.patients import patients
    from .routes.pdf import pdf
    from .routes.appointments import router as appointments
    from .routes.messaging import router as messaging
    from .routes.txagent import router as txagent, get_txagent, ChatRequest
    from fastapi import Depends, HTTPException
    from core.security import get_current_user
    import asyncio
    from fastapi.responses import StreamingResponse
    
    # Include sub-routers with correct prefixes
    router.include_router(patients, tags=["patients"])  # Remove prefix since routes already have /patients
    router.include_router(auth, prefix="/auth", tags=["auth"])
    router.include_router(pdf, prefix="/patients", tags=["pdf"])  # Keep prefix for PDF routes
    router.include_router(appointments, prefix="/appointments", tags=["appointments"])
    router.include_router(messaging, tags=["messaging"])
    router.include_router(txagent, tags=["txagent"])
    
    # Add the chat-stream endpoint at the root level
    @router.post("/chat-stream")
    async def chat_stream(
        request: ChatRequest,
        current_user: dict = Depends(get_current_user)
    ):
        """Streaming chat endpoint for TxAgent"""
        import logging
        logger = logging.getLogger(__name__)
        
        try:
            if not any(role in current_user.get('roles', []) for role in ['doctor', 'admin']):
                raise HTTPException(status_code=403, detail="Only doctors and admins can use TxAgent")

            # Get TxAgent instance
            txagent = get_txagent()
            
            if txagent is None:
                # Fallback response if TxAgent is not available
                response_text = f"TxAgent is currently unavailable. Please try again later. Your message was: {request.message}"
            else:
                # Get real response from TxAgent
                try:
                    response_text = txagent.chat(
                        message=request.message,
                        temperature=request.temperature,
                        max_new_tokens=request.max_new_tokens
                    )
                except Exception as e:
                    logger.error(f"TxAgent chat failed: {e}")
                    response_text = f"Sorry, I encountered an error processing your request: {request.message}. Please try again."

            # Return streaming text without data: prefix
            async def generate():
                # Send the complete response as plain text
                yield response_text

            return StreamingResponse(
                generate(),
                media_type="text/plain",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive"
                }
            )
        except Exception as e:
            logger.error(f"Error in chat stream: {e}")
            raise HTTPException(status_code=500, detail="Failed to process chat stream")
    
    return router

# Export the router
api_router = get_router()