from flask import request, jsonify
import logging
from routes import app
import re
from collections import deque

logger = logging.getLogger(__name__)

from flask import request, jsonify

@app.route('/lab_work', methods=['POST'])
def lab_work():
    data = request.get_json()

    # Adjusting to handle both list and dict inputs
    if isinstance(data, dict):
        # If data is a dict, extract the markdown tables
        test_cases = []
        for value in data.values():
            if isinstance(value, list):
                test_cases.extend(value)
            else:
                test_cases.append(value)
    else:
        test_cases = data

    output = []

    for test_case in test_cases:
        labs = parse_markdown_table(test_case)
        inspections = simulate_labs(labs)
        output.append(inspections)

    return jsonify(output)


def parse_markdown_table(test_case):
    import re

    labs = {}
    lines = test_case.strip().split('\n')

    # Skip header lines
    data_lines = [line for line in lines if not re.match(r'^\|[-\s]*\|$', line) and not line.startswith('|Lab')]

    for line in data_lines:
        # Split the line into columns
        cols = [col.strip() for col in line.strip('|').split('|')]

        if len(cols) != 4:
            continue  # Skip malformed lines

        lab_number = int(cols[0])
        cell_counts = list(map(int, cols[1].split()))
        increment_str = cols[2]
        condition_str = cols[3]

        # Parse increment operation
        increment_op = parse_increment(increment_str)

        # Parse condition
        divisor, true_lab, false_lab = map(int, condition_str.split())

        labs[lab_number] = {
            'cell_counts': cell_counts,
            'increment_op': increment_op,
            'divisor': divisor,
            'true_lab': true_lab,
            'false_lab': false_lab,
            'inspections': 0,
            'queue': cell_counts.copy()
        }

    return labs


def parse_increment(increment_str):
    import re
    # Example increments: 'count * 2', 'count + 5', 'count * count'
    if 'count * count' in increment_str:
        return lambda count: count * count
    else:
        match = re.match(r'count ([\*\+]) (\d+|count)', increment_str)
        if match:
            operator, operand = match.groups()
            if operand == 'count':
                if operator == '*':
                    return lambda count: count * count
                elif operator == '+':
                    return lambda count: count + count
            else:
                operand = int(operand)
                if operator == '*':
                    return lambda count: count * operand
                elif operator == '+':
                    return lambda count: count + operand
    # If increment cannot be parsed, return identity function
    return lambda count: count


def simulate_labs(labs):
    from collections import deque
    from math import gcd
    from functools import reduce

    # Initialize labs' queues and inspections
    for lab in labs.values():
        lab['queue'] = deque(lab['cell_counts'])
        lab['inspections'] = 0

    # Compute the LCM of all the divisors to use for modulus
    divisors = [lab['divisor'] for lab in labs.values()]
    lcm_value = reduce(lambda a, b: a * b // gcd(a, b), divisors)

    days = 10000
    intervals = range(1000, days + 1, 1000)
    inspection_records = {}
    for day in range(1, days + 1):
        # Prepare a temporary storage for new dishes to be added to labs
        new_dishes = {lab_number: deque() for lab_number in labs}

        for lab_number, lab in labs.items():
            while lab['queue']:
                count = lab['queue'].popleft()
                lab['inspections'] += 1

                # Apply increment operation
                new_count = lab['increment_op'](count)
                new_count %= lcm_value  # Apply modulus to prevent counts from growing too large

                # Apply condition
                if new_count % lab['divisor'] == 0:
                    next_lab_number = lab['true_lab']
                else:
                    next_lab_number = lab['false_lab']

                # Pass the dish to the next lab
                new_dishes[next_lab_number].append(new_count)

        # Update labs' queues with new dishes
        for lab_number, lab in labs.items():
            lab['queue'].extend(new_dishes.get(lab_number, []))

        # Record inspections at intervals
        if day in intervals:
            inspection_records[str(day)] = [lab['inspections'] for lab_number, lab in sorted(labs.items())]

    return inspection_records