#!/usr/bin/env python3
"""Script to shift dates in MQL5 CSV file by specified number of weeks"""

import os
import sys
import glob
from datetime import datetime, timedelta

# Determine input file
if len(sys.argv) > 1:
    input_file = sys.argv[1]
else:
    # Search for CSV files in current directory (excluding already processed)
    csv_files = [f for f in glob.glob("*.csv") if not f.endswith("_shifted.csv")]
    if len(csv_files) == 0:
        print("Error: No CSV files found in current directory")
        sys.exit(1)
    elif len(csv_files) == 1:
        input_file = csv_files[0]
        print(f"Found file: {input_file}")
    else:
        print("Multiple CSV files found. Select a file:")
        for i, f in enumerate(csv_files, 1):
            print(f"  [{i}] {f}")
        while True:
            try:
                selection = input(f"Enter number (1-{len(csv_files)}): ")
                selected_index = int(selection) - 1
                if 0 <= selected_index < len(csv_files):
                    input_file = csv_files[selected_index]
                    print(f"Selected: {input_file}")
                    break
                else:
                    print(f"Invalid selection. Enter a number from 1 to {len(csv_files)}")
            except ValueError:
                print("Please enter a number")
            except KeyboardInterrupt:
                print("\nCancelled by user")
                sys.exit(1)

# Check if file exists
if not os.path.exists(input_file):
    print(f"Error: File '{input_file}' not found")
    sys.exit(1)

# Request number of weeks to shift
default_weeks = 52  # Approximately a year
print(f"\nHow many weeks to shift dates? (default: {default_weeks})")
print("You can specify a negative value to shift backward (e.g., -52)")
while True:
    try:
        weeks_input = input(f"Enter number of weeks [{default_weeks}]: ").strip()
        if not weeks_input:
            weeks_shift = default_weeks
        else:
            weeks_shift = int(weeks_input)
            if weeks_shift == 0:
                print("Number of weeks cannot be zero")
                continue
        break
    except ValueError:
        print("Please enter an integer")
    except KeyboardInterrupt:
        print("\nCancelled by user")
        sys.exit(1)

direction = "forward" if weeks_shift > 0 else "backward"
print(f"Will shift {direction} by {abs(weeks_shift)} weeks ({abs(weeks_shift) * 7} days)")

# Form output file name
base_name = os.path.splitext(input_file)[0]
output_file = f"{base_name}_shifted.csv"

# Get file size for progress
file_size = os.path.getsize(input_file)
file_size_mb = file_size / (1024 * 1024)

print(f"Input file: {input_file}")
print(f"File size: {file_size_mb:.2f} MB")
print(f"Output file: {output_file}")
print("Processing...")

# Use buffering for large files
BUFFER_SIZE = 1024 * 1024  # 1MB buffer

with open(input_file, 'r', encoding='utf-8', buffering=BUFFER_SIZE) as f_in, \
     open(output_file, 'w', encoding='utf-8', buffering=BUFFER_SIZE) as f_out:
    
    # Header
    header = f_in.readline()
    f_out.write(header)
    
    # Determine file type by header
    header_upper = header.upper()
    if '<OPEN>' in header_upper and '<HIGH>' in header_upper:
        file_type = "Bars"
    elif '<BID>' in header_upper and '<ASK>' in header_upper:
        file_type = "Ticks"
    else:
        file_type = "Unknown"
    
    print(f"File type: {file_type}")
    
    processed = 0
    bytes_read = len(header.encode('utf-8'))
    
    # Process data
    for line in f_in:
        # Preserve original line ending
        line_ending = ''
        if line.endswith('\r\n'):
            line_ending = '\r\n'
            line = line[:-2]
        elif line.endswith('\n'):
            line_ending = '\n'
            line = line[:-1]
        elif line.endswith('\r'):
            line_ending = '\r'
            line = line[:-1]
        
        # Quick check and process date
        # Date format is the same for both types: YYYY.MM.DD
        if line and len(line) > 0 and line[0].isdigit():
            # Find position of first tabulator (separator between date and time)
            tab_pos = line.find('\t')
            if tab_pos >= 10:  # Minimum YYYY.MM.DD (10 characters)
                try:
                    # Parse date
                    date_str = line[:tab_pos]
                    dt = datetime.strptime(date_str, "%Y.%m.%d")
                    # Add weeks
                    dt_new = dt + timedelta(weeks=weeks_shift)
                    # Form new line with updated date
                    new_date_str = dt_new.strftime("%Y.%m.%d")
                    line = new_date_str + line[tab_pos:]
                except ValueError:
                    # If failed to parse date, skip line without changes
                    pass
        
        # Write line with original ending
        f_out.write(line + line_ending)
        
        processed += 1
        bytes_read += len((line + line_ending).encode('utf-8'))
        
        # Update progress every 10000 lines
        if processed % 10000 == 0:
            percent = min(100, int(bytes_read / file_size * 100))
            bar_length = 50
            filled = int(bar_length * bytes_read / file_size)
            bar = '█' * filled + '░' * (bar_length - filled)
            mb_processed = bytes_read / (1024 * 1024)
            sys.stdout.write(f'\r[{bar}] {percent}% ({mb_processed:.1f}/{file_size_mb:.1f} MB, {processed:,} lines)')
            sys.stdout.flush()

# Final progress
percent = 100
bar = '█' * 50
mb_processed = file_size / (1024 * 1024)
sys.stdout.write(f'\r[{bar}] {percent}% ({mb_processed:.1f}/{file_size_mb:.1f} MB, {processed:,} lines)')
sys.stdout.flush()

print(f"\nDone! Result saved to {output_file}")
print(f"Processed lines: {processed:,}")