"""
ROTATING STDOUT - Dual logs: 5000 lines (main) + 2500 lines (compact)
"""
import sys
import os
from datetime import datetime

MAX_LINES = 5000
MAX_LINES_COMPACT = 2500
LOG_FILE = os.environ.get('TRAINING_LOG', '/root/training.log')
LOG_FILE_COMPACT = LOG_FILE.replace('.log', '_compact.log')
ARCHIVE_DIR = LOG_FILE.replace('.log', '_archives')

class RotatingStdout:
    def __init__(self, max_lines=MAX_LINES, max_compact=MAX_LINES_COMPACT):
        self.max_lines = max_lines
        self.max_compact = max_compact
        self.line_count = 0
        self.line_count_compact = 0
        self.original_stdout = sys.stdout
        self.rotating = False
        os.makedirs(ARCHIVE_DIR, exist_ok=True)
        self.log_file = open(LOG_FILE, 'a', buffering=1)
        self.log_compact = open(LOG_FILE_COMPACT, 'a', buffering=1)
    
    def _rotate_main(self):
        self.log_file.close()
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        archive_path = os.path.join(ARCHIVE_DIR, f'archive_{timestamp}.log')
        with open(LOG_FILE, 'r') as f:
            lines = f.readlines()
        keep = min(self.max_lines // 2, len(lines))
        with open(archive_path, 'w') as f:
            f.writelines(lines[:-keep] if keep else lines)
        with open(LOG_FILE, 'w') as f:
            f.writelines(lines[-keep:] if keep else [])
        self.line_count = keep
        self.log_file = open(LOG_FILE, 'a', buffering=1)
    
    def _rotate_compact(self):
        self.log_compact.close()
        with open(LOG_FILE_COMPACT, 'r') as f:
            lines = f.readlines()
        keep = min(self.max_compact // 2, len(lines))
        with open(LOG_FILE_COMPACT, 'w') as f:
            f.writelines(lines[-keep:] if keep else [])
        self.line_count_compact = keep
        self.log_compact = open(LOG_FILE_COMPACT, 'a', buffering=1)
    
    def _rotate_if_needed(self):
        if self.rotating:
            return
        self.rotating = True
        if self.line_count > self.max_lines:
            self._rotate_main()
        if self.line_count_compact > self.max_compact:
            self._rotate_compact()
        self.rotating = False
    
    def write(self, text):
        self.original_stdout.write(text)
        if not self.rotating:
            self.log_file.write(text)
            self.log_compact.write(text)
            newlines = text.count('\n')
            self.line_count += newlines
            self.line_count_compact += newlines
            self._rotate_if_needed()
    
    def flush(self):
        self.original_stdout.flush()
        if self.log_file:
            self.log_file.flush()
        if self.log_compact:
            self.log_compact.flush()

def install_rotating_log():
    rotating = RotatingStdout()
    sys.stdout = rotating
    sys.stderr = rotating
    print(f"[rotating_log] Active. Main={MAX_LINES}, Compact={MAX_LINES_COMPACT} lines.")