import logging
import secrets
import time
from collections.abc import Generator

from fastapi import APIRouter, Depends, FastAPI, Header, HTTPException, Response
from fastapi.responses import StreamingResponse

from progtc.server.code_execution_manager import CodeExecutionManager
from progtc.server.config import server_config
from progtc.types import (
    CodeTimeoutError,
    ExecuteCodeRequest,
    ExecuteCodeStepMessage,
    ExecuteCodeStepRequest,
    ExecuteCodeStepState,
    ExecuteCodeStepToolCalls,
    ToolCallResult,
)

logger = logging.getLogger(__name__)

code_execution_manager = CodeExecutionManager()


def authenticate(
    x_progtc_api_key: str = Header(),
) -> None:
    if not secrets.compare_digest(x_progtc_api_key, server_config.api_key):
        raise HTTPException(status_code=401, detail="Unauthorized")


app = FastAPI()


@app.get("/ping")
def ping() -> Response:
    return Response(status_code=200)


router = APIRouter(dependencies=[Depends(authenticate)])


@router.post("/tool-results")
def add_tool_result(tool_call_result: ToolCallResult) -> None:
    code_execution_manager.send_tool_result(tool_call_result)


@router.post("/execute-code")
def execute_code(body: ExecuteCodeRequest) -> StreamingResponse:
    def stream() -> Generator[str, None, None]:
        try:
            with code_execution_manager.run(body.code, body.tool_names) as process:
                t_start = time.monotonic()
                t_last_event = time.monotonic()

                # Yield tool calls to the client until the process
                # is done or the timeout is reached
                while (
                    not process.is_done()
                    and time.monotonic() - t_start
                    < server_config.code_execution_timeout
                ):
                    tool_call = process.next_tool_call()
                    if tool_call is not None:
                        yield f"data: {tool_call.model_dump_json()}\n\n"
                        t_last_event = time.monotonic()
                    if time.monotonic() - t_last_event > 15.0:
                        yield ": keep alive\n\n"
                        t_last_event = time.monotonic()

                # Yield the final result to the client.
                if process.is_done():
                    result = process.result()
                else:
                    result = CodeTimeoutError(
                        stdout="",
                        stderr="",
                    )
                yield f"data: {result.model_dump_json()}\n\n"
                logger.debug("stream completed successfully")
        except Exception:
            logger.exception("server error during streaming")
            raise

    return StreamingResponse(stream(), media_type="text/event-stream")


@router.post("/execute-code-step")
def execute_code_step(
    body: ExecuteCodeStepRequest,
) -> ExecuteCodeStepMessage:
    state = (
        ExecuteCodeStepState.decode(body.state)
        if body.state
        else ExecuteCodeStepState()
    )
    if body.tool_results:
        state.tool_results.extend(body.tool_results)
    with code_execution_manager.run(body.code, body.tool_names, state) as process:
        t_start = time.monotonic()

        # Return tool calls to the client until the process
        # is done or the timeout is reached
        while (
            not process.is_done()
            and time.monotonic() - t_start < server_config.code_execution_timeout
        ):
            tool_calls = process.next_tool_calls()
            if tool_calls:
                state.tool_call_hashes.extend(tc.call_hash() for tc in tool_calls)
                return ExecuteCodeStepToolCalls(
                    tool_calls=tool_calls,
                    state=state.encode(),
                )

        # Return the final result.
        if process.is_done():
            result = process.result()
        else:
            result = CodeTimeoutError(
                stdout="",
                stderr="",
            )
        return result


app.include_router(router)
