import json
from typing import List, Dict, Any, Set
from mcp.server.fastmcp.tools import Tool
from ..dependency_manager import get_dependency_manager


def get_job_execution_sequence() -> str:
    """Analyze complete job execution sequence and dependencies for scheduled tasks."""
    try:
        dm = get_dependency_manager()
        scheduled_tasks_data = dm.get_components_by_type("SCHEDULED_TASK")
        scheduled_tasks = [task['name'] for task in scheduled_tasks_data]
        
        execution_tree = {}
        for task in scheduled_tasks:
            dependencies = dm.get_dependencies(task)
            
            # Find JCL script (look for "Exec JCL" type)
            jcl_script = next((dep['name'] for dep in dependencies 
                             if dep.get('type') == 'Exec JCL'), None)
            
            # Find next scheduled tasks
            next_jobs = [dep['name'] for dep in dependencies 
                        if dep.get('type') == 'Exec Next Scheduled Task']
            
            execution_tree[task] = {
                'jcl_script': jcl_script,
                'next_jobs': next_jobs
            }
        
        # Find critical path
        critical_path = _find_critical_path(execution_tree)
        
        result = {
            'total_scheduled_tasks': len(scheduled_tasks),
            'execution_tree': execution_tree,
            'critical_path': critical_path,
            'critical_path_length': len(critical_path)
        }
        
        return json.dumps(result, indent=2)
        
    except Exception as e:
        return f"Error analyzing job execution sequence: {str(e)}"


def _find_critical_path(execution_tree: Dict[str, Dict]) -> List[str]:
    """Find the longest dependency chain."""
    def get_chain_length(job: str, visited: Set[str] = None) -> int:
        if visited is None:
            visited = set()
        if job in visited:
            return 0
        visited.add(job)
        
        next_jobs = execution_tree.get(job, {}).get('next_jobs', [])
        if not next_jobs:
            return 1
        
        return 1 + max(get_chain_length(next_job, visited.copy()) for next_job in next_jobs)
    
    if not execution_tree:
        return []
    
    # Find job with longest chain
    longest_job = max(execution_tree.keys(), key=lambda job: get_chain_length(job))
    
    # Build the actual path
    path = []
    current = longest_job
    visited = set()
    
    while current and current not in visited:
        path.append(current)
        visited.add(current)
        next_jobs = execution_tree.get(current, {}).get('next_jobs', [])
        current = next_jobs[0] if next_jobs else None
    
    return path


def get_job_impact_boundary(job_name: str) -> str:
    """Analyze the impact boundary for a specific job - shows forward and backward dependencies."""
    try:
        dm = get_dependency_manager()
        
        # Get job info
        job_info = dm.get_component_info(job_name)
        if not job_info:
            return json.dumps({'error': f'Job {job_name} not found'}, indent=2)
        
        # Forward dependencies (what depends on this job)
        forward_deps = dm.get_recursive_dependents(job_name)
        
        # Backward dependencies (what this job depends on)
        backward_deps = dm.get_recursive_dependencies(job_name)
        
        # Categorize by type
        forward_by_type = _categorize_by_type(forward_deps)
        backward_by_type = _categorize_by_type(backward_deps)
        
        result = {
            'target_job': job_name,
            'job_type': job_info.get('type'),
            'forward_impact': {
                'description': 'Components that depend on this job',
                'total_count': len(forward_deps),
                'by_type': forward_by_type
            },
            'backward_dependencies': {
                'description': 'Components this job depends on',
                'total_count': len(backward_deps),
                'by_type': backward_by_type
            }
        }
        
        return json.dumps(result, indent=2)
        
    except Exception as e:
        return f"Error analyzing job impact boundary for {job_name}: {str(e)}"


def _categorize_by_type(dependencies: List[Dict]) -> Dict[str, int]:
    """Categorize dependencies by type."""
    categories = {}
    for dep in dependencies:
        dep_type = dep.get('type', 'Unknown')
        categories[dep_type] = categories.get(dep_type, 0) + 1
    
    return dict(sorted(categories.items(), key=lambda x: x[1], reverse=True))


# Tool definitions for MCP server
get_job_execution_sequence_tool = Tool.from_function(
    fn=get_job_execution_sequence,
    name="get_job_execution_sequence",
    description="Analyze complete job execution sequence and dependencies for scheduled tasks. Includes critical path analysis."
)

get_job_impact_boundary_tool = Tool.from_function(
    fn=get_job_impact_boundary,
    name="get_job_impact_boundary",
    description="Analyze the impact boundary for a specific job - shows what depends on it (forward) and what it depends on (backward)."
)
