"""
BDiasLangChainTools: LangChain Integration Module for Bart.dIAs

This module provides LangChain tool wrappers for the Bart.dIAs core functionality,
enabling seamless integration with LangChain agents and workflows.

Features:
- AnalyzeCodeTool: Analyzes Python code for parallelization opportunities
- GenerateParallelCodeTool: Generates parallelized code based on identified patterns

Tools:
- AnalyzeCodeTool: Multi-mode code analysis (blocks, critical path, patterns, integrated)
- GenerateParallelCodeTool: Pattern-based parallel code generation

Dependencies:
- langchain-core: For BaseTool implementation
- bdias_parser: Code parsing
- bdias_critical_path: Critical path analysis
- bdias_pattern_analyzer: Pattern recognition
- bdias_pattern_codegen: Code generation
"""

import os
from typing import Optional, Type, Literal, Tuple
from pydantic import BaseModel, Field

try:
    from langchain_core.tools import BaseTool
except ImportError:
    raise ImportError(
        "langchain-core is required for LangChain tools. "
        "Install it with: pip install langchain-core"
    )

from .bdias_parser import BDiasParser
from .bdias_critical_path import BDiasCriticalPathAnalyzer
from .bdias_pattern_analyzer import BDiasPatternAnalyzer
from .bdias_pattern_codegen import generate_parallel_code
from .bdias_pattern_presenter import present_transformation


class AnalyzeCodeInput(BaseModel):
    """Input schema for AnalyzeCodeTool."""
    
    code_input: str = Field(
        description="Python code as a string or path to a Python file to analyze"
    )
    analysis_type: Literal["blocks", "critical_path", "patterns", "integrated"] = Field(
        default="integrated",
        description=(
            "Type of analysis to perform: "
            "'blocks' for block-based opportunities, "
            "'critical_path' for DAG-based metrics, "
            "'patterns' for parallel pattern recognition, "
            "'integrated' for combined critical path and pattern analysis"
        )
    )


class GenerateParallelCodeInput(BaseModel):
    """Input schema for GenerateParallelCodeTool."""
    
    code_input: str = Field(
        description="Python code as a string or path to a Python file"
    )
    pattern: Optional[str] = Field(
        default=None,
        description=(
            "Parallel pattern to use (e.g., 'map_reduce', 'pipeline', 'master_slave', 'pool_workers'). "
            "If not specified, will auto-detect from code"
        )
    )
    line_number: Optional[int] = Field(
        default=None,
        description="Specific line number to target for transformation. If not specified, uses first detected pattern"
    )


class AnalyzeCodeTool(BaseTool):
    """
    LangChain tool for analyzing Python code to identify parallelization opportunities.
    
    This tool provides multiple analysis modes:
    - Block-based: Identifies parallelizable loops, functions, and list comprehensions
    - Critical path: Performs DAG-based analysis to calculate theoretical parallelism metrics
    - Patterns: Recognizes common parallel programming patterns
    - Integrated: Combines critical path analysis with pattern recognition
    """
    
    name: str = "analyze_python_code"
    description: str = (
        "Analyzes Python code to identify parallelization opportunities. "
        "Supports multiple analysis modes: blocks (basic opportunities), "
        "critical_path (DAG metrics), patterns (parallel patterns), or "
        "integrated (comprehensive analysis). Input can be code string or file path."
    )
    args_schema: Type[BaseModel] = AnalyzeCodeInput
    
    def _run(self, code_input: str, analysis_type: str = "integrated") -> str:
        """Execute the analysis."""
        try:
            # Detect if input is file path or code string
            code = self._load_code(code_input)
            
            # Initialize parser and parse code to populate parser.tree
            parser = BDiasParser()
            structured_code = parser.parse(code)
            
            if structured_code is None:
                return "Error: Failed to parse code. Please check for syntax errors."
            
            # Route to appropriate analysis method
            if analysis_type == "blocks":
                return self._analyze_blocks(parser, code)
            elif analysis_type == "critical_path":
                return self._analyze_critical_path(parser, code)
            elif analysis_type == "patterns":
                return self._analyze_patterns(parser, code)
            elif analysis_type == "integrated":
                return self._analyze_integrated(parser, code)
            else:
                return f"Error: Invalid analysis_type '{analysis_type}'. Must be one of: blocks, critical_path, patterns, integrated"
                
        except Exception as e:
            return f"Error analyzing code: {str(e)}"
    
    def _load_code(self, code_input: str) -> str:
        """Load code from file or return as-is if it's a code string."""
        # Check if it's a file path
        if os.path.isfile(code_input):
            with open(code_input, 'r') as f:
                return f.read()
        return code_input
    
    def _analyze_blocks(self, parser: BDiasParser, code: str) -> str:
        """Perform block-based analysis."""
        structured_code = parser.parse(code)
        
        if structured_code is None:
            return "Error: Failed to parse code. Please check for syntax errors."
        
        output = ["## Block-Based Parallelization Analysis\n"]
        
        # Check for opportunities
        has_opportunities = any(structured_code.get(key) for key in structured_code)
        
        if not has_opportunities:
            return "No parallelization opportunities identified in the code."
        
        # Loops
        if structured_code.get('loops'):
            output.append("### Parallelizable Loops")
            for loop in structured_code['loops']:
                loop_type = loop.get('type', 'loop')
                lineno = loop.get('lineno', '?')
                output.append(f"- **Line {lineno}**: {loop_type}")
                if loop_type == 'loop':
                    output.append(f"  - Loop variable: `{loop.get('loop_var', '?')}`")
                    output.append(f"  - Iterates over: `{loop.get('iterable_name', '?')}`")
                output.append("")
        
        # Functions
        if structured_code.get('functions'):
            output.append("### Parallelizable Functions")
            for func in structured_code['functions']:
                func_type = func.get('type', 'function')
                lineno = func.get('lineno', '?')
                name = func.get('name', '?')
                output.append(f"- **Line {lineno}**: {func_type} `{name}`")
                output.append("")
        
        # List comprehensions
        if structured_code.get('list_comprehensions'):
            output.append("### Parallelizable List Comprehensions")
            for lc in structured_code['list_comprehensions']:
                lineno = lc.get('lineno', '?')
                output.append(f"- **Line {lineno}**: List comprehension")
                output.append("")
        
        # Combo patterns
        if structured_code.get('combos'):
            output.append("### Complex Patterns")
            for combo in structured_code['combos']:
                combo_type = combo.get('type', 'combo')
                lineno = combo.get('lineno', '?')
                output.append(f"- **Line {lineno}**: {combo_type}")
                output.append("")
        
        return "\n".join(output)
    
    def _analyze_critical_path(self, parser: BDiasParser, code: str) -> str:
        """Perform critical path analysis."""
        analyzer = BDiasCriticalPathAnalyzer()
        results = analyzer.analyze(parser, code)
        
        output = ["## Critical Path Analysis (DAG-Based)\n"]
        output.append("### Theoretical Metrics")
        output.append(f"- **Total Work (T₁)**: {results['total_work']:.2f}")
        output.append(f"- **Critical Path Length (T∞)**: {results['critical_path_length']:.2f}")
        output.append(f"- **Theoretical Parallelism (T₁/T∞)**: {results['parallelism']:.2f}x")
        output.append(f"- **Sequential Fraction**: {results['sequential_fraction']:.2%}")
        output.append(f"- **Amdahl's Max Speedup**: {results['amdahl_max_speedup']:.2f}x\n")
        
        # Bottlenecks
        bottlenecks = results.get('bottlenecks', [])
        if bottlenecks:
            output.append("### Critical Path Bottlenecks")
            for i, bottleneck in enumerate(bottlenecks[:5], 1):  # Top 5
                output.append(f"{i}. **{bottleneck['type'].replace('_', ' ').title()}** (Line {bottleneck['lineno']})")
                output.append(f"   - Work: {bottleneck['work']:.2f}")
                output.append(f"   - Span: {bottleneck['span']:.2f}")
                source = bottleneck['source'][:60]
                output.append(f"   - Code: `{source}{'...' if len(bottleneck['source']) > 60 else ''}`")
                output.append("")
        else:
            output.append("No significant bottlenecks identified.")
        
        # Recommendations
        if bottlenecks:
            output.append("### Recommendations")
            output.append("- Focus on parallelizing the identified bottlenecks")
            output.append("- Consider breaking sequential sections into independent tasks")
            output.append("- Apply algorithmic transformations to reduce dependencies")
        
        return "\n".join(output)
    
    def _analyze_patterns(self, parser: BDiasParser, code: str) -> str:
        """Perform pattern recognition analysis."""
        analyzer = BDiasPatternAnalyzer(parser)
        analysis = analyzer.analyze(code)
        
        output = ["## Parallel Pattern Recognition\n"]
        
        identified = analysis.get("identified_patterns", {})
        if not identified:
            return "No specific parallel patterns identified in the code."
        
        output.append("### Identified Patterns\n")
        
        for pattern_name, instances in identified.items():
            output.append(f"#### {pattern_name.upper()} Pattern")
            output.append(f"**Instances found**: {len(instances)}\n")
            
            for i, instance in enumerate(instances[:3], 1):  # Show top 3
                output.append(f"{i}. Line {instance['lineno']}")
                output.append(f"   - Type: {instance['type']}")
                output.append(f"   - Confidence: {instance['confidence']:.2f}")
            
            if len(instances) > 3:
                output.append(f"   _(... and {len(instances) - 3} more)_")
            
            # Partitioning recommendations
            partitioning = analysis["recommended_partitioning"].get(pattern_name, {})
            if partitioning:
                strategies = partitioning.get('strategies', [])
                output.append(f"\n**Recommended Partitioning**: {', '.join(strategies)}")
                output.append(f"**Rationale**: {partitioning.get('rationale', '')}")
            
            # Performance characteristics
            perf = analysis["performance_characteristics"].get(pattern_name, {})
            if perf:
                output.append(f"\n**Performance Characteristics**:")
                output.append(f"- Work: {perf.get('work', 'Unknown')}")
                output.append(f"- Span: {perf.get('span', 'Unknown')}")
                output.append(f"- Parallelism: {perf.get('parallelism', 'Unknown')}")
                output.append(f"- Communication Overhead: {perf.get('communication_overhead', 'Unknown')}")
            
            output.append("")
        
        return "\n".join(output)
    
    def _analyze_integrated(self, parser: BDiasParser, code: str) -> str:
        """Perform integrated critical path and pattern analysis."""
        # Critical path analysis
        cp_analyzer = BDiasCriticalPathAnalyzer()
        cp_results = cp_analyzer.analyze(parser, code)
        
        # Pattern analysis
        pattern_analyzer = BDiasPatternAnalyzer(parser)
        pattern_analysis = pattern_analyzer.analyze(code)
        
        output = ["## Integrated Analysis: Critical Path + Pattern Recognition\n"]
        
        # Metrics
        output.append("### Parallelism Metrics")
        output.append(f"- **Total Work (T₁)**: {cp_results['total_work']:.2f}")
        output.append(f"- **Critical Path (T∞)**: {cp_results['critical_path_length']:.2f}")
        output.append(f"- **Theoretical Parallelism**: {cp_results['parallelism']:.2f}x")
        output.append(f"- **Amdahl's Max Speedup**: {cp_results['amdahl_max_speedup']:.2f}x\n")
        
        # Bottlenecks with pattern suggestions
        output.append("### Critical Path Bottlenecks with Suggested Patterns\n")
        
        bottlenecks = cp_results.get('bottlenecks', [])
        identified = pattern_analysis.get('identified_patterns', {})
        recommended = pattern_analysis.get('recommended_partitioning', {})
        
        if not bottlenecks:
            output.append("No significant bottlenecks identified.")
        else:
            for i, bottleneck in enumerate(bottlenecks[:5], 1):
                output.append(f"#### {i}. {bottleneck['type'].replace('_', ' ').title()} (Line {bottleneck['lineno']})")
                output.append(f"**Computational Cost**: Work={bottleneck['work']:.2f}, Span={bottleneck['span']:.2f}")
                
                source = bottleneck['source'].strip().splitlines()[0][:60]
                output.append(f"**Code**: `{source}...`\n")
                
                # Find matching patterns
                matches = []
                for pattern_name, instances in identified.items():
                    for inst in instances:
                        if inst.get('lineno') == bottleneck['lineno']:
                            matches.append((pattern_name, inst))
                
                if matches:
                    output.append("**Suggested Parallel Patterns**:")
                    for j, (pattern_name, inst) in enumerate(matches[:2], 1):
                        part = recommended.get(pattern_name, {})
                        output.append(f"{j}. **{pattern_name.upper()}** (Confidence: {inst['confidence']:.2f})")
                        output.append(f"   - Partitioning: {', '.join(part.get('strategies', []))}")
                        output.append(f"   - Rationale: {part.get('rationale', '')}")
                else:
                    output.append("**Suggested Patterns**: Consider general parallelization strategies")
                
                output.append("")
        
        return "\n".join(output)


class GenerateParallelCodeTool(BaseTool):
    """
    LangChain tool for generating parallelized Python code based on identified patterns.
    
    This tool transforms sequential code into parallel implementations using
    pattern-specific strategies and multiprocessing templates.
    """
    
    name: str = "generate_parallel_code"
    description: str = (
        "Generates parallelized Python code from sequential code. "
        "Can auto-detect patterns or use a specified pattern. "
        "Supports map_reduce, pipeline, master_slave, and pool_workers patterns. "
        "Input can be code string or file path."
    )
    args_schema: Type[BaseModel] = GenerateParallelCodeInput
    
    def _run(
        self, 
        code_input: str, 
        pattern: Optional[str] = None, 
        line_number: Optional[int] = None
    ) -> str:
        """Generate parallel code."""
        try:
            # Load code
            code = self._load_code(code_input)
            
            # Parse code to populate parser.tree
            parser = BDiasParser()
            structured_code = parser.parse(code)
            
            if structured_code is None:
                return "Error: Failed to parse code. Please check for syntax errors."
            
            # If pattern not specified, detect it
            if pattern is None:
                pattern, line_number = self._auto_detect_pattern(parser, code, line_number)
                if pattern is None:
                    return "Error: No parallelizable patterns detected in the code."
            
            # Get pattern analysis for partitioning strategy
            pattern_analyzer = BDiasPatternAnalyzer(parser)
            pattern_analysis = pattern_analyzer.analyze(code)
            
            # Get partitioning strategy for the pattern
            recommended = pattern_analysis.get('recommended_partitioning', {})
            pattern_info = recommended.get(pattern, {})
            
            if not pattern_info:
                # Default strategies if pattern not recognized
                strategies = ['default']
            else:
                strategies = pattern_info.get('strategies', ['default'])
            
            # Create bottleneck dict for code generation
            bottleneck = {
                'source': code,
                'lineno': line_number or 1,
                'type': pattern
            }
            
            # Generate parallel code
            original_code, transformed_code, context = generate_parallel_code(
                bottleneck,
                pattern,
                strategies
            )
            
            # Format output
            pattern_display_info = {
                'pattern': pattern,
                'partitioning_strategy': strategies,
                'context': context
            }
            
            output = present_transformation(
                original_code,
                transformed_code,
                pattern_display_info
            )
            
            return output
            
        except Exception as e:
            return f"Error generating parallel code: {str(e)}"
    
    def _load_code(self, code_input: str) -> str:
        """Load code from file or return as-is."""
        if os.path.isfile(code_input):
            with open(code_input, 'r') as f:
                return f.read()
        return code_input
    
    def _auto_detect_pattern(
        self, 
        parser: BDiasParser, 
        code: str, 
        preferred_line: Optional[int]
    ) -> Tuple[Optional[str], Optional[int]]:
        """Auto-detect the most appropriate pattern."""
        pattern_analyzer = BDiasPatternAnalyzer(parser)
        analysis = pattern_analyzer.analyze(code)
        
        identified = analysis.get('identified_patterns', {})
        if not identified:
            return None, None
        
        # If preferred line specified, find pattern at that line
        if preferred_line:
            for pattern_name, instances in identified.items():
                for inst in instances:
                    if inst.get('lineno') == preferred_line:
                        return pattern_name, preferred_line
        
        # Otherwise, return highest confidence pattern
        best_pattern = None
        best_confidence = 0
        best_line = None
        
        for pattern_name, instances in identified.items():
            for inst in instances:
                conf = inst.get('confidence', 0)
                if conf > best_confidence:
                    best_confidence = conf
                    best_pattern = pattern_name
                    best_line = inst.get('lineno')
        
        return best_pattern, best_line


# Convenience function to get both tools
def get_bartdias_tools():
    """
    Get both Bart.dIAs LangChain tools.
    
    Returns:
        list: [AnalyzeCodeTool(), GenerateParallelCodeTool()]
    """
    return [AnalyzeCodeTool(), GenerateParallelCodeTool()]

