import pytest
from fastapi import FastAPI, HTTPException
from starlette.requests import Request
from starlette.responses import JSONResponse
from unittest.mock import MagicMock
import asyncio

from aare.daq.server_exception_handler import register_exception_handlers
from aare.common.exception_handler import (
    MountingFailed, WarningTellException, CriticalTellException,
    LoopCenteringFailed, TransformationInvalidException, BeamlineBusyException,
    AuthenticationException, SampleException, UserRightsException,
    SmargonCommunicationError, TellCommunicationError, JFJochCommunicationError,
    AerotechCommunicationError, DataCollectionException, RasterScanException,
    UnmountingFailed, AXCFailed, AareDBCommunicationError,
    MagnetPositionSensorErorr, ManualMountException, SmartMagnetFaultException,
    TellMountFailedException, TellCommandWhileBusyException, TellConnectionException
)

@pytest.fixture
def app():
    app = FastAPI()
    register_exception_handlers(app)
    return app

@pytest.fixture
def mock_request():
    return MagicMock(spec=Request)

@pytest.mark.asyncio
async def test_http_exception_handler(app, mock_request):
    exc = HTTPException(status_code=418, detail="I'm a teapot")
    handler = app.exception_handlers[HTTPException]
    response = await handler(mock_request, exc)
    assert response.status_code == 418
    assert response.body == b'{"code":"HTTP_ERROR","message":"I\'m a teapot"}'

@pytest.mark.asyncio
async def test_mounting_failed_handler(app, mock_request):
    exc = MountingFailed("Mount failed")
    handler = app.exception_handlers[MountingFailed]
    response = await handler(mock_request, exc)
    assert response.status_code == 404
    assert b"MOUNTING_FAILED" in response.body

@pytest.mark.asyncio
async def test_unmounting_failed_handler(app, mock_request):
    exc = UnmountingFailed("Unmount failed")
    handler = app.exception_handlers[UnmountingFailed]
    response = await handler(mock_request, exc)
    assert response.status_code == 500
    assert b"UNMOUNTING_FAILED" in response.body

@pytest.mark.asyncio
async def test_smargon_comm_handler(app, mock_request):
    exc = SmargonCommunicationError("Smargon dead", operation="MOVE", endpoint="/move")
    handler = app.exception_handlers[SmargonCommunicationError]
    response = await handler(mock_request, exc)
    assert response.status_code == 503
    assert b"SMARGON_UNAVAILABLE" in response.body
    assert b"MOVE" in response.body

@pytest.mark.asyncio
async def test_unhandled_exception_handler(app, mock_request):
    exc = ValueError("Something went wrong")
    handler = app.exception_handlers[Exception]
    response = await handler(mock_request, exc)
    assert response.status_code == 500
    assert b"INTERNAL_SERVER_ERROR" in response.body

@pytest.mark.asyncio
async def test_all_handlers_and_payloads(app, mock_request):
    # Test all registered handlers to ensure they return JSONResponse and cover the code
    for exc_class, handler in app.exception_handlers.items():
        if exc_class in (HTTPException, Exception, Request):
            continue
            
        # Try to instantiate the exception
        try:
            # Some exceptions might need specific args, but most in aare.common.exception_handler 
            # have defaults or take a message
            if exc_class in (SmargonCommunicationError, AareDBCommunicationError, TellCommunicationError, 
                             JFJochCommunicationError, AerotechCommunicationError):
                exc = exc_class("error", operation="OP", endpoint="/EP")
            else:
                exc = exc_class("error")
                
            response = await handler(mock_request, exc)
            assert isinstance(response, JSONResponse)
            assert response.status_code != 200
        except Exception as e:
            print(f"Skipping {exc_class} due to {e}")

    # Special case for HTTPException with dict detail
    exc = HTTPException(status_code=400, detail={"code": "CUSTOM", "message": "Msg"})
    handler = app.exception_handlers[HTTPException]
    response = await handler(mock_request, exc)
    assert response.status_code == 400
    assert b"CUSTOM" in response.body