"""
Schema Exporter - Generate JSON schema from cleaned DataFrames.
"""

import json
import pandas as pd
import numpy as np
from typing import Dict, Any, List, Optional
from datetime import datetime


class SchemaExporter:
    """Exports DataFrame schema to JSON format for ML/analytics pipelines."""
    
    def __init__(self):
        self.schema = {}
    
    def generate_schema(
        self,
        df: pd.DataFrame,
        dataset_name: str = "cleaned_dataset"
    ) -> str:
        """
        Generate JSON schema from DataFrame.
        
        Args:
            df: Cleaned DataFrame
            dataset_name: Name of the dataset
            
        Returns:
            JSON schema as string
        """
        schema = {
            "$schema": "http://json-schema.org/draft-07/schema#",
            "title": dataset_name,
            "type": "object",
            "metadata": {
                "generated_at": datetime.now().isoformat(),
                "generator": "DataShaper AI v1.0",
                "row_count": len(df),
                "column_count": len(df.columns)
            },
            "properties": {},
            "required": []
        }
        
        # Analyze each column
        for col in df.columns:
            col_schema = self._analyze_column(df[col], col)
            schema["properties"][col] = col_schema
            
            # Add to required if not nullable
            if not col_schema.get("nullable", True):
                schema["required"].append(col)
        
        return json.dumps(schema, indent=2)
    
    def _analyze_column(self, series: pd.Series, col_name: str) -> Dict[str, Any]:
        """Analyze a single column and generate its schema."""
        col_schema = {
            "description": f"Column: {col_name}",
            "nullable": series.isna().any()
        }
        
        # Determine data type
        dtype = series.dtype
        
        if pd.api.types.is_integer_dtype(dtype):
            col_schema.update(self._analyze_integer(series))
        elif pd.api.types.is_float_dtype(dtype):
            col_schema.update(self._analyze_float(series))
        elif pd.api.types.is_bool_dtype(dtype):
            col_schema.update(self._analyze_boolean(series))
        elif pd.api.types.is_datetime64_any_dtype(dtype):
            col_schema.update(self._analyze_datetime(series))
        else:
            col_schema.update(self._analyze_string(series))
        
        return col_schema
    
    def _analyze_integer(self, series: pd.Series) -> Dict[str, Any]:
        """Analyze integer column."""
        clean_series = series.dropna()
        
        return {
            "type": "integer",
            "minimum": int(clean_series.min()) if len(clean_series) > 0 else None,
            "maximum": int(clean_series.max()) if len(clean_series) > 0 else None,
            "statistics": {
                "mean": float(clean_series.mean()) if len(clean_series) > 0 else None,
                "median": float(clean_series.median()) if len(clean_series) > 0 else None,
                "std": float(clean_series.std()) if len(clean_series) > 0 else None
            }
        }
    
    def _analyze_float(self, series: pd.Series) -> Dict[str, Any]:
        """Analyze float column."""
        clean_series = series.dropna()
        
        return {
            "type": "number",
            "minimum": float(clean_series.min()) if len(clean_series) > 0 else None,
            "maximum": float(clean_series.max()) if len(clean_series) > 0 else None,
            "statistics": {
                "mean": float(clean_series.mean()) if len(clean_series) > 0 else None,
                "median": float(clean_series.median()) if len(clean_series) > 0 else None,
                "std": float(clean_series.std()) if len(clean_series) > 0 else None
            }
        }
    
    def _analyze_boolean(self, series: pd.Series) -> Dict[str, Any]:
        """Analyze boolean column."""
        clean_series = series.dropna()
        
        return {
            "type": "boolean",
            "statistics": {
                "true_count": int(clean_series.sum()) if len(clean_series) > 0 else 0,
                "false_count": int((~clean_series).sum()) if len(clean_series) > 0 else 0,
                "true_percentage": float(clean_series.mean() * 100) if len(clean_series) > 0 else 0
            }
        }
    
    def _analyze_datetime(self, series: pd.Series) -> Dict[str, Any]:
        """Analyze datetime column."""
        clean_series = series.dropna()
        
        schema = {
            "type": "string",
            "format": "date-time"
        }
        
        if len(clean_series) > 0:
            schema["minimum"] = clean_series.min().isoformat()
            schema["maximum"] = clean_series.max().isoformat()
        
        return schema
    
    def _analyze_string(self, series: pd.Series) -> Dict[str, Any]:
        """Analyze string column."""
        clean_series = series.dropna()
        
        schema = {
            "type": "string"
        }
        
        if len(clean_series) > 0:
            # Calculate string lengths
            lengths = clean_series.astype(str).str.len()
            schema["minLength"] = int(lengths.min())
            schema["maxLength"] = int(lengths.max())
            
            # Get cardinality
            unique_count = clean_series.nunique()
            total_count = len(clean_series)
            
            schema["statistics"] = {
                "unique_count": int(unique_count),
                "cardinality_ratio": float(unique_count / total_count)
            }
            
            # If low cardinality, include enum values
            if unique_count <= 20:
                schema["enum"] = sorted(clean_series.unique().tolist())
            
            # If very low cardinality, likely categorical
            if unique_count / total_count < 0.1:
                schema["suggested_type"] = "categorical"
                
            # Sample values
            schema["examples"] = clean_series.head(5).tolist()
        
        return schema
    
    def generate_schema_from_multiple(
        self,
        dataframes: Dict[str, pd.DataFrame]
    ) -> str:
        """
        Generate a combined schema from multiple DataFrames.
        
        Args:
            dataframes: Dictionary of name -> DataFrame
            
        Returns:
            JSON schema covering all DataFrames
        """
        combined_schema = {
            "$schema": "http://json-schema.org/draft-07/schema#",
            "title": "DataShaper AI - Multi-File Schema",
            "type": "object",
            "metadata": {
                "generated_at": datetime.now().isoformat(),
                "generator": "DataShaper AI v1.0",
                "file_count": len(dataframes)
            },
            "datasets": {}
        }
        
        for name, df in dataframes.items():
            dataset_schema = json.loads(self.generate_schema(df, name))
            combined_schema["datasets"][name] = dataset_schema
        
        return json.dumps(combined_schema, indent=2)
