"""
BDias Analysis Module

This module contains high-level functions for analyzing code using Bart.dIAs components.
It provides programmatic access to block-based analysis, critical path analysis, 
pattern recognition, and code generation.

Supports two output formats:
- "structured": Returns dicts with raw data (ideal for LLM consumption)
- "markdown":   Returns formatted Markdown strings (for CLI/human reading)
"""

import os
import sys
import logging
from typing import Optional, Literal, Tuple, List, Dict, Any

logger = logging.getLogger(__name__)

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 as bdias_generate_parallel_code
from .bdias_pattern_presenter import present_transformation
from .bdias_code_gen import BDiasCodeGen
from .constants import EXPLANATIONS, PARTITIONING_SUGGESTIONS


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def load_code(code_input: str) -> str:
    """
    Load code from file or return as-is if it's a code string.
    
    Args:
        code_input: Python code as a string or path to a Python file
        
    Returns:
        str: The code content
        
    Raises:
        FileNotFoundError: If the file path doesn't exist
        IOError: If there's an error reading the file
    """
    if os.path.isfile(code_input):
        logger.info(f"Loading code from file: {code_input}")
        with open(code_input, 'r') as f:
            return f.read()
    logger.info("Using provided code string")
    return code_input


def _filter_by_line_range(
    items: List[Dict],
    start_line: Optional[int],
    end_line: Optional[int],
    line_key: str = "lineno"
) -> List[Dict]:
    """Filter a list of dicts keeping only those whose line is within [start_line, end_line]."""
    if start_line is None and end_line is None:
        return items
    filtered = []
    for item in items:
        lineno = item.get(line_key)
        if lineno is None:
            continue
        if start_line is not None and lineno < start_line:
            continue
        if end_line is not None and lineno > end_line:
            continue
        filtered.append(item)
    return filtered


def _serialize_dag(dag, critical_path: List[str]) -> Dict[str, Any]:
    """Serialize a NetworkX DAG into a JSON-friendly dict."""
    nodes = []
    for node_id, data in dag.nodes(data=True):
        nodes.append({
            "id": str(node_id),
            "work": data.get("work", 0),
            "span": data.get("span", 0),
            "type": data.get("type", ""),
            "lineno": data.get("lineno", None),
            "source": data.get("source", ""),
        })
    edges = []
    for u, v in dag.edges():
        edges.append({"from": str(u), "to": str(v)})
    return {
        "nodes": nodes,
        "edges": edges,
        "critical_path_nodes": [str(n) for n in critical_path],
    }


def _try_generate_code(
    source_code: str,
    pattern: str,
    strategies: List[str]
) -> Dict[str, Any]:
    """
    Safely attempt to generate parallel code using Jinja2 pattern templates.
    Returns dict with original_code, transformed_code, template_context, or error.
    """
    try:
        bottleneck = {'source': source_code, 'lineno': 1, 'type': pattern}
        original, transformed, context = bdias_generate_parallel_code(
            bottleneck, pattern, strategies
        )
        # Make context JSON-serializable (remove non-serializable items)
        safe_context = {}
        for k, v in context.items():
            if isinstance(v, (str, int, float, bool, list, dict, type(None))):
                safe_context[k] = v
        return {
            "original_code": original,
            "transformed_code": transformed,
            "template_context": safe_context,
        }
    except Exception as e:
        return {
            "original_code": source_code,
            "transformed_code": None,
            "error": str(e),
        }


# ---------------------------------------------------------------------------
# Structured output functions (return dicts)
# ---------------------------------------------------------------------------

def _analyze_blocks_structured(
    parser: BDiasParser,
    code: str,
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> Dict[str, Any]:
    """Return block-based analysis as structured data."""
    structured_code = parser.parse(code)

    if structured_code is None:
        return {"opportunities": [], "total_count": 0, "error": "Failed to parse code."}

    opportunities: List[Dict] = []

    for loop in structured_code.get('loops', []):
        opportunities.append({
            "category": "loop",
            "type": loop.get('type', 'loop'),
            "lineno": loop.get('lineno'),
            "loop_var": loop.get('loop_var'),
            "iterable_name": loop.get('iterable_name'),
        })

    for func in structured_code.get('functions', []):
        opportunities.append({
            "category": "function",
            "type": func.get('type', 'function'),
            "lineno": func.get('lineno'),
            "name": func.get('name'),
        })

    for lc in structured_code.get('list_comprehensions', []):
        opportunities.append({
            "category": "list_comprehension",
            "type": "list_comprehension",
            "lineno": lc.get('lineno'),
        })

    for combo in structured_code.get('combos', []):
        opportunities.append({
            "category": "combo",
            "type": combo.get('type', 'combo'),
            "lineno": combo.get('lineno'),
        })

    opportunities = _filter_by_line_range(opportunities, start_line, end_line)

    return {
        "opportunities": opportunities,
        "total_count": len(opportunities),
    }


def _analyze_blocks_with_suggestions_structured(
    parser: BDiasParser,
    code: str,
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> Dict[str, Any]:
    """Return block-based analysis with code suggestions from BDiasCodeGen templates."""
    structured_code = parser.parse(code)

    if structured_code is None:
        return {"suggestions": [], "total_count": 0, "error": "Failed to parse code."}

    code_gen = BDiasCodeGen(EXPLANATIONS, PARTITIONING_SUGGESTIONS)
    raw_suggestions = code_gen.generate_suggestions(structured_code)

    suggestions: List[Dict] = []
    for s in raw_suggestions:
        explanation_key = s.get("explanation_index", "")
        explanation = ""
        try:
            explanation = EXPLANATIONS.get(explanation_key, "").format(**s)
        except (KeyError, IndexError):
            explanation = EXPLANATIONS.get(explanation_key, "")

        partitioning = PARTITIONING_SUGGESTIONS.get(explanation_key, 
            "Consider data or task partitioning based on the specific pattern.")

        entry = {
            "lineno": s.get("lineno"),
            "opportunity_type": s.get("opportunity_type", "unknown"),
            "explanation": explanation,
            "partitioning_suggestion": partitioning,
            "code_suggestion": s.get("code_suggestion", ""),
        }
        # Copy extra fields if present
        for key in ("loop_var", "iterable_name", "func_name", "parent_func_name"):
            if key in s:
                entry[key] = s[key]

        suggestions.append(entry)

    suggestions = _filter_by_line_range(suggestions, start_line, end_line)

    return {
        "suggestions": suggestions,
        "total_count": len(suggestions),
    }


def _analyze_critical_path_structured(
    parser: BDiasParser,
    code: str,
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> Dict[str, Any]:
    """Return critical path analysis as structured data."""
    analyzer = BDiasCriticalPathAnalyzer()
    results = analyzer.analyze(parser, code)

    # Pattern analysis for auto-detecting patterns per bottleneck
    pattern_analyzer = BDiasPatternAnalyzer(parser)
    pattern_analysis = pattern_analyzer.analyze(code)
    identified = pattern_analysis.get('identified_patterns', {})
    recommended = pattern_analysis.get('recommended_partitioning', {})

    raw_bottlenecks = results.get('bottlenecks', [])
    raw_bottlenecks = _filter_by_line_range(raw_bottlenecks, start_line, end_line)

    enriched = []
    for b in raw_bottlenecks:
        entry = {
            "type": b['type'],
            "lineno": b['lineno'],
            "work": b['work'],
            "span": b['span'],
            "source_code": b['source'],
        }
        # Try to find a matching pattern and generate template code
        matched_pattern = None
        for pattern_name, instances in identified.items():
            for inst in instances:
                if inst.get('lineno') == b['lineno']:
                    matched_pattern = pattern_name
                    break
            if matched_pattern:
                break

        if matched_pattern:
            strategies = recommended.get(matched_pattern, {}).get('strategies', ['default'])
            entry["detected_pattern"] = matched_pattern
            entry["code_generation"] = _try_generate_code(b['source'], matched_pattern, strategies)
        else:
            entry["detected_pattern"] = None
            entry["code_generation"] = None

        enriched.append(entry)

    return {
        "metrics": {
            "total_work": results['total_work'],
            "critical_path_length": results['critical_path_length'],
            "parallelism": results['parallelism'],
            "sequential_fraction": results['sequential_fraction'],
            "amdahl_max_speedup": results['amdahl_max_speedup'],
        },
        "bottlenecks": enriched,
        "dag_data": _serialize_dag(
            results['dag'],
            results.get('critical_path', [])
        ),
    }


def _analyze_patterns_structured(
    parser: BDiasParser,
    code: str,
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> Dict[str, Any]:
    """Return pattern recognition analysis as structured data."""
    analyzer = BDiasPatternAnalyzer(parser)
    analysis = analyzer.analyze(code)

    identified = analysis.get("identified_patterns", {})
    recommended = analysis.get("recommended_partitioning", {})
    performance = analysis.get("performance_characteristics", {})

    # Build structured output, filtering by line range
    patterns_out: Dict[str, Any] = {}
    for pattern_name, instances in identified.items():
        filtered = _filter_by_line_range(instances, start_line, end_line)
        if not filtered:
            continue

        perf = performance.get(pattern_name, {})
        part = recommended.get(pattern_name, {})
        strategies = part.get('strategies', ['default'])

        # Generate template code for this pattern using the full source
        code_generation = _try_generate_code(code, pattern_name, strategies)

        instances_out = []
        for inst in filtered:
            instances_out.append({
                "lineno": inst.get("lineno"),
                "type": inst.get("type"),
                "confidence": inst.get("confidence"),
            })

        patterns_out[pattern_name] = {
            "instances": instances_out,
            "recommended_partitioning": {
                "strategies": strategies,
                "rationale": part.get("rationale", ""),
            },
            "performance_characteristics": {
                "work": perf.get("work", "Unknown"),
                "span": perf.get("span", "Unknown"),
                "parallelism": perf.get("parallelism", "Unknown"),
                "communication_overhead": perf.get("communication_overhead", "Unknown"),
                "synchronization_points": perf.get("synchronization_points", "Unknown"),
            },
            "code_generation": code_generation,
        }

    return {
        "identified_patterns": patterns_out,
        "total_patterns": len(patterns_out),
    }


def _analyze_integrated_structured(
    parser: BDiasParser,
    code: str,
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> Dict[str, Any]:
    """Return integrated critical path + pattern analysis as structured data."""
    cp_data = _analyze_critical_path_structured(parser, code, start_line, end_line)
    pattern_data = _analyze_patterns_structured(parser, code, start_line, end_line)

    # Cross-reference: for each bottleneck, find matching patterns + code generation
    identified = pattern_data.get("identified_patterns", {})
    enriched_bottlenecks = []
    for b in cp_data.get("bottlenecks", []):
        matches = []
        for pattern_name, pinfo in identified.items():
            for inst in pinfo.get("instances", []):
                if inst.get("lineno") == b.get("lineno"):
                    strategies = pinfo.get("recommended_partitioning", {}).get("strategies", ["default"])
                    match_entry = {
                        "pattern": pattern_name,
                        "confidence": inst.get("confidence"),
                        "partitioning": pinfo.get("recommended_partitioning", {}),
                        "code_generation": _try_generate_code(
                            b.get("source_code", code), pattern_name, strategies
                        ),
                    }
                    matches.append(match_entry)
        entry = dict(b)
        entry["suggested_patterns"] = matches
        enriched_bottlenecks.append(entry)

    return {
        "metrics": cp_data["metrics"],
        "bottlenecks": enriched_bottlenecks,
        "patterns": pattern_data["identified_patterns"],
        "dag_data": cp_data["dag_data"],
    }


def _analyze_full_structured(
    parser: BDiasParser,
    code: str,
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> Dict[str, Any]:
    """Return all analyses combined: blocks with suggestions + integrated (critical path + patterns)."""
    blocks_data = _analyze_blocks_with_suggestions_structured(parser, code, start_line, end_line)
    integrated_data = _analyze_integrated_structured(parser, code, start_line, end_line)

    return {
        "block_suggestions": blocks_data.get("suggestions", []),
        "metrics": integrated_data.get("metrics", {}),
        "bottlenecks": integrated_data.get("bottlenecks", []),
        "patterns": integrated_data.get("patterns", {}),
        "dag_data": integrated_data.get("dag_data", {}),
    }


# ---------------------------------------------------------------------------
# Markdown output functions (return strings — backward compatible)
# ---------------------------------------------------------------------------

def analyze_blocks(parser: BDiasParser, code: str) -> str:
    """Perform block-based analysis (Markdown output)."""
    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"]
    
    has_opportunities = any(structured_code.get(key) for key in structured_code)
    
    if not has_opportunities:
        return "No parallelization opportunities identified in the code."
    
    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("")
    
    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("")
    
    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("")
    
    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_blocks_with_suggestions(parser: BDiasParser, code: str) -> str:
    """Perform block-based analysis with code suggestions (Markdown output)."""
    structured_code = parser.parse(code)

    if structured_code is None:
        return "Error: Failed to parse code. Please check for syntax errors."

    code_gen = BDiasCodeGen(EXPLANATIONS, PARTITIONING_SUGGESTIONS)
    raw_suggestions = code_gen.generate_suggestions(structured_code)

    if not raw_suggestions:
        return "No parallelization opportunities identified in the code."

    output = ["## Block-Based Parallelization Analysis (with Suggestions)\n"]

    for i, s in enumerate(raw_suggestions, 1):
        lineno = s.get('lineno', '?')
        opp_type = s.get('opportunity_type', 'unknown')
        explanation_key = s.get('explanation_index', '')

        output.append(f"### {i}. **{opp_type.replace('_', ' ').title()}** (Line {lineno})\n")

        # Explanation
        try:
            explanation = EXPLANATIONS.get(explanation_key, '').format(**s)
        except (KeyError, IndexError):
            explanation = EXPLANATIONS.get(explanation_key, '')
        if explanation:
            output.append(f"**Explanation**: {explanation}\n")

        # Partitioning suggestion
        partitioning = PARTITIONING_SUGGESTIONS.get(explanation_key, '')
        if partitioning:
            output.append(f"**Partitioning**: {partitioning}\n")

        # Code suggestion
        code_suggestion = s.get('code_suggestion', '')
        if code_suggestion:
            output.append("**Suggested Parallel Code**:")
            output.append(f"```python\n{code_suggestion.strip()}\n```")

        output.append("")

    return "\n".join(output)


def analyze_full(parser: BDiasParser, code: str) -> str:
    """Perform full analysis: blocks with suggestions + integrated (Markdown output)."""
    output = []

    # Part 1: Block suggestions
    blocks_md = analyze_blocks_with_suggestions(parser, code)
    output.append(blocks_md)
    output.append("\n---\n")

    # Part 2: Integrated (critical path + patterns)
    integrated_md = analyze_integrated(parser, code)
    output.append(integrated_md)

    return "\n".join(output)


def analyze_critical_path(parser: BDiasParser, code: str) -> str:
    """Perform critical path analysis (Markdown output)."""
    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 = results.get('bottlenecks', [])
    if bottlenecks:
        output.append("### Critical Path Bottlenecks")
        for i, bottleneck in enumerate(bottlenecks[:10], 1):
            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.")
    
    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(parser: BDiasParser, code: str) -> str:
    """Perform pattern recognition analysis (Markdown output)."""
    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):
            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 = 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', '')}")
        
        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(f"- Synchronization Points: {perf.get('synchronization_points', 'Unknown')}")
        
        output.append("")
    
    return "\n".join(output)


def analyze_integrated(parser: BDiasParser, code: str) -> str:
    """Perform integrated critical path and pattern analysis (Markdown output)."""
    cp_analyzer = BDiasCriticalPathAnalyzer()
    cp_results = cp_analyzer.analyze(parser, code)
    
    pattern_analyzer = BDiasPatternAnalyzer(parser)
    pattern_analysis = pattern_analyzer.analyze(code)
    
    output = ["## Integrated Analysis: Critical Path + Pattern Recognition\n"]
    
    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")
    
    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[:10], 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")
            
            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)


# ---------------------------------------------------------------------------
# Main public API
# ---------------------------------------------------------------------------

def analyze_code(
    code_input: str,
    analysis_type: Literal["blocks", "blocks_with_suggestions", "critical_path", "patterns", "integrated"] = "integrated",
    output_format: Literal["structured", "markdown"] = "structured",
    start_line: Optional[int] = None,
    end_line: Optional[int] = None
) -> dict:
    """
    Analyze Python code to identify parallelization opportunities.
    
    Args:
        code_input: Python code as a string or path to a Python file to analyze
        analysis_type: Type of analysis to perform:
            - 'blocks': basic block-based opportunities
            - 'blocks_with_suggestions': block-based with Jinja2 template code suggestions
            - 'critical_path': DAG-based metrics
            - 'patterns': parallel pattern recognition
            - 'integrated': combined critical path and pattern analysis
        output_format: Output format:
            - 'structured': Returns dict with raw data (ideal for LLMs)
            - 'markdown': Returns formatted Markdown string (for CLI/humans)
        start_line: Optional start line for filtering results
        end_line: Optional end line for filtering results
            
    Returns:
        dict: Structured response with format:
            {"ok": True, "result": dict|str}  # Success
            {"ok": False, "error": str, "details": dict}  # Error
    """
    try:
        logger.info(f"Starting code analysis with type: {analysis_type}, format: {output_format}")
        
        code = load_code(code_input)
        
        parser = BDiasParser()
        structured_code = parser.parse(code)
        
        if structured_code is None:
            return {
                "ok": False,
                "error": "Failed to parse code. Please check for syntax errors.",
                "details": {"analysis_type": analysis_type}
            }
        
        valid_types = ["blocks", "blocks_with_suggestions", "critical_path", "patterns", "integrated", "full"]
        if analysis_type not in valid_types:
            return {
                "ok": False,
                "error": f"Invalid analysis_type '{analysis_type}'. Must be one of: {', '.join(valid_types)}",
                "details": {"analysis_type": analysis_type, "valid_types": valid_types}
            }

        if output_format == "structured":
            # Structured output (dicts)
            if analysis_type == "blocks":
                result = _analyze_blocks_structured(parser, code, start_line, end_line)
            elif analysis_type == "blocks_with_suggestions":
                result = _analyze_blocks_with_suggestions_structured(parser, code, start_line, end_line)
            elif analysis_type == "critical_path":
                result = _analyze_critical_path_structured(parser, code, start_line, end_line)
            elif analysis_type == "patterns":
                result = _analyze_patterns_structured(parser, code, start_line, end_line)
            elif analysis_type == "integrated":
                result = _analyze_integrated_structured(parser, code, start_line, end_line)
            elif analysis_type == "full":
                result = _analyze_full_structured(parser, code, start_line, end_line)
        else:
            if analysis_type == "blocks":
                result = analyze_blocks(parser, code)
            elif analysis_type == "blocks_with_suggestions":
                result = analyze_blocks_with_suggestions(parser, code)
            elif analysis_type == "critical_path":
                result = analyze_critical_path(parser, code)
            elif analysis_type == "patterns":
                result = analyze_patterns(parser, code)
            elif analysis_type == "integrated":
                result = analyze_integrated(parser, code)
            elif analysis_type == "full":
                result = analyze_full(parser, code)
        
        logger.info(f"Analysis completed successfully")
        return {"ok": True, "result": result}
        
    except FileNotFoundError as e:
        logger.error(f"File not found: {e}")
        return {
            "ok": False,
            "error": f"File not found: {str(e)}",
            "details": {"type": "FileNotFoundError", "code_input": code_input}
        }
    except Exception as e:
        logger.error(f"Error analyzing code: {e}", exc_info=True)
        return {
            "ok": False,
            "error": f"Error analyzing code: {str(e)}",
            "details": {"type": type(e).__name__, "analysis_type": analysis_type}
        }


# ---------------------------------------------------------------------------
# Code generation API
# ---------------------------------------------------------------------------

def auto_detect_pattern(
    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:
        for pattern_name, instances in identified.items():
            for inst in instances:
                if inst.get('lineno') == preferred_line:
                    return pattern_name, preferred_line
    
    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


def generate_parallel_code_logic(
    code_input: str,
    pattern: Optional[str] = None,
    line_number: Optional[int] = None
) -> dict:
    """
    Generate parallelized Python code from sequential code.
    
    Args:
        code_input: Python code as a string or path to a Python file
        pattern: Parallel pattern to use (e.g., 'map_reduce', 'pipeline', 'master_slave', 'pool_workers').
                If not specified, will auto-detect from code
        line_number: Specific line number to target for transformation.
                    If not specified, uses first detected pattern
                    
    Returns:
        dict: Structured response with format:
            {"ok": True, "result": str}  # Success with transformation output
            {"ok": False, "error": str, "details": dict}  # Error
    """
    try:
        logger.info(f"Starting parallel code generation with pattern: {pattern}, line: {line_number}")
        
        code = load_code(code_input)
        
        parser = BDiasParser()
        structured_code = parser.parse(code)
        
        if structured_code is None:
            return {
                "ok": False,
                "error": "Failed to parse code. Please check for syntax errors.",
                "details": {"pattern": pattern, "line_number": line_number}
            }
        
        if pattern is None:
            pattern, line_number = auto_detect_pattern(parser, code, line_number)
            if pattern is None:
                return {
                    "ok": False,
                    "error": "No parallelizable patterns detected in the code.",
                    "details": {"code_length": len(code)}
                }
            logger.info(f"Auto-detected pattern: {pattern} at line {line_number}")
        
        pattern_analyzer = BDiasPatternAnalyzer(parser)
        pattern_analysis = pattern_analyzer.analyze(code)
        
        recommended = pattern_analysis.get('recommended_partitioning', {})
        pattern_info = recommended.get(pattern, {})
        
        if not pattern_info:
            strategies = ['default']
        else:
            strategies = pattern_info.get('strategies', ['default'])
        
        bottleneck = {
            'source': code,
            'lineno': line_number or 1,
            'type': pattern
        }
        
        original_code, transformed_code, context = bdias_generate_parallel_code(
            bottleneck,
            pattern,
            strategies
        )
        
        pattern_display_info = {
            'pattern': pattern,
            'partitioning_strategy': strategies,
            'context': context
        }
        
        output = present_transformation(
            original_code,
            transformed_code,
            pattern_display_info
        )
        
        logger.info(f"Parallel code generation completed successfully")
        return {"ok": True, "result": output}
        
    except FileNotFoundError as e:
        logger.error(f"File not found: {e}")
        return {
            "ok": False,
            "error": f"File not found: {str(e)}",
            "details": {"type": "FileNotFoundError", "code_input": code_input}
        }
    except Exception as e:
        logger.error(f"Error generating parallel code: {e}", exc_info=True)
        return {
            "ok": False,
            "error": f"Error generating parallel code: {str(e)}",
            "details": {"type": type(e).__name__, "pattern": pattern, "line_number": line_number}
        }
