#!/usr/bin/env node

// fmt.js - Baba Yaga code formatter
// Similar to Go's fmt tool, formats Baba Yaga source code according to standard style

import { createLexer } from './lexer.js';
import { createParser } from './parser.js';
import fs from 'fs';
import path from 'path';

/**
 * Baba Yaga code formatter
 * Formats code according to consistent style rules
 */
class BabaYagaFormatter {
  constructor(options = {}) {
    this.indentSize = options.indentSize || 2;
    this.maxLineLength = options.maxLineLength || 100;
    this.preserveComments = options.preserveComments !== false;
  }

  /**
   * Format source code string
   */
  format(source) {
    try {
      const lexer = createLexer(source);
      const tokens = lexer.allTokens();
      
      // Extract comments before parsing
      const comments = this.extractComments(source);
      
      const parser = createParser(tokens);
      const ast = parser.parse();
      
      return this.formatAST(ast, comments, source);
    } catch (error) {
      throw new Error(`Formatting failed: ${error.message}`);
    }
  }

  /**
   * Extract comments from source with their positions
   */
  extractComments(source) {
    const comments = [];
    const lines = source.split('\n');
    
    lines.forEach((line, lineIndex) => {
      const commentMatch = line.match(/\/\/(.*)$/);
      if (commentMatch) {
        const column = line.indexOf('//');
        comments.push({
          line: lineIndex + 1,
          column: column,
          text: commentMatch[0],
          content: commentMatch[1].trim()
        });
      }
    });
    
    return comments;
  }

  /**
   * Format AST node
   */
  formatAST(ast, comments = [], originalSource = '') {
    return this.visitNode(ast, 0, comments);
  }

  /**
   * Visit and format a node
   */
  visitNode(node, depth = 0, comments = []) {
    if (!node) return '';

    switch (node.type) {
      case 'Program':
        return this.formatProgram(node, depth, comments);
      case 'TypeDeclaration':
        return this.formatTypeDeclaration(node, depth);
      case 'VariableDeclaration':
        return this.formatVariableDeclaration(node, depth, comments);
      case 'FunctionDeclaration':
        return this.formatFunctionDeclaration(node, depth, comments);
      case 'CurriedFunctionDeclaration':
        return this.formatCurriedFunctionDeclaration(node, depth, comments);
      case 'WithHeader':
        return this.formatWithHeader(node, depth, comments);
      case 'WhenExpression':
        return this.formatWhenExpression(node, depth, comments);
      case 'BinaryExpression':
        return this.formatBinaryExpression(node, depth, comments);
      case 'UnaryExpression':
        return this.formatUnaryExpression(node, depth, comments);
      case 'FunctionCall':
        return this.formatFunctionCall(node, depth, comments);
      case 'AnonymousFunction':
        return this.formatAnonymousFunction(node, depth, comments);
      case 'ListLiteral':
        return this.formatListLiteral(node, depth, comments);
      case 'TableLiteral':
        return this.formatTableLiteral(node, depth, comments);
      case 'MemberExpression':
        return this.formatMemberExpression(node, depth, comments);
      case 'ResultExpression':
        return this.formatResultExpression(node, depth, comments);
      case 'NumberLiteral':
        return this.formatNumberLiteral(node);
      case 'StringLiteral':
        return this.formatStringLiteral(node);
      case 'BooleanLiteral':
        return this.formatBooleanLiteral(node);
      case 'Identifier':
        return this.formatIdentifier(node);
      default:
        // Fallback for unknown node types - avoid infinite recursion
        if (typeof node === 'string') {
          return node;
        }
        if (typeof node === 'number') {
          return node.toString();
        }
        if (typeof node === 'boolean') {
          return node.toString();
        }
        if (node && typeof node === 'object') {
          // Try to handle as a literal value
          if (node.value !== undefined) {
            return node.value.toString();
          }
          if (node.name !== undefined) {
            return node.name;
          }
        }
        return JSON.stringify(node);
    }
  }

  /**
   * Format program (top level)
   */
  formatProgram(node, depth, comments) {
    const statements = [];
    let lastWasFunction = false;

    node.body.forEach((stmt, index) => {
      const formatted = this.visitNode(stmt, depth, comments);
      const isFunction = stmt.type === 'FunctionDeclaration' || 
                        stmt.type === 'CurriedFunctionDeclaration';
      
      // Add extra spacing between functions and other statements
      if (index > 0 && (isFunction || lastWasFunction)) {
        statements.push('');
      }
      
      statements.push(formatted);
      lastWasFunction = isFunction;
    });

    return statements.join('\n') + (statements.length > 0 ? '\n' : '');
  }

  /**
   * Format type declaration
   */
  formatTypeDeclaration(node, depth) {
    const indent = this.getIndent(depth);
    return `${indent}${node.name} ${node.typeAnnotation};`;
  }

  /**
   * Format variable declaration
   */
  formatVariableDeclaration(node, depth, comments) {
    const indent = this.getIndent(depth);
    
    // Check if the value is a complex expression that should be on its own line
    if (node.value.type === 'WhenExpression' || node.value.type === 'WithHeader') {
      const value = this.visitNode(node.value, depth + 1, comments);
      return `${indent}${node.name} :\n${value};`;
    } else {
      const value = this.visitNode(node.value, depth, comments);
      return `${indent}${node.name} : ${value};`;
    }
  }

  /**
   * Format function declaration
   */
  formatFunctionDeclaration(node, depth, comments) {
    const indent = this.getIndent(depth);
    let result = `${indent}${node.name} : `;
    
    // Format parameters
    if (node.params && node.params.length > 0) {
      if (this.hasTypedParams(node.params)) {
        result += this.formatTypedParameters(node.params);
      } else {
        result += node.params.map(p => typeof p === 'string' ? p : p.name).join(' ');
      }
    }
    
    // Add return type if present
    if (node.returnType) {
      result += ` -> ${this.formatType(node.returnType)}`;
    }
    
    result += ' ->\n';
    
    // Format body with proper indentation
    const body = this.visitNode(node.body, depth + 1, comments);
    // If the body doesn't start with indentation, add it
    if (body && !body.startsWith(this.getIndent(depth + 1))) {
      result += this.getIndent(depth + 1) + body;
    } else {
      result += body;
    }
    
    result += ';';
    return result;
  }

  /**
   * Format curried function declaration
   */
  formatCurriedFunctionDeclaration(node, depth, comments) {
    const indent = this.getIndent(depth);
    let result = `${indent}${node.name} : `;
    
    // Format first typed parameter
    result += `(${node.param.name}: ${this.formatType(node.param.type)})`;
    
    // Format return type
    if (node.returnType) {
      result += ` -> ${this.formatType(node.returnType)}`;
    }
    
    result += ' ->\n';
    
    // Format body with proper indentation
    const body = this.visitNode(node.body, depth + 1, comments);
    result += body + ';';
    
    return result;
  }

  /**
   * Format with header
   */
  formatWithHeader(node, depth, comments) {
    const indent = this.getIndent(depth);
    let result = `${indent}with`;
    
    if (node.recursive) {
      result += ' rec';
    }
    
    result += ' (\n';
    
    // Format entries
    node.entries.forEach((entry, index) => {
      const entryIndent = this.getIndent(depth + 1);
      if (entry.type === 'WithTypeDecl') {
        result += `${entryIndent}${entry.name} ${this.formatType(entry.typeAnnotation)};`;
      } else if (entry.type === 'WithAssign') {
        const value = this.visitNode(entry.value, depth + 1, comments);
        result += `${entryIndent}${entry.name} : ${value};`;
      }
      
      if (index < node.entries.length - 1) {
        result += '\n';
      }
    });
    
    result += `\n${indent}) ->\n`;
    const body = this.visitNode(node.body, depth + 1, comments);
    // Ensure the body is properly indented
    if (body && !body.startsWith(this.getIndent(depth + 1))) {
      result += this.getIndent(depth + 1) + body;
    } else {
      result += body;
    }
    
    return result;
  }

  /**
   * Format when expression
   */
  formatWhenExpression(node, depth, comments) {
    const indent = this.getIndent(depth);
    let result = `${indent}when `;
    
    // Format discriminants
    const discriminants = node.discriminants.map(d => 
      this.visitNode(d, 0, comments)
    ).join(' ');
    result += `${discriminants} is\n`;
    
    // Calculate the maximum pattern width to align 'then' keywords
    const caseIndent = this.getIndent(depth + 1);
    const formattedCases = node.cases.map(caseNode => {
      const patterns = caseNode.patterns.map(p => 
        this.formatPattern(p, depth + 1, comments)
      ).join(' ');
      return {
        patterns,
        consequent: caseNode.consequent,
        originalCase: caseNode
      };
    });
    
    // Find the maximum pattern length for alignment
    const maxPatternLength = Math.max(
      ...formattedCases.map(c => c.patterns.length)
    );
    
    // Format cases with aligned 'then' keywords
    formattedCases.forEach((formattedCase, index) => {
      const { patterns, consequent } = formattedCase;
      
      // Pad patterns to align 'then' keywords
      const paddedPatterns = patterns.padEnd(maxPatternLength);
      result += `${caseIndent}${paddedPatterns} then `;
      
      // Format consequent - handle nested when expressions specially
      if (consequent.type === 'WhenExpression') {
        // For nested when expressions, add newline and proper indentation
        result += '\n' + this.visitNode(consequent, depth + 2, comments);
      } else {
        // For simple consequents, add inline
        const consequentFormatted = this.visitNode(consequent, 0, comments);
        result += consequentFormatted;
      }
      
      // Add newline between cases (but not after the last one)
      if (index < formattedCases.length - 1) {
        result += '\n';
      }
    });
    
    return result;
  }

  /**
   * Format pattern
   */
  formatPattern(pattern, depth, comments) {
    if (!pattern) return '';
    
    switch (pattern.type) {
      case 'WildcardPattern':
        return '_';
      case 'TypePattern':
        return pattern.name;
      case 'ResultPattern':
        return `${pattern.variant} ${pattern.identifier.name}`;
      case 'ListPattern':
        const elements = pattern.elements.map(e => 
          this.formatPattern(e, depth, comments)
        ).join(', ');
        return `[${elements}]`;
      case 'TablePattern':
        const properties = pattern.properties.map(prop => 
          `${prop.key}: ${this.formatPattern(prop.value, depth, comments)}`
        ).join(', ');
        return `{${properties}}`;
      case 'NumberLiteral':
        return pattern.value.toString();
      case 'StringLiteral':
        return `"${pattern.value}"`;
      case 'BooleanLiteral':
        return pattern.value.toString();
      case 'Identifier':
        return pattern.name;
      default:
        // For literal patterns, try to format them directly
        if (typeof pattern === 'string') {
          return pattern;
        }
        if (typeof pattern === 'number') {
          return pattern.toString();
        }
        return this.visitNode(pattern, depth, comments);
    }
  }

  /**
   * Format binary expression
   */
  formatBinaryExpression(node, depth, comments) {
    const left = this.visitNode(node.left, depth, comments);
    const right = this.visitNode(node.right, depth, comments);
    
    // Add spaces around operators
    const needsSpaces = !['.', '..'].includes(node.operator);
    if (needsSpaces) {
      return `${left} ${node.operator} ${right}`;
    } else {
      return `${left}${node.operator}${right}`;
    }
  }

  /**
   * Format unary expression
   */
  formatUnaryExpression(node, depth, comments) {
    const operand = this.visitNode(node.operand, depth, comments);
    return `${node.operator}${operand}`;
  }

  /**
   * Format function call
   */
  formatFunctionCall(node, depth, comments) {
    const callee = this.visitNode(node.callee, depth, comments);
    const args = node.arguments.map(arg => 
      this.visitNode(arg, depth, comments)
    );
    
    if (args.length === 0) {
      return callee;
    }
    
    // Handle parentheses for complex expressions
    const formattedArgs = args.map(arg => {
      // If argument contains operators or is complex, wrap in parentheses
      if (arg.includes(' -> ') || (arg.includes(' ') && !arg.startsWith('"') && !arg.startsWith('['))) {
        return `(${arg})`;
      }
      return arg;
    });
    
    return `${callee} ${formattedArgs.join(' ')}`;
  }

  /**
   * Format anonymous function
   */
  formatAnonymousFunction(node, depth, comments) {
    // Handle both string parameters and object parameters
    const params = node.params.map(param => {
      if (typeof param === 'string') {
        return param;
      } else if (param && typeof param === 'object' && param.name) {
        return param.name;
      } else if (param && typeof param === 'object' && param.type === 'Identifier') {
        return param.name;
      } else {
        return String(param);
      }
    }).join(' ');
    const body = this.visitNode(node.body, depth, comments);
    return `${params} -> ${body}`;
  }

  /**
   * Format list literal
   */
  formatListLiteral(node, depth, comments) {
    if (node.elements.length === 0) {
      return '[]';
    }
    
    const elements = node.elements.map(el => 
      this.visitNode(el, depth, comments)
    );
    
    // Single line if short, multi-line if long
    const singleLine = `[${elements.join(', ')}]`;
    if (singleLine.length <= 50) {
      return singleLine;
    }
    
    const indent = this.getIndent(depth);
    const elementIndent = this.getIndent(depth + 1);
    let result = '[\n';
    elements.forEach((el, index) => {
      result += `${elementIndent}${el}`;
      if (index < elements.length - 1) {
        result += ',';
      }
      result += '\n';
    });
    result += `${indent}]`;
    return result;
  }

  /**
   * Format table literal
   */
  formatTableLiteral(node, depth, comments) {
    if (node.properties.length === 0) {
      return '{}';
    }
    
    const properties = node.properties.map(prop => {
      const value = this.visitNode(prop.value, depth + 1, comments);
      return `${prop.key}: ${value}`;
    });
    
    // Single line if short, multi-line if long
    const singleLine = `{${properties.join(', ')}}`;
    if (singleLine.length <= 50 && !properties.some(p => p.includes('\n'))) {
      return singleLine;
    }
    
    const indent = this.getIndent(depth);
    const propIndent = this.getIndent(depth + 1);
    let result = '{\n';
    properties.forEach((prop, index) => {
      result += `${propIndent}${prop}`;
      if (index < properties.length - 1) {
        result += ',';
      }
      result += '\n';
    });
    result += `${indent}}`;
    return result;
  }

  /**
   * Format member expression
   */
  formatMemberExpression(node, depth, comments) {
    const object = this.visitNode(node.object, depth, comments);
    const property = this.visitNode(node.property, depth, comments);
    return `${object}.${property}`;
  }

  /**
   * Format result expression
   */
  formatResultExpression(node, depth, comments) {
    const value = this.visitNode(node.value, depth, comments);
    return `${node.variant} ${value}`;
  }

  /**
   * Format number literal
   */
  formatNumberLiteral(node) {
    return node.value.toString();
  }

  /**
   * Format string literal
   */
  formatStringLiteral(node) {
    return `"${node.value}"`;
  }

  /**
   * Format boolean literal
   */
  formatBooleanLiteral(node) {
    return node.value.toString();
  }

  /**
   * Format identifier
   */
  formatIdentifier(node) {
    return node.name;
  }

  // Helper methods

  /**
   * Get indentation string
   */
  getIndent(depth) {
    return ' '.repeat(depth * this.indentSize);
  }

  /**
   * Check if parameters have type annotations
   */
  hasTypedParams(params) {
    return params.some(p => 
      typeof p === 'object' && p.type && p.type !== 'Identifier'
    );
  }

  /**
   * Format typed parameters
   */
  formatTypedParameters(params) {
    const formatted = params.map(p => {
      if (typeof p === 'string') {
        return p;
      } else if (p.type && p.type !== 'Identifier') {
        return `${p.name}: ${this.formatType(p.type)}`;
      } else {
        return p.name;
      }
    });
    return `(${formatted.join(', ')})`;
  }

  /**
   * Format type annotation
   */
  formatType(type) {
    if (typeof type === 'string') {
      return type;
    }
    
    if (type.type === 'PrimitiveType') {
      return type.name;
    }
    
    if (type.type === 'FunctionType') {
      const paramTypes = type.paramTypes.map(t => this.formatType(t)).join(', ');
      const returnType = this.formatType(type.returnType);
      return `(${paramTypes}) -> ${returnType}`;
    }
    
    return 'Unknown';
  }
}

/**
 * CLI interface
 */
function main() {
  const args = process.argv.slice(2);
  
  if (args.length === 0) {
    console.error('Usage: node fmt.js <file.baba> [options]');
    console.error('Options:');
    console.error('  --write, -w    Write result to file instead of stdout');
    console.error('  --check, -c    Check if file is already formatted');
    console.error('  --indent=N     Set indentation size (default: 2)');
    process.exit(1);
  }
  
  const options = {
    write: args.includes('--write') || args.includes('-w'),
    check: args.includes('--check') || args.includes('-c'),
    indentSize: 2
  };
  
  // Parse indent option
  const indentArg = args.find(arg => arg.startsWith('--indent='));
  if (indentArg) {
    options.indentSize = parseInt(indentArg.split('=')[1]) || 2;
  }
  
  const filename = args.find(arg => 
    !arg.startsWith('-') && !arg.startsWith('--')
  );
  
  if (!filename) {
    console.error('Error: No input file specified');
    process.exit(1);
  }
  
  if (!fs.existsSync(filename)) {
    console.error(`Error: File '${filename}' not found`);
    process.exit(1);
  }
  
  try {
    const source = fs.readFileSync(filename, 'utf8');
    const formatter = new BabaYagaFormatter(options);
    const formatted = formatter.format(source);
    
    if (options.check) {
      if (source.trim() !== formatted.trim()) {
        console.error(`File '${filename}' is not formatted`);
        process.exit(1);
      } else {
        console.log(`File '${filename}' is already formatted`);
        process.exit(0);
      }
    }
    
    if (options.write) {
      fs.writeFileSync(filename, formatted);
      console.log(`Formatted '${filename}'`);
    } else {
      process.stdout.write(formatted);
    }
    
  } catch (error) {
    console.error(`Error formatting '${filename}': ${error.message}`);
    process.exit(1);
  }
}

// Export for use as module
export { BabaYagaFormatter };

// Run CLI if called directly
if (import.meta.url === `file://${process.argv[1]}`) {
  main();
}