app.py

import json
import gradio as gr
import time
import logging
from typing import Dict, List, Tuple, Optional
from agents import judge_agent, opponent_agent, narrator_agent, set_api_key
from dotenv import load_dotenv
import os
from files_extraction import (
    extract_text_from_pdf_bytes,extract_text_from_docx_bytes, extract_text_from_txt_bytes)
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()

# ==================== CORE SYSTEM STATE ====================
class LawTrainingSystem:
    def __init__(self):
        self.session_active = False
        self.current_case = None
        self.student_side = None
        self.conversation_history = []
        self.processing_lock = False
        self.api_key_set = False
        
    def set_api_key(self, api_key: str):
        """Set the API key for the agents"""
        try:
            if api_key and api_key.strip():
                success = set_api_key(api_key)
                if success:
                    self.api_key_set = True
                    logger.info("API key set successfully in system")
                    return "✅ API Key set successfully!"
                else:
                    self.api_key_set = False
                    logger.error("Failed to set API key")
                    return "❌ Invalid API key format"
            else:
                self.api_key_set = False
                logger.error("Empty API key provided")
                return "❌ Please enter a valid API key"
        except Exception as e:
            logger.error(f"Error setting API key: {e}", exc_info=True)
            return f"❌ Error setting API key: {str(e)}"
    
    def initialize_session(self, case_title: str, case_description: str, student_side: str):
        """Initialize a new training session"""
        try:
            self.session_active = True
            self.current_case = {"title": case_title, "description": case_description}
            self.student_side = student_side
            self.conversation_history = []
            logger.info(f"Session initialized: {case_title} - Student side: {student_side}")
            return f"Session initialized: {case_title} - You are the {student_side}"
        except Exception as e:
            logger.error(f"Error initializing session: {e}", exc_info=True)
            raise
    
    def is_processing(self) -> bool:
        """Check if system is currently processing a response"""
        return self.processing_lock
    
    def set_processing(self, status: bool):
        """Set processing lock status"""
        self.processing_lock = status
        logger.info(f"Processing lock set to: {status}")

# Global system instance
system = LawTrainingSystem()

# ==================== AGENT RESPONSE HANDLERS ====================
def generate_judge_response(student_message: str, opponent_message: str = "") -> str:
    """Generate Judge AI response"""
    try:
        logger.info("Generating judge response")
        return judge_agent(student_message, opponent_message)
    except Exception as e:
        logger.error(f"Error in judge response: {e}", exc_info=True)
        return f"Judge: I'm having trouble processing your argument. Error: {str(e)}"

def generate_opponent_response(student_message: str) -> str:
    """Generate Opponent AI response"""
    try:
        logger.info("Generating opponent response")
        case_context = f"Case: {system.current_case['title']} - {system.current_case['description']}" if system.current_case else ""
        return opponent_agent(student_message, case_context)
    except Exception as e:
        logger.error(f"Error in opponent response: {e}", exc_info=True)
        return f"Opponent: I'm experiencing technical difficulties. Error: {str(e)}"

def generate_narrator_response(context: str, student_message: str = "", opponent_message: str = "") -> str:
    """Generate Narrator AI response - handles side character arguments too"""
    try:
        logger.info("Generating narrator response")
        return narrator_agent(context, student_message, opponent_message)
    except Exception as e:
        logger.error(f"Error in narrator response: {e}", exc_info=True)
        return f"Narrator: The courtroom atmosphere is tense as technical difficulties arise. Error: {str(e)}"

# ==================== UI RESPONSE HANDLERS ====================
def handle_student_response(student_message: str, judge_history: List, opponent_history: List, narrator_history: List) -> Tuple:
    """Process student message and generate AI responses"""
    
    try:
        logger.info(f"Processing student response: {student_message[:50]}...")
        
        # Initialize histories if they are None
        if judge_history is None:
            judge_history = []
        if opponent_history is None:
            opponent_history = []
        if narrator_history is None:
            narrator_history = []
        
        # Prevent multiple simultaneous processing
        if system.is_processing():
            logger.warning("System is already processing, skipping request")
            return student_message, judge_history, opponent_history, narrator_history, gr.update(interactive=False)
        
        if not system.session_active:
            logger.warning("No active session")
            return student_message, judge_history, opponent_history, narrator_history, gr.update(interactive=True)
        
        if not system.api_key_set:
            logger.warning("API key not set")
            return student_message, judge_history, opponent_history, narrator_history, gr.update(interactive=True)
        
        # Set processing lock
        system.set_processing(True)
        
        # Generate opponent response first
        logger.info("Generating opponent response")
        opponent_response = generate_opponent_response(student_message)
        
        # Generate judge response (evaluating both student and opponent)
        logger.info("Generating judge response")
        judge_response = generate_judge_response(student_message, opponent_response)
        
        # Generate narrator response (describing the courtroom scene)
        logger.info("Generating narrator response")
        narrator_response = generate_narrator_response("Arguments presented", student_message, opponent_response)
        
        # Update chat histories
        judge_history.append({"role": "assistant", "content": judge_response})
        opponent_history.append({"role": "assistant", "content": opponent_response})
        narrator_history.append({"role": "assistant", "content": narrator_response})
        
        logger.info("All responses generated successfully")
        
        # Clear student input and re-enable
        return "", judge_history, opponent_history, narrator_history, gr.update(interactive=True)
        
    except Exception as e:
        logger.error(f"Error in handle_student_response: {e}", exc_info=True)
        return student_message, judge_history, opponent_history, narrator_history, gr.update(interactive=True)
    finally:
        # Always release processing lock
        system.set_processing(False)

def start_case_session(case_title: str, case_description: str, student_side: str) -> Tuple:
    """Initialize a new case session"""
    try:
        logger.info(f"Starting case session: {case_title}")
        
        if not system.api_key_set:
            logger.warning("Cannot start session - API key not set")
            return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(interactive=False)
        
        if not case_title or not case_description:
            logger.warning("Cannot start session - missing case details")
            return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(interactive=False)
        
        # Initialize system
        init_message = system.initialize_session(case_title, case_description, student_side)
        
        # Generate initial context from Narrator
        logger.info("Generating initial narrator context")
        initial_context = generate_narrator_response(f"Case session begins: {case_title}")
        
        # Create initial chat histories
        judge_history = [{"role": "assistant", "content": "Judge: I am ready to evaluate your arguments based on legal precedent and constitutional law."}]
        opponent_side = 'prosecution' if student_side == 'defense' else 'defense'
        opponent_history = [{"role": "assistant", "content": f"Opponent: I will argue for the {opponent_side} side. Present your case."}]
        narrator_history = [{"role": "assistant", "content": initial_context}]
        
        logger.info("Case session started successfully")
        
        return (
            judge_history,
            opponent_history,
            narrator_history,
            gr.update(visible=True),
            gr.update(interactive=True)
        )
        
    except Exception as e:
        logger.error(f"Error starting case session: {e}", exc_info=True)
        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(interactive=False)

def upload_case_documents(files) -> str:
    """Handle case document uploads"""
    try:
        if not files:
            return "No documents uploaded"

        all_extracted_text = []
        document_data_for_agent = {"documents": {}}

        for idx, file_path in enumerate(files):
            # Read bytes from the file path
            with open(file_path, "rb") as f:
                file_bytes = f.read()

            # Determine extension
            file_extension = os.path.splitext(file_path)[1].lstrip(".").lower()

            # Extract text based on file type
            if file_extension == "pdf":
                text = extract_text_from_pdf_bytes(file_bytes)
            elif file_extension == "docx":
                text = extract_text_from_docx_bytes(file_bytes)
            elif file_extension == "pptx":
                text = extract_text_from_pptx_bytes(file_bytes)
            elif file_extension == "txt":
                text = extract_text_from_txt_bytes(file_bytes)
            else:
                text = ""

            # Collect results
            if text:
                all_extracted_text.append(text)
                document_data_for_agent["documents"][f"doc_{idx+1}"] = {
                    "title": os.path.basename(file_path),
                    "content": text
                }

        # Save JSON only once, after processing
        if document_data_for_agent["documents"]:
            json_filename = "document_data.json"
            with open(json_filename, "w", encoding="utf-8") as json_file:
                json.dump(document_data_for_agent, json_file, indent=4, ensure_ascii=False)

        # Log and return summary
        uploaded_names = [os.path.basename(p) for p in files]
        logger.info(f"Uploaded {len(files)} documents: {', '.join(uploaded_names)}")
        return f"Uploaded {len(files)} documents: {', '.join(uploaded_names)}"

    except Exception as e:
        logger.error("Error uploading documents", exc_info=True)
        return f"Error uploading documents: {e}"

# ==================== GRADIO INTERFACE ====================
def create_law_training_interface():
    """Create the main Gradio interface"""
    
    with gr.Blocks(
        title="Law Training Agentic System",
        theme=gr.themes.Soft(),
        css="""
        .chat-container { height: 400px; }
        .upload-section { border: 2px dashed #ccc; padding: 20px; margin: 10px 0; }
        .character-panel { border: 1px solid #ddd; border-radius: 8px; margin: 5px; }
        .student-panel { border: 2px solid #4CAF50; }
        """
    ) as interface:
        
        # Header
        gr.Markdown("# ⚖️ Law Training Agentic System")
        gr.Markdown("Practice your courtroom argumentation skills with AI-powered opponents, judges, and narrators")
        
        # ==================== API KEY SECTION ====================
        with gr.Row():
            with gr.Column():
                gr.Markdown("## 🔑 API Configuration")
                api_key_input = gr.Textbox(
                    label="Claude API Key",
                    placeholder="Enter your Anthropic Claude API key...",
                    type="password",
                    lines=1
                )
                api_key_btn = gr.Button("Set API Key", variant="secondary")
                api_status = gr.Textbox(
                    label="API Status",
                    interactive=False,
                    lines=1
                )
        
        # ==================== CASE SETUP SECTION ====================
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("## 📋 Case Setup")
                case_title = gr.Textbox(
                    label="Case Title",
                    placeholder="e.g., Freedom of Speech vs. Public Safety",
                    lines=1
                )
                case_description = gr.Textbox(
                    label="Case Description",
                    placeholder="Describe the legal scenario you want to practice...",
                    lines=3
                )
                student_side = gr.Radio(
                    choices=["prosecution", "defense"],
                    label="Your Role",
                    value="prosecution"
                )
                
            with gr.Column(scale=1):
                gr.Markdown("## 📁 Documents")
                file_upload = gr.File(
                    label="Upload Case Documents",
                    file_count="multiple",
                    file_types=[".pdf", ".txt", ".docx"],
                    type="filepath",   
                    elem_classes="upload-section"
                )
                upload_status = gr.Textbox(
                    label="Upload Status",
                    interactive=False,
                    lines=2
                )
        
        # Start session button
        start_btn = gr.Button("🚀 Start Training Session", variant="primary", size="lg")
        
        # ==================== COURTROOM INTERFACE ====================
        courtroom_section = gr.Group(visible=False)
        
        with courtroom_section:
            gr.Markdown("## 🏛️ Courtroom")
            
            # Four-panel layout for characters
            with gr.Row(equal_height=True):
                # Student Panel (Interactive)
                with gr.Column(scale=1, elem_classes="character-panel student-panel"):
                    gr.Markdown("### 🎓 **Student (You)**")
                    student_input = gr.Textbox(
                        label="Your Argument",
                        placeholder="Present your legal argument here...",
                        lines=4,
                        interactive=False
                    )
                    submit_btn = gr.Button("📝 Submit Argument", variant="primary")
                
                # Judge Panel (AI Response Only)
                with gr.Column(scale=1, elem_classes="character-panel"):
                    gr.Markdown("### ⚖️ **Judge**")
                    judge_chat = gr.Chatbot(
                        label="Judge Feedback",
                        elem_classes="chat-container",
                        show_label=False,
                        type='messages'
                    )
            
            with gr.Row(equal_height=True):
                # Opponent Panel (AI Response Only)
                with gr.Column(scale=1, elem_classes="character-panel"):
                    gr.Markdown("### 🥊 **Opponent**")
                    opponent_chat = gr.Chatbot(
                        label="Opponent Arguments",
                        elem_classes="chat-container",
                        show_label=False,
                        type='messages'
                    )
                
                # Narrator Panel (AI Response Only)
                with gr.Column(scale=1, elem_classes="character-panel"):
                    gr.Markdown("### 📖 **Narrator**")
                    narrator_chat = gr.Chatbot(
                        label="Courtroom Narration & Side Characters",
                        elem_classes="chat-container",
                        show_label=False,
                        type='messages'
                    )
        
        # ==================== EVENT HANDLERS ====================
        
        # API key handler
        api_key_btn.click(
            fn=lambda key: system.set_api_key(key),
            inputs=[api_key_input],
            outputs=[api_status]
        )
        
        # File upload handler
        file_upload.change(
            fn=upload_case_documents,
            inputs=[file_upload],
            outputs=[upload_status]
        )
        
        # Start session handler
        start_btn.click(
            fn=start_case_session,
            inputs=[case_title, case_description, student_side],
            outputs=[judge_chat, opponent_chat, narrator_chat, courtroom_section, student_input]
        )
        
        # Student response handler
        submit_btn.click(
            fn=handle_student_response,
            inputs=[student_input, judge_chat, opponent_chat, narrator_chat],
            outputs=[student_input, judge_chat, opponent_chat, narrator_chat, submit_btn]
        )
        
        # Enter key submission
        student_input.submit(
            fn=handle_student_response,
            inputs=[student_input, judge_chat, opponent_chat, narrator_chat],
            outputs=[student_input, judge_chat, opponent_chat, narrator_chat, submit_btn]
        )
    
    return interface

# ==================== MAIN EXECUTION ====================
if __name__ == "__main__":
    try:
        logger.info("Starting Law Training System")
        # Create and launch the interface
        app = create_law_training_interface()
        
        # Launch configuration
        app.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            debug=True
        )
    except Exception as e:
        logger.error(f"Failed to start application: {e}", exc_info=True)
        raise