main.py

"""
Main Entry Point - Telco Troubleshooting Agentic Challenge
Phase 3 Evaluation: Autonomous agent execution with time constraints
"""

import os
import sys
import json
import time
import argparse
import logging
from pathlib import Path

# Add agent directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'agent'))

from agent.llm_engine import get_llm_engine
from agent.react_loop import get_react_agent
from agent.trace_logger import get_trace_logger
from agent.tools import get_tool_executor

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class TelcoAgent:
    """
    Main agent class for the Telco Troubleshooting Challenge
    """
    
    def __init__(self, server_url: str = "http://localhost:8000"):
        self.server_url = server_url
        self.llm_engine = None
        self.react_agent = None
        self.trace_logger = None
        self.tool_executor = None
        
        # Performance tracking
        self.start_time = None
        self.max_execution_time = 300  # 5 minutes for Phase 3
        
    def initialize(self) -> bool:
        """
        Initialize all components
        """
        try:
            logger.info("Initializing Telco Agent...")
            
            # Initialize LLM engine
            logger.info("Loading LLM engine...")
            self.llm_engine = get_llm_engine()
            
            # Initialize tool executor with server URL
            logger.info(f"Initializing tool executor for {self.server_url}")
            self.tool_executor = get_tool_executor()
            self.tool_executor.server_url = self.server_url
            
            # Initialize ReAct agent
            logger.info("Initializing ReAct agent...")
            self.react_agent = get_react_agent()
            
            # Initialize trace logger
            logger.info("Initializing trace logger...")
            self.trace_logger = get_trace_logger()
            
            logger.info("Telco Agent initialized successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize agent: {e}")
            return False
    
    def load_scenario(self, scenario_file: str) -> Dict:
        """
        Load scenario from file
        """
        try:
            with open(scenario_file, 'r') as f:
                scenario = json.load(f)
            logger.info(f"Loaded scenario from {scenario_file}")
            return scenario
        except Exception as e:
            logger.error(f"Failed to load scenario: {e}")
            return {}
    
    def run_scenario(self, scenario: Dict) -> Dict:
        """
        Run a single scenario
        """
        logger.info("Running scenario...")
        self.start_time = time.time()
        
        # Log conversation start
        self.trace_logger.log_conversation_start(scenario)
        
        # Run the ReAct loop
        try:
            result = self.react_agent.run_react_loop(scenario)
            
            # Check time constraint
            execution_time = time.time() - self.start_time
            if execution_time > self.max_execution_time:
                logger.warning(f"Scenario exceeded time limit: {execution_time:.2f}s > {self.max_execution_time}s")
                result["time_exceeded"] = True
            else:
                logger.info(f"Scenario completed in {execution_time:.2f}s")
            
            # Log conversation end
            self.trace_logger.log_conversation_end(result)
            
            return result
            
        except Exception as e:
            logger.error(f"Scenario execution failed: {e}")
            error_result = {
                "error": str(e),
                "execution_time": time.time() - self.start_time,
                "success": False
            }
            self.trace_logger.log_conversation_end(error_result)
            return error_result
    
    def run_batch_scenarios(self, scenarios_file: str) -> Dict:
        """
        Run multiple scenarios from a file
        """
        try:
            with open(scenarios_file, 'r') as f:
                scenarios = json.load(f)
            
            logger.info(f"Running {len(scenarios)} scenarios...")
            
            results = []
            total_start_time = time.time()
            
            for i, scenario in enumerate(scenarios):
                logger.info(f"Running scenario {i+1}/{len(scenarios)}")
                
                result = self.run_scenario(scenario)
                result["scenario_index"] = i
                results.append(result)
                
                # Check if we're approaching time limit
                elapsed = time.time() - total_start_time
                remaining = self.max_execution_time - elapsed
                
                if remaining < 60:  # Less than 1 minute remaining
                    logger.warning(f"Approaching time limit, stopping early. Remaining: {remaining:.2f}s")
                    break
            
            total_time = time.time() - total_start_time
            logger.info(f"Batch completed in {total_time:.2f}s")
            
            return {
                "total_scenarios": len(results),
                "total_execution_time": total_time,
                "results": results,
                "trace_summary": self.trace_logger.get_trace_summary()
            }
            
        except Exception as e:
            logger.error(f"Batch execution failed: {e}")
            return {"error": str(e)}
    
    def generate_submission_file(self, results: Dict, output_file: str = "result.csv"):
        """
        Generate submission file from results with robust formatting
        """
        try:
            import pandas as pd
            
            # Convert results to submission format
            submission_data = []
            
            for result in results.get("results", []):
                scenario = result.get("scenario", {})
                solution = result.get("final_state", {}).get("solution", {})
                
                # Extract and format Track A answer
                track_a_answer = solution.get("track_a_answer", "")
                if not track_a_answer:
                    # Fallback: derive from conversation history
                    track_a_answer = self._extract_track_a_from_history(result)
                
                # Ensure Track A format (C1|C2|C3)
                track_a_answer = self._format_track_a(track_a_answer)
                
                # Extract and format Track B answer
                track_b_answer = solution.get("track_b_answer", "")
                if not track_b_answer:
                    # Fallback: derive from conversation history
                    track_b_answer = self._extract_track_b_from_history(result)
                
                # Ensure Track B is string (can be empty)
                track_b_answer = str(track_b_answer) if track_b_answer else ""
                
                submission_data.append({
                    "ID": scenario.get("id", f"scenario_{result.get('scenario_index', 'unknown')}"),
                    "Track A": track_a_answer,
                    "Track B": track_b_answer
                })
            
            # Create DataFrame with proper column order
            df = pd.DataFrame(submission_data, columns=["ID", "Track A", "Track B"])
            
            # Remove any NaN values
            df = df.fillna("")
            
            # Validate format before saving
            from utils.format_validator import get_format_validator
            validator = get_format_validator()
            
            # Save to temporary file first
            temp_file = output_file.replace('.csv', '_temp.csv')
            df.to_csv(temp_file, index=False)
            
            # Validate the format
            is_valid, validation_result = validator.validate_csv_format(temp_file)
            
            if not is_valid:
                logger.error(f"CSV validation failed: {validation_result}")
                return False
            
            # Move to final location
            import os
            os.rename(temp_file, output_file)
            
            logger.info(f"Submission file saved and validated: {output_file}")
            logger.info(f"Total entries: {len(df)}")
            logger.info(f"Track A non-empty: {(df['Track A'] != '').sum()}")
            logger.info(f"Track B non-empty: {(df['Track B'] != '').sum()}")
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to generate submission file: {e}")
            return False
    
    def _extract_track_a_from_history(self, result: Dict) -> str:
        """Extract Track A answer from conversation history"""
        history = result.get("conversation_history", [])
        
        # Look for cell selections in conversation
        for entry in reversed(history):
            thought = entry.get("thought", "")
            action = entry.get("action", "")
            
            # Look for cell patterns like "C1", "C2", etc.
            import re
            cells = re.findall(r'C\d+', thought)
            if cells:
                return "|".join(sorted(set(cells)))
        
        return ""
    
    def _extract_track_b_from_history(self, result: Dict) -> str:
        """Extract Track B answer from conversation history"""
        history = result.get("conversation_history", [])
        
        # Look for diagnosis or solution in conversation
        for entry in reversed(history):
            thought = entry.get("thought", "")
            observation = entry.get("observation", {})
            
            # Look for diagnosis patterns
            if "diagnosis" in thought.lower() or "root cause" in thought.lower():
                return thought[:200]  # First 200 chars
        
        return ""
    
    def _format_track_a(self, track_a: str) -> str:
        """Format Track A answer to ensure proper format"""
        if not track_a:
            return ""
        
        # Extract cell references
        import re
        cells = re.findall(r'C\d+', str(track_a))
        
        # Remove duplicates and sort
        unique_cells = sorted(set(cells))
        
        # Join with pipe
        return "|".join(unique_cells) if unique_cells else ""
    
    def cleanup(self):
        """
        Clean up resources
        """
        try:
            if self.llm_engine:
                self.llm_engine.cleanup()
            logger.info("Agent cleanup completed")
        except Exception as e:
            logger.error(f"Cleanup failed: {e}")


def main():
    """
    Main entry point for Phase 3 evaluation
    """
    parser = argparse.ArgumentParser(description="Telco Troubleshooting Agent")
    parser.add_argument("--server", default="http://localhost:8000", help="Server URL")
    parser.add_argument("--scenario", help="Single scenario file")
    parser.add_argument("--scenarios", help="Batch scenarios file")
    parser.add_argument("--output", default="result.csv", help="Output submission file")
    parser.add_argument("--validate", action="store_true", help="Validate trace file only")
    
    args = parser.parse_args()
    
    # Initialize agent
    agent = TelcoAgent(server_url=args.server)
    
    if not agent.initialize():
        sys.exit(1)
    
    try:
        if args.validate:
            # Validate trace file
            is_valid = agent.trace_logger.validate_trace_file()
            print(f"Trace file validation: {'PASS' if is_valid else 'FAIL'}")
            return
        
        if args.scenarios:
            # Run batch scenarios
            results = agent.run_batch_scenarios(args.scenarios)
            
            # Generate submission file
            agent.generate_submission_file(results, args.output)
            
            # Print summary
            print(f"\nEXECUTION SUMMARY")
            print(f"Total scenarios: {results.get('total_scenarios', 0)}")
            print(f"Total time: {results.get('total_execution_time', 0):.2f}s")
            print(f"Trace entries: {results.get('trace_summary', {}).get('total_entries', 0)}")
            print(f"Success rate: {results.get('trace_summary', {}).get('success_rate', 0):.1f}%")
            
        elif args.scenario:
            # Run single scenario
            scenario = agent.load_scenario(args.scenario)
            if scenario:
                result = agent.run_scenario(scenario)
                print(f"Scenario completed: {result.get('final_state', {}).get('completed', False)}")
                print(f"Execution time: {result.get('execution_time', 0):.2f}s")
        
        else:
            # Interactive mode for testing
            logger.info("Running in interactive mode...")
            test_scenario = {
                "id": "test_001",
                "description": "Network connectivity issue",
                "symptoms": ["ping fails", "high latency"],
                "environment": "test"
            }
            
            result = agent.run_scenario(test_scenario)
            print(f"Test completed: {result.get('final_state', {}).get('completed', False)}")
    
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    except Exception as e:
        logger.error(f"Execution failed: {e}")
    finally:
        agent.cleanup()


if __name__ == "__main__":
    main()