#!/usr/bin/env python3
"""
Log file JSON parser - Parse log files in statement_logs directory and write to Delta Lake table
"""

import os
import json
import re
import pandas as pd
from pathlib import Path
from datetime import datetime, timezone
from deltalake import DeltaTable, write_deltalake
import argparse
import pyarrow as pa

class StatementLogParser:
    def __init__(self, input_dir="./statement_logs", output_dir="./statement_delta"):
        """
        Initialize log parser

        Args:
            input_dir: Input log file directory
            output_dir: Output Delta Lake table directory
        """
        self.input_dir = Path(input_dir)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)

        # Regular expression for JSON extraction
        self.json_pattern = re.compile(r'\{.*\}')

        # Statistics
        self.stats = {
            'total_files': 0,
            'processed_files': 0,
            'total_lines': 0,
            'parsed_lines': 0,
            'failed_lines': 0
        }

    def extract_json_from_line(self, line):
        """
        Extract JSON part from log line

        Args:
            line: Log line

        Returns:
            dict: Extracted JSON data, returns None if extraction fails
        """
        try:
            # Use regular expression to extract JSON part (content surrounded by curly braces)
            json_match = self.json_pattern.search(line)
            if not json_match:
                return None

            json_str = json_match.group()
            return json.loads(json_str)
        except (json.JSONDecodeError, ValueError):
            return None

    def parse_log_file(self, file_path):
        """
        Parse a single log file

        Args:
            file_path: Log file path

        Returns:
            list: List of successfully parsed JSON data
        """
        parsed_data = []

        try:
            print(f"Processing file: {file_path}")
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    self.stats['total_lines'] += 1

                    # Extract and parse JSON
                    json_data = self.extract_json_from_line(line.strip())
                    if json_data:
                        parsed_data.append(json_data)
                        self.stats['parsed_lines'] += 1
                    else:
                        self.stats['failed_lines'] += 1

        except Exception as e:
            print(f"Error processing file {file_path}: {e}")

        return parsed_data

    def process_all_files(self, batch_size=1000):
        """
        Process all log files in directory (batch processing to avoid memory overflow)

        Args:
            batch_size: Number of files to process per batch before writing to Delta Lake

        Returns:
            int: Total number of records processed
        """
        if not self.input_dir.exists():
            print(f"Input directory does not exist: {self.input_dir}")
            return 0

        # Find all .log files
        log_files = list(self.input_dir.glob("*.log"))
        self.stats['total_files'] = len(log_files)

        print(f"Found {len(log_files)} log files, starting batch processing...")

        total_records = 0
        current_batch = []

        for i, log_file in enumerate(log_files):
            file_data = self.parse_log_file(log_file)
            if file_data:
                current_batch.extend(file_data)
                self.stats['processed_files'] += 1
            
            if current_batch:
                self._write_batch_to_delta(current_batch, i // batch_size + 1)
                total_records += len(current_batch)
                current_batch = []  # Clear current batch

        print(f"Total records written: {total_records}")
        return total_records

    def _write_batch_to_delta(self, batch_data, batch_num):
        """
        Write a batch of data to Delta Lake table

        Args:
            batch_data: Batch data
            batch_num: Batch number for progress display
        """
        try:
            # Preprocess data, handle None values and empty structures
            processed_data = []
            for item in batch_data:
                if item is None:
                    continue
                    
                # Handle None values in nested structures
                processed_item = {}
                for key, value in item.items():
                    if value is None:
                        processed_item[key] = None
                    elif isinstance(value, dict):
                        # Recursively process nested dictionaries
                        processed_item[key] = {k: v for k, v in value.items() if v is not None} or None
                    elif isinstance(value, list):
                        # Handle None values in lists
                        processed_item[key] = [v for v in value if v is not None] or None
                    else:
                        processed_item[key] = value
                processed_data.append(processed_item)
            
            if not processed_data:
                print(f"Warning: Batch {batch_num} has no valid data, skipping")
                return

            # Convert to DataFrame
            df = pd.DataFrame(processed_data)

            # Add processing timestamp and batch number
            df['processed_at'] = datetime.now(timezone.utc)
            df['batch_num'] = batch_num

            # Convert data types and handle null values
            for column in df.columns:
                if column in ['processed_at', 'batch_num']:
                    continue
                    
                # Handle list types (e.g., index_names column)
                if any(isinstance(x, (list, dict)) for x in df[column].dropna()):
                    # Convert lists/dictionaries to JSON strings
                    df[column] = df[column].apply(
                        lambda x: json.dumps(x, ensure_ascii=False) 
                        if x is not None and isinstance(x, (list, dict)) 
                        else x
                    )
                # Try to convert to numeric types
                elif not df[column].empty:
                    non_null = df[column].dropna()
                    if not non_null.empty:
                        try:
                            # First try to convert to integer
                            if all(isinstance(x, (int, bool)) or 
                                 (isinstance(x, str) and x.replace('-', '').isdigit()) 
                                 for x in non_null if x not in ['None', '']):
                                df[column] = pd.to_numeric(df[column])
                            # Then try to convert to float
                            elif all(isinstance(x, (int, float)) or 
                                   (isinstance(x, str) and self._is_float(x)) 
                                   for x in non_null if x not in ['None', '']):
                                df[column] = pd.to_numeric(df[column])
                        except (ValueError, TypeError):
                            # Conversion failed, keep as is
                            pass
                
                # Handle null values
                df[column] = df[column].replace(['None', ''], None)
                
                # Limit string length
                if not df[column].empty and df[column].dtype == 'object':
                    def truncate_string(x):
                        if x is not None:
                            s = str(x)
                            if len(s) > 100:
                                print(f"Truncating long string in column '{column}': {s[:100]}... (length: {len(s)} -> 100)")
                                return s[:100]
                            return x
                        return x
                    
                    df[column] = df[column].apply(truncate_string)

            # Generate Delta Lake schema
            delta_schema = self._get_delta_schema(df)

            # Save to Delta Lake
            table_path = self.output_dir / "statement_logs"
            table_path_str = str(table_path.absolute())

            print(f"Writing batch {batch_num} with {len(df)} records...")
            #print(f"Using schema: {[field.name + ':' + str(field.type) for field in delta_schema]}")

            # Ensure output directory exists
            table_path.parent.mkdir(parents=True, exist_ok=True)

            # Try to write to Delta Lake using PyArrow Table
            try:
                # Convert to PyArrow Table with predefined schema
                pa_table = pa.Table.from_pandas(df, schema=delta_schema, preserve_index=False)

                write_deltalake(
                    table_path_str,
                    pa_table,
                    mode="append",  # Append mode
                )
                print(f"Batch {batch_num} write completed")

            except Exception as e:
                print(f"Delta Lake write failed: {str(e)}")

        except ImportError as e:
            if "pyarrow" in str(e):
                print(f"⚠️  Missing pyarrow dependency")
            else:
                raise
        except Exception as e:
            print(f"Error writing batch {batch_num}: {e}")
            raise

    def _infer_column_type(self, series):
        """
        Infer column data type
        
        Args:
            series: pandas Series to infer type for
            
        Returns:
            pyarrow.DataType: Inferred data type
        """
        # Skip null values
        non_null = series.dropna()
        if len(non_null) == 0:
            return pa.string()  # Default to string type
            
        # Check if contains list or dict types
        if any(isinstance(x, (list, dict)) for x in non_null):
            return pa.string()  # Store lists/dictionaries as JSON strings
            
        # Check if all values are integers
        try:
            if all(isinstance(x, (int, bool)) or 
                  (isinstance(x, str) and x.replace('-', '').isdigit()) 
                  for x in non_null if x not in ['None', '']):
                return pa.int64()
                
            # Check if all values are numbers (integers or floats)
            if all(isinstance(x, (int, float)) or 
                  (isinstance(x, str) and self._is_float(x)) 
                  for x in non_null if x not in ['None', '']):
                return pa.float64()
                
        except (ValueError, TypeError):
            pass
            
        # Default to string type
        return pa.string()
        
    def _is_float(self, value):
        """Check if string can be converted to float"""
        try:
            float(value)
            return True
        except (ValueError, TypeError):
            return False

    def _get_delta_schema(self, df):
        """
        Generate Delta Lake schema from DataFrame

        Args:
            df: Input DataFrame

        Returns:
            pa.Schema: PyArrow schema
        """
        schema_fields = []

        for column in df.columns:
            if column in ['processed_at', 'batch_num']:
                # Timestamp columns
                schema_fields.append(pa.field(column, pa.timestamp('us', tz='UTC'), nullable=True))
            else:
                # Infer column type
                col_type = self._infer_column_type(df[column])
                schema_fields.append(pa.field(column, col_type, nullable=True))

        return pa.schema(schema_fields)

    def _is_empty_struct(self, value):
        """
        Check if value is an empty struct (will cause Parquet write errors)
        Recursively check nested structures

        Args:
            value: Value to check

        Returns:
            bool: Whether it's an empty struct or struct containing only empty values
        """
        if value is None:
            return False

        # Check if it's an empty dictionary
        if isinstance(value, dict):
            if not value:  # Empty dictionary
                return True
            # Recursively check dictionary values
            return all(self._is_empty_struct(v) for v in value.values())

        # Check if it's an empty list
        if isinstance(value, list):
            if not value:  # Empty list
                return True
            # Recursively check list elements
            return all(self._is_empty_struct(item) for item in value)

        # Check if it's an empty string
        if isinstance(value, str) and not value.strip():
            return True

        # Other types (numbers, booleans, etc.) are not considered empty structs
        return False


    def print_stats(self):
        """
        Print processing statistics
        """
        print("\n" + "="*50)
        print("Processing Statistics:")
        print(f"Total files: {self.stats['total_files']}")
        print(f"Processed files: {self.stats['processed_files']}")
        print(f"Total lines: {self.stats['total_lines']}")
        print(f"Successfully parsed lines: {self.stats['parsed_lines']}")
        print(f"Failed lines: {self.stats['failed_lines']}")
        print(f"Parse success rate: {self.stats['parsed_lines']/max(self.stats['total_lines'], 1)*100:.2f}%")
        print("="*50)

def main():
    parser = argparse.ArgumentParser(description='Parse statement log files and write to Delta Lake table')
    parser.add_argument('--input-dir', default='./statement_logs', help='Input log file directory')
    parser.add_argument('--output-dir', default='./statement_delta', help='Output Delta Lake table directory')
    parser.add_argument('--table-name', default='statement_logs', help='Delta Lake table name')
    parser.add_argument('--batch-size', type=int, default=1000, help='Number of files to process per batch before writing to Delta Lake')

    args = parser.parse_args()

    # Create parser and process files
    log_parser = StatementLogParser(args.input_dir, args.output_dir)

    print("Starting to parse log files...")
    total_records = log_parser.process_all_files(args.batch_size)

    # Print statistics
    log_parser.print_stats()

    if total_records > 0:
        print(f"\n✅ Successfully processed {total_records} log records to Delta Lake table")
        print(f"📁 Results saved in: {log_parser.output_dir / args.table_name}")
    else:
        print("\n❌ No valid log data found")

if __name__ == "__main__":
    main()