app.py

"""12 Angry Agents - Main Gradio Application.

AI-powered jury deliberation simulation where 11 AI agents + 1 human player
debate criminal cases. Uses smolagents CodeAgent with LlamaIndex-powered tools.
"""

import asyncio
from pathlib import Path

import gradio as gr
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Add project root to path for imports
import sys
sys.path.insert(0, str(Path(__file__).parent))

from core.game_state import GameState, GamePhase
from core.orchestrator import OrchestratorAgent
from case_db import CaseLoader, CriminalCase, CaseIndexFactory
from agents import load_juror_configs, SmolagentJuror
from ui.components import render_jury_box, render_vote_tally
from services.mcp import MCPService, init_mcp_service, register_mcp_tools
from services.tts import get_tts_service
from services.narration import JudgeNarration


# CSS file path for scalable styling
CSS_PATH = Path(__file__).parent / "ui" / "static" / "styles.css"


class JuryDeliberation:
    """Main game orchestrator - delegates to OrchestratorAgent."""

    def __init__(self):
        self.case_loader = CaseLoader()
        self.juror_configs = load_juror_configs()
        self.juror_agents: dict[str, SmolagentJuror] = {}
        self.orchestrator: OrchestratorAgent | None = None
        self.current_case: CriminalCase | None = None
        self.mcp_service: MCPService | None = None
        self._case_index = None  # Store reference for MCP service

    @property
    def game_state(self) -> GameState | None:
        """Get game state from orchestrator."""
        return self.orchestrator.game_state if self.orchestrator else None

    async def initialize_game(self, case_id: str | None = None) -> tuple[str, str, str, list, str, bytes | None, str]:
        """Initialize a new game.

        Returns:
            Tuple of (case_summary, evidence_html, jury_box_html, chat_history, vote_html, audio_bytes, transcript)
        """
        # Load case
        if case_id:
            self.current_case = self.case_loader.get_case(case_id)
        else:
            self.current_case = self.case_loader.get_random_case()

        if not self.current_case:
            return (
                "No cases available. Please add case files to case_db/cases/",
                "",
                render_jury_box(self.juror_configs),
                [],
                "",
                None,
                ""
            )

        # Create LlamaIndex for this case
        case_index = CaseIndexFactory.get_index(self.current_case)
        self._case_index = case_index  # Store reference for MCP service

        # Initialize juror agents with case index (skip player seat 7)
        self.juror_agents = {}
        for config in self.juror_configs:
            if config.archetype != "player":
                try:
                    agent = SmolagentJuror(
                        config,
                        case_index=case_index,
                        case=self.current_case
                    )
                    agent.set_initial_conviction(self.current_case)
                    self.juror_agents[config.juror_id] = agent
                except Exception as e:
                    print(f"Warning: Failed to initialize {config.name}: {e}")

        # Create orchestrator with agents
        self.orchestrator = OrchestratorAgent(
            juror_configs=self.juror_configs,
            juror_agents=self.juror_agents,
            case=self.current_case
        )
        self.orchestrator.state.phase = GamePhase.PRESENTATION

        # Initialize MCP service for external agent participation
        self.mcp_service = MCPService(
            orchestrator=self.orchestrator,
            case_index=self._case_index,
            juror_configs=self.juror_configs,
            juror_agents=self.juror_agents
        )
        init_mcp_service(self.mcp_service)

        # Format case info
        case_summary = f"""## {self.current_case.title}

**Year:** {self.current_case.year} | **Jurisdiction:** {self.current_case.jurisdiction}

**Charges:** {self.current_case.get_charges_text()}

---

{self.current_case.summary}
"""

        evidence_html = f"""### Evidence

{self.current_case.get_evidence_summary()}

### Witnesses

{self.current_case.get_witness_summary()}
"""

        jury_html = render_jury_box(self.juror_configs, self.game_state)

        # Generate judge narration
        transcript = JudgeNarration.case_introduction(self.current_case)
        audio_bytes = None
        tts = get_tts_service()
        if tts.is_available:
            audio_bytes = await tts.asynthesize(transcript, style="formal", use_cache=True)

        chat_history = [
            {"role": "assistant", "content": f"**Judge:** {transcript}"}
        ]

        vote_html = render_vote_tally(self.game_state)

        return case_summary, evidence_html, jury_html, chat_history, vote_html, audio_bytes, transcript

    def select_player_side(self, side: str) -> tuple[str, str, list]:
        """Player selects prosecution or defense side."""
        if not self.orchestrator:
            return "", "", []

        self.orchestrator.set_player_side(side)

        vote_html = render_vote_tally(self.game_state)
        jury_html = render_jury_box(self.juror_configs, self.game_state)

        side_text = "PROSECUTION (Guilty)" if side == "prosecute" else "DEFENSE (Not Guilty)"
        chat_update = [{"role": "user", "content": f"*You have chosen to argue for the {side_text}*"}]

        return vote_html, jury_html, chat_update

    async def run_deliberation_round(self, prev_tally: tuple | None = None) -> tuple[str, str, list[dict], bytes | None, str, tuple]:
        """Run a single round of AI juror deliberation using fair speaker queue.

        Args:
            prev_tally: Previous vote tally (guilty, not_guilty) for shift detection.

        Returns:
            Tuple of (vote_html, jury_html, new_messages, audio_bytes, transcript, new_tally)
        """
        if not self.orchestrator or not self.current_case:
            return "", "", [], None, "", (0, 0)

        new_messages = []

        # Use orchestrator's fair speaker selection
        results = await self.orchestrator.run_deliberation_round()

        # Format results for chat
        for result in results:
            turn = result.turn
            config = next(c for c in self.juror_configs if c.juror_id == turn.speaker_id)

            # Speaker's argument
            new_messages.append(
                {"role": "assistant", "content": f"**{config.emoji} {config.name}:** {turn.content}"}
            )

            # Show reasoning steps if available (smolagents feature)
            if result.reasoning_steps:
                reasoning_text = "\n".join([f"  {step}" for step in result.reasoning_steps])
                new_messages.append({
                    "role": "assistant",
                    "content": f"*{config.emoji} {config.name}'s reasoning:*\n```\n{reasoning_text}\n```"
                })

            # Show tool calls if any
            if result.tool_calls:
                tools_text = ", ".join(result.tool_calls)
                new_messages.append({
                    "role": "assistant",
                    "content": f"*Tools used: {tools_text}*"
                })

            # Vote changes
            for juror_id, old_vote, new_vote in result.vote_changes:
                other_config = next(c for c in self.juror_configs if c.juror_id == juror_id)
                vote_text = "GUILTY" if new_vote == "guilty" else "NOT GUILTY"
                new_messages.append(
                    {"role": "assistant", "content": f"*{other_config.emoji} {other_config.name} changes their vote to {vote_text}*"}
                )

        vote_html = render_vote_tally(self.game_state)
        jury_html = render_jury_box(self.juror_configs, self.game_state)

        # Check for significant vote shift and generate narration
        new_guilty, new_not_guilty = self.game_state.get_vote_tally()
        new_tally = (new_guilty, new_not_guilty)
        audio_bytes = None
        transcript = ""

        if prev_tally:
            old_g, old_ng = prev_tally
            shift = abs(new_guilty - old_g)
            if shift >= 2:
                direction = "conviction" if new_guilty > old_g else "acquittal"
                transcript = JudgeNarration.vote_shift(shift, direction, new_guilty, new_not_guilty)
                if transcript:
                    tts = get_tts_service()
                    if tts.is_available:
                        audio_bytes = await tts.asynthesize(transcript, style="dramatic", use_cache=False)
                    new_messages.append({"role": "assistant", "content": f"**Judge:** {transcript}"})

        return vote_html, jury_html, new_messages, audio_bytes, transcript, new_tally

    def make_player_argument(
        self,
        strategy: str,
        custom_text: str | None = None
    ) -> tuple[str, str, list[tuple]]:
        """Process player's argument using orchestrator."""
        if not self.orchestrator or not self.current_case:
            return "", "", []

        player_config = next(c for c in self.juror_configs if c.archetype == "player")

        # Create argument content
        if custom_text:
            content = custom_text
        else:
            strategy_templates = {
                "Challenge Evidence": "I'd like to point out some issues with the evidence presented...",
                "Question Witness Credibility": "Can we really trust the witness testimony here?",
                "Appeal to Reasonable Doubt": "Remember, we need proof beyond reasonable doubt. Do we have that?",
                "Present Alternative Theory": "What if there's another explanation for what happened?",
                "Call for Vote": "I think we should take a vote and see where everyone stands.",
            }
            content = strategy_templates.get(strategy, "I have some concerns about this case...")

        # Use orchestrator to process
        argument_type = strategy.lower().replace(" ", "_")
        result = self.orchestrator.process_player_argument(content, argument_type)

        new_messages = [{"role": "user", "content": f"**{player_config.emoji} You:** {content}"}]

        # Vote changes
        for juror_id, old_vote, new_vote in result.vote_changes:
            config = next(c for c in self.juror_configs if c.juror_id == juror_id)
            vote_text = "GUILTY" if new_vote == "guilty" else "NOT GUILTY"
            new_messages.append(
                {"role": "assistant", "content": f"*{config.emoji} {config.name} changes their vote to {vote_text}*"}
            )

        vote_html = render_vote_tally(self.game_state)
        jury_html = render_jury_box(self.juror_configs, self.game_state)

        return vote_html, jury_html, new_messages


# Create global game instance
game = JuryDeliberation()


def create_ui():
    """Create the Gradio UI."""

    with gr.Blocks() as demo:

        # State
        chat_history = gr.State([])

        # Header
        gr.HTML("<h1 class='jury-title'>12 ANGRY AGENTS</h1>", elem_id="app-title")
        gr.HTML("<p class='jury-subtitle'>AI-Powered Jury Deliberation</p>", elem_id="app-subtitle")

        with gr.Row(elem_id="main-row"):
            # Left Column: Jury Box
            with gr.Column(scale=1, elem_id="jury-column"):
                gr.Markdown("### The Jury")
                jury_box = gr.HTML(
                    render_jury_box(game.juror_configs),
                    elem_id="jury-box-container"
                )
                vote_tally = gr.HTML("", elem_id="vote-tally-container")

                # Judge narration section
                gr.Markdown("### Judge")
                with gr.Group(elem_id="judge-narration"):
                    judge_audio = gr.Audio(
                        label="",
                        autoplay=True,
                        elem_id="judge-audio",
                        visible=True
                    )
                    judge_transcript = gr.Textbox(
                        label="",
                        lines=2,
                        interactive=False,
                        elem_id="judge-transcript",
                        show_label=False
                    )

                gr.Markdown("### Your Side")
                with gr.Row(elem_id="side-selection-row"):
                    defend_btn = gr.Button(
                        "DEFEND\n(Not Guilty)",
                        variant="secondary",
                        elem_id="defend-btn"
                    )
                    prosecute_btn = gr.Button(
                        "PROSECUTE\n(Guilty)",
                        variant="primary",
                        elem_id="prosecute-btn"
                    )

            # Center Column: Deliberation
            with gr.Column(scale=2, elem_id="deliberation-column"):
                gr.Markdown("### Deliberation Room")
                deliberation_chat = gr.Chatbot(
                    label="Deliberation",
                    height=400,
                    show_label=False,
                    elem_id="deliberation-chat",
                )

                # Player input
                with gr.Row(elem_id="player-input-row"):
                    strategy_select = gr.Dropdown(
                        choices=[
                            "Challenge Evidence",
                            "Question Witness Credibility",
                            "Appeal to Reasonable Doubt",
                            "Present Alternative Theory",
                            "Call for Vote",
                        ],
                        label="Your Strategy",
                        value="Challenge Evidence",
                        elem_id="strategy-select",
                    )
                    speak_btn = gr.Button(
                        "Speak",
                        variant="primary",
                        elem_id="speak-btn",
                        elem_classes=["primary-btn"]
                    )

                custom_text = gr.Textbox(
                    label="Custom argument (optional)",
                    placeholder="Type your own argument here...",
                    max_lines=2,
                    elem_id="custom-argument",
                )

                with gr.Row(elem_id="action-buttons-row"):
                    pass_btn = gr.Button("Pass Turn", elem_id="pass-btn")
                    next_round_btn = gr.Button(
                        "Next Round (AI Speaks)",
                        variant="secondary",
                        elem_id="next-round-btn"
                    )

            # Right Column: Case File
            with gr.Column(scale=1, elem_id="case-column", elem_classes=["case-file"]):
                gr.Markdown("### Case File")
                case_summary = gr.Markdown(
                    "*Click 'Start New Case' to begin*",
                    elem_id="case-summary"
                )

                with gr.Accordion("Evidence & Witnesses", open=False, elem_id="evidence-accordion"):
                    evidence_display = gr.Markdown("", elem_id="evidence-display")

        # Start Game Button
        start_btn = gr.Button(
            "Start New Case",
            variant="primary",
            size="lg",
            elem_id="start-btn",
            elem_classes=["primary-btn"]
        )

        # State for tracking vote tally for shift detection
        vote_tally_state = gr.State(None)

        # Event handlers
        async def start_game():
            case_md, evidence_md, jury_html, chat, vote_html, audio_bytes, transcript = await game.initialize_game()
            initial_tally = game.game_state.get_vote_tally() if game.game_state else (0, 0)
            return case_md, evidence_md, jury_html, chat, vote_html, audio_bytes, transcript, initial_tally

        start_btn.click(
            fn=start_game,
            inputs=[],
            outputs=[case_summary, evidence_display, jury_box, deliberation_chat, vote_tally, judge_audio, judge_transcript, vote_tally_state]
        )

        def select_defend(chat):
            vote_html, jury_html, new_msgs = game.select_player_side("defend")
            return vote_html, jury_html, chat + new_msgs

        def select_prosecute(chat):
            vote_html, jury_html, new_msgs = game.select_player_side("prosecute")
            return vote_html, jury_html, chat + new_msgs

        defend_btn.click(
            fn=select_defend,
            inputs=[deliberation_chat],
            outputs=[vote_tally, jury_box, deliberation_chat]
        )

        prosecute_btn.click(
            fn=select_prosecute,
            inputs=[deliberation_chat],
            outputs=[vote_tally, jury_box, deliberation_chat]
        )

        async def run_ai_round(chat, prev_tally):
            vote_html, jury_html, new_msgs, audio_bytes, transcript, new_tally = await game.run_deliberation_round(prev_tally)
            return vote_html, jury_html, chat + new_msgs, audio_bytes, transcript, new_tally

        next_round_btn.click(
            fn=run_ai_round,
            inputs=[deliberation_chat, vote_tally_state],
            outputs=[vote_tally, jury_box, deliberation_chat, judge_audio, judge_transcript, vote_tally_state]
        )

        def player_speak(strategy, custom, chat):
            text = custom if custom else None
            vote_html, jury_html, new_msgs = game.make_player_argument(strategy, text)
            return vote_html, jury_html, chat + new_msgs, ""

        speak_btn.click(
            fn=player_speak,
            inputs=[strategy_select, custom_text, deliberation_chat],
            outputs=[vote_tally, jury_box, deliberation_chat, custom_text]
        )

        def pass_turn(chat):
            chat.append({"role": "user", "content": "*You pass your turn*"})
            return chat

        pass_btn.click(
            fn=pass_turn,
            inputs=[deliberation_chat],
            outputs=[deliberation_chat]
        )

        # Register MCP tools for external AI agent participation
        register_mcp_tools(demo)

        return demo


# Main entry point
if __name__ == "__main__":
    demo = create_ui()
    demo.launch(
        mcp_server=True,
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        css_paths=[CSS_PATH],
        theme=gr.themes.Base(),
    )