app.py

from __future__ import annotations

import os
import socket
from pathlib import Path
from typing import Any

import gradio as gr

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = None

from evd_agent.conversation import ConversationManager
from evd_agent.explainability import (
    render_alert_banner,
    render_classification_panel,
)


def _load_environment_file() -> None:
    if load_dotenv is not None:
        load_dotenv()
        return

    env_path = Path(".env")
    if not env_path.exists():
        return

    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue

        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")

        if key and key not in os.environ:
            os.environ[key] = value


_load_environment_file()


def _require_llm_configuration() -> None:
    if not (
        os.getenv("EVD_LLM_API_KEY")
        or os.getenv("OPENAI_API_KEY")
    ):
        raise RuntimeError(
            "LLM credentials are required. Set EVD_LLM_API_KEY or OPENAI_API_KEY before starting the app."
        )


manager = ConversationManager(context_path=os.getenv("EVD_CONTEXT_PATH"))


def _resolve_server_port() -> int:
    preferred_port = int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860")))
    for port in range(preferred_port, preferred_port + 100):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                sock.bind(("0.0.0.0", port))
            except OSError:
                continue
            return port
    raise RuntimeError(f"Could not find a free port starting at {preferred_port}.")


def _history_to_chatbot(state) -> list[dict[str, Any]]:
    messages: list[dict[str, Any]] = []
    for turn in state.history:
        messages.append({"role": turn.role, "content": turn.content})
    return messages


def _context_panel() -> str:
    return manager.context_engine.context_summary()


def initialize_session():
    state, _ = manager.start_session()
    return (
        state,
        _history_to_chatbot(state),
        "",
        render_classification_panel(state.decision),
        _context_panel(),
    )


def submit_message(user_message: str, state):
    if state is None:
        state = manager.new_state()

    cleaned = (user_message or "").strip()
    if not cleaned:
        return (
            state,
            _history_to_chatbot(state),
            "",
            render_classification_panel(state.decision),
            _context_panel(),
            "",
        )

    result = manager.process_turn(state, cleaned)

    alert = render_alert_banner(result.decision)
    classification = render_classification_panel(result.decision)

    return (
        state,
        _history_to_chatbot(state),
        alert,
        classification,
        _context_panel(),
        "",
    )


def build_app() -> gr.Blocks:
    with gr.Blocks(title="EVD Clinical Screening Agent") as demo:
        gr.Markdown("# EVD Clinical Screening AI Agent")
        gr.Markdown(
            "Adaptive clinical reasoning assistant for rapid Ebola suspected/probable case screening. "
            "This tool is decision support and does not replace national case management protocols."
        )

        interview_state = gr.State()

        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(label="Clinical Interview", height=500)
                message_box = gr.Textbox(
                    label="Clinician Input",
                    placeholder="Enter findings, symptoms, exposures, travel, and context."
                )
                with gr.Row():
                    submit_btn = gr.Button("Send", variant="primary")
                    reset_btn = gr.Button("Reset Session")

            with gr.Column(scale=1):
                alert_banner = gr.Markdown(label="Alert")
                classification_panel = gr.Markdown(label="Classification")
                context_panel = gr.Markdown(label="Epidemiological Context")

        demo.load(
            initialize_session,
            inputs=[],
            outputs=[
                interview_state,
                chatbot,
                alert_banner,
                classification_panel,
                context_panel,
            ],
        )

        submit_btn.click(
            submit_message,
            inputs=[message_box, interview_state],
            outputs=[
                interview_state,
                chatbot,
                alert_banner,
                classification_panel,
                context_panel,
                message_box,
            ],
        )

        message_box.submit(
            submit_message,
            inputs=[message_box, interview_state],
            outputs=[
                interview_state,
                chatbot,
                alert_banner,
                classification_panel,
                context_panel,
                message_box,
            ],
        )

        reset_btn.click(
            initialize_session,
            inputs=[],
            outputs=[
                interview_state,
                chatbot,
                alert_banner,
                classification_panel,
                context_panel,
            ],
        )

    return demo


if __name__ == "__main__":
    _require_llm_configuration()
    app = build_app()
    running_in_space = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID"))
    app.launch(
        server_name="0.0.0.0",
        server_port=_resolve_server_port(),
        share=running_in_space,
        theme=gr.themes.Soft(),
    )