import pandas as pd
import numpy as np
import io
import contextlib
import traceback
import json
import logging
import math
import re
from datetime import datetime, date, timedelta
from typing import Any, Dict, Optional, Union

# --- Robust Library Loading ---
# We try to import common analysis libraries. 
# If they are missing, the service still runs, just without those specifics.
try:
    import scipy
    import scipy.stats as stats
except ImportError:
    scipy = None
    stats = None

try:
    import sklearn
    from sklearn.linear_model import LinearRegression, LogisticRegression
    from sklearn.model_selection import train_test_split
    from sklearn import metrics
except ImportError:
    sklearn = None

try:
    import statsmodels.api as sm
    import statsmodels.formula.api as smf
except ImportError:
    sm = None
    smf = None

# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("CSV_Analysis_Executor")

def robust_json_serializer(obj: Any) -> Any:
    """
    Universal Serializer.
    1. Handles DataFrames/Series (returns FULL data).
    2. Handles NumPy types (int/float/arrays).
    3. Handles Dates.
    4. Handles Unknown Objects (Models, Classes) -> returns String Repr.
    """
    # 1. Pandas Types
    if isinstance(obj, pd.DataFrame):
        # UNCONSTRAINED: Return the full dataset as list of dicts
        return obj.to_dict(orient='records')
    elif isinstance(obj, pd.Series):
        return obj.to_dict()
    elif isinstance(obj, pd.Index):
        return obj.tolist()
    
    # 2. NumPy Types
    elif isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        if np.isnan(obj) or np.isinf(obj):
            return None
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, np.bool_):
        return bool(obj)
    
    # 3. Standard Python Dates
    elif isinstance(obj, (datetime, date)):
        return obj.isoformat()
    elif isinstance(obj, timedelta):
        return str(obj)
    
    # 4. Fallback for Complex Objects (e.g., sklearn model, statsmodels result)
    # Instead of crashing, we return the string representation
    if hasattr(obj, '__dict__'):
        return str(obj)
    
    return str(obj)

class CsvAnalysisExecutor:
    """
    A 'Natural' Executor. 
    It mimics a local Jupyter notebook environment with common libraries pre-loaded.
    """
    def __init__(self, df: pd.DataFrame):
        # NATURAL MODE: We do NOT sanitize column names. 
        # We rely on the code generator to handle keys like df['User ID'] correctly.
        self.df = df
        self.exec_locals = {}

    def execute(self, code: str) -> Dict[str, Any]:
        """
        Executes code with a rich data science context.
        """
        output_buffer = io.StringIO()
        error_result = None
        success = False

        # --- Rich Environment Injection ---
        exec_globals = {
            # Core
            'pd': pd,
            'np': np,
            'df': self.df,
            'math': math,
            're': re,
            'json': json,
            'datetime': datetime,
            'timedelta': timedelta,
            
            # Statistics & ML (if available)
            'scipy': scipy,
            'stats': stats,
            'sklearn': sklearn,
            'sm': sm,
            'smf': smf,
            
            # Common shortcuts (makes generated code more natural)
            'LinearRegression': LinearRegression if sklearn else None,
            'train_test_split': train_test_split if sklearn else None,
        }

        try:
            # Capture standard output (print statements)
            with contextlib.redirect_stdout(output_buffer):
                exec(code, exec_globals, self.exec_locals)
            
            success = True

        except Exception:
            # Clean traceback for the user
            error_result = traceback.format_exc()
            logger.error(f"Analysis Execution Error:\n{error_result}")

        # --- Extract "Natural" Results ---
        # We capture everything the user defined, excluding imports and modules
        final_vars = {}
        
        # List of system modules to ignore in output
        ignored_types = (type(pd), type(np), type(math), type(json))
        
        for k, v in self.exec_locals.items():
            # Skip hidden vars
            if k.startswith('_'): continue
            
            # Skip the input dataframe ref (unless they made a copy)
            if k == 'df': continue
            
            # Skip modules (e.g., if user did 'import random', don't return the random module)
            if isinstance(v, ignored_types) or hasattr(v, '__name__') and 'module' in str(type(v)):
                continue

            # Capture the variable
            final_vars[k] = v

        return {
            "success": success,
            "output_log": output_buffer.getvalue(),
            "error": error_result,
            "data_vars": final_vars
        }

def execute_analysis_logic(code: str, csv_url: str) -> Dict[str, Any]:
    """
    Entry point. Loads CSV (handling errors) and executes logic.
    """
    try:
        # 1. Load Data
        logger.info(f"Loading CSV from: {csv_url}")
        try:
            # Use 'on_bad_lines' to be robust against messy CSVs
            df = pd.read_csv(csv_url, on_bad_lines='skip')
        except Exception as e:
            return {
                "success": False,
                "error": f"Failed to load CSV: {str(e)}",
                "output_log": "",
                "results": {}
            }
        
        # 2. Execute
        executor = CsvAnalysisExecutor(df)
        result = executor.execute(code)
        
        # 3. Serialize (Unconstrained)
        clean_vars = {}
        if result["data_vars"]:
            try:
                # Force strictly valid JSON
                serialized_str = json.dumps(result["data_vars"], default=robust_json_serializer)
                clean_vars = json.loads(serialized_str)
            except Exception as e:
                # Fallback mechanism
                logger.error(f"Serialization Warning: {e}")
                clean_vars = {"serialization_error": str(e), "raw_str_dump": str(result["data_vars"])}

        return {
            "success": result["success"],
            "error": result["error"],
            "output_log": result["output_log"],
            "results": clean_vars
        }

    except Exception as e:
        err_msg = f"System Error: {str(e)}\n{traceback.format_exc()}"
        return {
            "success": False,
            "error": err_msg,
            "output_log": "",
            "results": {}
        }