#!/bin/bash

# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Source the logging system using absolute path
source "${SCRIPT_DIR}/logging.sh"

# Source the Lil tester for secure code testing
source "${SCRIPT_DIR}/lil_tester.sh"

# Source the quality guard for output quality protection
source "${SCRIPT_DIR}/quality_guard.sh"

# Source the RAG integration for corpus queries
source "${SCRIPT_DIR}/rag_integration.sh"

# Get mechanism name automatically
MECHANISM_NAME=$(get_mechanism_name "$0")

# --- Lil Knowledge Base ---
get_lil_knowledge() {
    cat << 'EOF'
Lil is a multi-paradigm scripting language used in the Decker multimedia tool. Here is comprehensive knowledge about Lil:

## Core Language Features

### Types and Values
- **Numbers**: Floating-point values (42, 37.5, -29999)
- **Strings**: Double-quoted with escape sequences ("hello\nworld", "foo\"bar")
- **Lists**: Ordered sequences using comma: (1,2,3), empty list ()
- **Dictionaries**: Key-value pairs: ("a","b") dict (11,22) or {}.fruit:"apple"
- **Tables**: Rectangular data with named columns (created with insert or table)
- **Functions**: Defined with 'on name do ... end', called with []
- **Interfaces**: Opaque values for system resources

### Basic Syntax
- **Variables**: Assignment uses ':' (x:42, y:"hello")
- **Indexing**: Lists with [index], dicts with .key or ["key"]
- **Operators**: Right-to-left precedence, 2*3+5 = 16, (2*3)+5 = 11
- **Comments**: # line comments only
- **Expressions**: Everything is an expression, returns values

### Control Flow
- **Conditionals**: if condition ... end, with optional elseif/else
- **Loops**: each value key index in collection ... end
- **While**: while condition ... end
- **Functions**: on func x y do ... end, called as func[args]

### Query Language (SQL-like)
- **select**: select columns from table where condition orderby column
- **update**: update column:value from table where condition
- **extract**: extract values from table (returns simple types)
- **insert**: insert columns with values end
- **Clauses**: where, by, orderby asc/desc
- **Joins**: table1 join table2 (natural), table1 cross table2 (cartesian)

### Vector Operations (Conforming)
- **Arithmetic spreads**: 1+(2,3,4) = (3,4,5)
- **List operations**: (1,2,3)+(4,5,6) = (5,7,9)
- **Equality**: 5=(1,5,10) = (0,1,0)  # Use ~ for exact match
- **Application**: func @ (1,2,3) applies func to each element

### Key Operators
- **Arithmetic**: + - * / % ^ & | (min/max)
- **Comparison**: < > = (conforming), ~ (exact match)
- **Logic**: ! (not), & | (and/or with numbers)
- **Data**: , (concat), @ (index each), dict, take, drop
- **String**: fuse (concat), split, format, parse, like (glob matching)
- **Query**: join, cross, limit, window

### Important Patterns
- **Function definition**: on add x y do x+y end
- **List comprehension**: each x in data x*2 end
- **Table query**: select name age where age>21 from people
- **Dictionary building**: d:(); d.key:"value"
- **String formatting**: "%i %s" format (42,"answer")

### Common Functions
- **Math**: cos, sin, tan, exp, ln, sqrt, floor, mag, unit, heading
- **Aggregation**: sum, prod, min, max, count, first, last
- **Data**: range, keys, list, table, flip, raze, typeof
- **IO**: read, write, show, print (in Lilt environment)

### Best Practices
- Use functional style when possible (immutable operations)
- Leverage vector operations for data manipulation
- Use queries for complex data transformations
- Functions are first-class values
- Lexical scoping with closures
- Tail-call optimization supported

### Common Patterns
- **Mode finding**: extract first value by value orderby count value desc from data
- **Filtering**: select from table where condition
- **Grouping**: select agg_func column by group_column from table
- **List processing**: each x in data transform[x] end
- **Dictionary operations**: keys dict, range dict, dict operations
- **String manipulation**: split, fuse, format, parse, like

Lil emphasizes expressive, concise code with powerful built-in operations for data manipulation, making it excellent for algorithmic puzzles and data processing tasks.
EOF
}

# --- Model Configuration ---
PUZZLE_MODEL="llama3:8b-instruct-q4_K_M"
ANALYSIS_MODEL="phi3:3.8b-mini-4k-instruct-q4_K_M"

# Validate and set models with fallbacks
PUZZLE_MODEL=$(validate_model "$PUZZLE_MODEL" "gemma3n:e2b")
if [ $? -ne 0 ]; then
    log_error "No valid puzzle model available"
    exit 1
fi

ANALYSIS_MODEL=$(validate_model "$ANALYSIS_MODEL" "gemma3n:e2b")
if [ $? -ne 0 ]; then
    log_error "No valid analysis model available"
    exit 1
fi

# --- Defaults ---
DEFAULT_ROUNDS=2

# --- Argument Validation ---
if [ "$#" -lt 1 ]; then
    echo -e "\n\tPuzzle"
    echo -e "\tThis script specializes in puzzle solving and coding challenges with Lil programming language expertise."
    echo -e "\n\tUsage: $0 [-f <file_path>] [-l <language>] \"<your puzzle/challenge>\" [number_of_rounds]"
    echo -e "\n\tExample: $0 -f ./challenge.txt -l lil \"How can I implement a sorting algorithm?\" 2"
    echo -e "\n\tIf number_of_rounds is not provided, the program will default to 2 rounds."
    echo -e "\n\t-f <file_path> (optional): Append the contents of the file to the prompt."
    echo -e "\n\t-l <language> (optional): Specify programming language focus (default: lil)."
    echo -e "\n"
    exit 1
fi

# --- Argument Parsing ---
FILE_PATH=""
LANGUAGE="lil"
while getopts "f:l:" opt; do
  case $opt in
    f)
      FILE_PATH="$OPTARG"
      ;;
    l)
      LANGUAGE="$OPTARG"
      ;;
    *)
      log_error "Invalid option: -$OPTARG"
      exit 1
      ;;
  esac
done
shift $((OPTIND -1))

PROMPT="$1"
if [ -z "$2" ]; then
    ROUNDS=2
else
    ROUNDS=$2
fi

# Validate prompt
PROMPT=$(validate_prompt "$PROMPT")
if [ $? -ne 0 ]; then
    exit 1
fi

# Validate file path if provided
if [ -n "$FILE_PATH" ]; then
    if ! validate_file_path "$FILE_PATH"; then
        exit 1
    fi
    FILE_CONTENTS=$(cat "$FILE_PATH")
    PROMPT="$PROMPT\n[FILE CONTENTS]\n$FILE_CONTENTS\n[END FILE]"
fi

# Validate rounds
if ! [[ "$ROUNDS" =~ ^[1-9][0-9]*$ ]] || [ "$ROUNDS" -gt 5 ]; then
    log_error "Invalid number of rounds: $ROUNDS (must be 1-5)"
    exit 1
fi

# --- File Initialization ---
mkdir -p ~/tmp
SESSION_FILE=~/tmp/puzzle_$(date +%Y%m%d_%H%M%S).txt

# Initialize timing
SESSION_ID=$(generate_session_id)
start_timer "$SESSION_ID" "puzzle"

echo "Puzzle Session Log: ${SESSION_FILE}"
echo "---------------------------------"

# Store the initial user prompt in the session file
echo "USER PROMPT: ${PROMPT}" >> "${SESSION_FILE}"
echo "LANGUAGE FOCUS: ${LANGUAGE}" >> "${SESSION_FILE}"
echo "NUMBER OF ROUNDS: ${ROUNDS}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Phase 1: Problem Analysis ---
echo "Phase 1: Analyzing the puzzle or coding challenge..."
echo "PHASE 1 - PROBLEM ANALYSIS:" >> "${SESSION_FILE}"

# Check for RAG context
RAG_CONTEXT=$(use_rag_if_available "${PROMPT}" "${MECHANISM_NAME}")

PROBLEM_ANALYSIS_PROMPT="You are an expert puzzle solver and programming mentor specializing in the ${LANGUAGE} programming language. Analyze the following puzzle or coding challenge.

CHALLENGE: ${PROMPT}

$(if [[ "$LANGUAGE" == "lil" ]]; then
echo "LIL PROGRAMMING LANGUAGE KNOWLEDGE:
$(get_lil_knowledge)

Use this comprehensive knowledge of Lil to provide accurate, helpful analysis."
fi)

$(if [[ "$RAG_CONTEXT" != "$PROMPT" ]]; then
echo "ADDITIONAL CONTEXT FROM KNOWLEDGE BASE:
$RAG_CONTEXT

Use this context to enhance your analysis if it's relevant to the challenge."
fi)

Please provide a comprehensive analysis including:
1. Problem type classification (algorithm, data structure, logic, etc.)
2. Complexity assessment (time/space requirements)
3. Key concepts and patterns involved
4. Relevant ${LANGUAGE} language features that could help
5. Potential solution approaches
6. Common pitfalls or edge cases to consider

Provide a clear, structured analysis that helps understand the problem."

problem_analysis=$(ollama run "${PUZZLE_MODEL}" "${PROBLEM_ANALYSIS_PROMPT}")
problem_analysis=$(guard_output_quality "$problem_analysis" "$PROMPT" "$MECHANISM_NAME" "$PUZZLE_MODEL")



echo "PROBLEM ANALYSIS:" >> "${SESSION_FILE}"
echo "${problem_analysis}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Phase 2: Solution Strategy ---
echo "Phase 2: Developing solution strategies..."
echo "PHASE 2 - SOLUTION STRATEGY:" >> "${SESSION_FILE}"

SOLUTION_STRATEGY_PROMPT="Based on the problem analysis, develop multiple solution strategies for this puzzle or coding challenge.

ORIGINAL CHALLENGE: ${PROMPT}

PROBLEM ANALYSIS: ${problem_analysis}

$(if [[ "$LANGUAGE" == "lil" ]]; then
echo "LIL PROGRAMMING LANGUAGE KNOWLEDGE:
$(get_lil_knowledge)

Use this comprehensive knowledge of Lil to provide accurate, helpful analysis."
fi)

$(if [[ "$RAG_CONTEXT" != "$PROMPT" ]]; then
echo "ADDITIONAL CONTEXT FROM KNOWLEDGE BASE:
$RAG_CONTEXT

Use this context to develop more informed and accurate solution strategies."
fi)

Please provide:
1. At least 2-3 different solution approaches
2. Algorithmic complexity analysis for each approach
3. Trade-offs between different solutions
4. Specific ${LANGUAGE} language constructs that would be useful
5. Implementation considerations and challenges
6. Testing and validation strategies

Focus on practical, implementable solutions with clear reasoning."

solution_strategy=$(ollama run "${PUZZLE_MODEL}" "${SOLUTION_STRATEGY_PROMPT}")
solution_strategy=$(guard_output_quality "$solution_strategy" "$PROMPT" "$MECHANISM_NAME" "$PUZZLE_MODEL")



echo "SOLUTION STRATEGY:" >> "${SESSION_FILE}"
echo "${solution_strategy}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Phase 3: Implementation Guidance ---
echo "Phase 3: Providing implementation guidance..."
echo "PHASE 3 - IMPLEMENTATION GUIDANCE:" >> "${SESSION_FILE}"

IMPLEMENTATION_PROMPT="Provide detailed implementation guidance for the best solution approach to this puzzle or coding challenge.

ORIGINAL CHALLENGE: ${PROMPT}
PROBLEM ANALYSIS: ${problem_analysis}
SOLUTION STRATEGY: ${solution_strategy}

$(if [[ "$LANGUAGE" == "lil" ]]; then
echo "LIL PROGRAMMING LANGUAGE KNOWLEDGE:
$(get_lil_knowledge)

Use this comprehensive knowledge of Lil to provide accurate, helpful analysis."
fi)

$(if [[ "$RAG_CONTEXT" != "$PROMPT" ]]; then
echo "ADDITIONAL CONTEXT FROM KNOWLEDGE BASE:
$RAG_CONTEXT

Use this context to provide more accurate and comprehensive implementation guidance."
fi)

Please provide:
1. Step-by-step implementation plan
2. Complete code example in ${LANGUAGE} (if applicable)
3. Explanation of key code sections and patterns
4. Variable naming and structure recommendations
5. Error handling and edge case considerations
6. Performance optimization tips
7. Testing and debugging guidance

Make the implementation clear and educational, explaining the reasoning behind each decision."

implementation_guidance=$(ollama run "${PUZZLE_MODEL}" "${IMPLEMENTATION_PROMPT}")
implementation_guidance=$(guard_output_quality "$implementation_guidance" "$PROMPT" "$MECHANISM_NAME" "$PUZZLE_MODEL")



echo "IMPLEMENTATION GUIDANCE:" >> "${SESSION_FILE}"
echo "${implementation_guidance}" >> "${SESSION_FILE}"
echo "" >> "${SESSION_FILE}"

# --- Phase 4: Code Testing (if applicable) ---
if [[ "$LANGUAGE" == "lil" ]] && [[ "$implementation_guidance" =~ (on [a-zA-Z_][a-zA-Z0-9_]* do|function|procedure) ]]; then
    echo "Phase 4: Testing the Lil code implementation..."
    echo "PHASE 4 - CODE TESTING:" >> "${SESSION_FILE}"
    
    # Simple code extraction - look for function definitions
    lil_code=""
    if echo "$implementation_guidance" | grep -q "on [a-zA-Z_][a-zA-Z0-9_]* do"; then
        lil_code=$(echo "$implementation_guidance" | grep -A 20 "on [a-zA-Z_][a-zA-Z0-9_]* do" | head -20)
    fi
    
    if [ -n "$lil_code" ]; then
        echo "Extracted Lil code for testing:"
        echo "----------------------------------------"
        echo "$lil_code"
        echo "----------------------------------------"
        
        # Test the extracted code
        test_result=$(test_lil_script "$lil_code" "puzzle_implementation")
        test_exit_code=$?
        
        echo "CODE TESTING RESULTS:" >> "${SESSION_FILE}"
        echo "$test_result" >> "${SESSION_FILE}"
        echo "" >> "${SESSION_FILE}"
        
        if [ $test_exit_code -eq 0 ]; then
            echo "✅ Lil code testing PASSED"
        else
            echo "❌ Lil code testing FAILED"
            echo "Note: The code may have syntax errors or runtime issues."
        fi
    else
        echo "No executable Lil code found in implementation guidance."
        echo "CODE TESTING: No executable code found" >> "${SESSION_FILE}"
    fi
else
    echo "Phase 4: Skipping code testing (not Lil language or no executable code)"
    echo "CODE TESTING: Skipped (not applicable)" >> "${SESSION_FILE}"
fi

# --- Phase 5: Solution Validation ---
echo "Phase 5: Validating and reviewing the solution..."
echo "PHASE 5 - SOLUTION VALIDATION:" >> "${SESSION_FILE}"

VALIDATION_PROMPT="Review and validate the proposed solution to ensure it's correct, efficient, and well-implemented.

ORIGINAL CHALLENGE: ${PROMPT}
PROBLEM ANALYSIS: ${problem_analysis}
SOLUTION STRATEGY: ${solution_strategy}
IMPLEMENTATION: ${implementation_guidance}

Please provide:
1. Code review and correctness verification
2. Edge case analysis and testing scenarios
3. Performance analysis and optimization opportunities
4. Alternative approaches or improvements
5. Common mistakes to avoid
6. Learning resources and next steps
7. Final recommendations and best practices

Ensure the solution is robust, maintainable, and follows ${LANGUAGE} best practices."

solution_validation=$(ollama run "${ANALYSIS_MODEL}" "${VALIDATION_PROMPT}")
solution_validation=$(guard_output_quality "$solution_validation" "$PROMPT" "$MECHANISM_NAME" "$ANALYSIS_MODEL")



echo "SOLUTION VALIDATION:" >> "${SESSION_FILE}"
echo "${solution_validation}" >> "${SESSION_FILE}"

# End timing
duration=$(end_timer "$SESSION_ID" "puzzle")

# --- Final Output ---
echo "---------------------------------"
echo "Puzzle-solving process complete."
echo "---------------------------------"
echo ""
echo "PROBLEM ANALYSIS:"
echo "================="
echo "${problem_analysis}"
echo ""
echo "SOLUTION STRATEGY:"
echo "=================="
echo "${solution_strategy}"
echo ""
echo "IMPLEMENTATION GUIDANCE:"
echo "========================"
echo "${implementation_guidance}"
echo ""
if [[ "$LANGUAGE" == "lil" ]] && [[ "$implementation_guidance" =~ (on [a-zA-Z_][a-zA-Z0-9_]* do|function|procedure) ]]; then
    echo "CODE TESTING:"
    echo "============="
    if [ -n "$lil_code" ]; then
        echo "✅ Lil code was tested successfully"
        echo "Test results logged in session file"
    else
        echo "No executable code found for testing"
    fi
    echo ""
fi
echo "SOLUTION VALIDATION:"
echo "===================="
echo "${solution_validation}"
echo ""
echo "Language focus: ${LANGUAGE}"
echo "Rounds completed: ${ROUNDS}"
echo "Execution time: ${duration} seconds"
echo ""
echo "Full puzzle-solving log: ${SESSION_FILE}"