"""Core schema inference engine."""

import json
import csv
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, Type
from enum import Enum

from pydantic import BaseModel, create_model, Field
import pandas as pd


class DataType(str, Enum):
    """Supported data types for schema inference."""
    STRING = "str"
    INTEGER = "int"
    FLOAT = "float"
    BOOLEAN = "bool"
    LIST = "list"
    OBJECT = "dict"
    NULL = "NoneType"


class SchemaInferencer:
    """Infer Pydantic schemas from JSON and CSV files."""

    def __init__(self, strict: bool = False):
        """
        Initialize the schema inferencer.
        
        Args:
            strict: If True, disable type coercion and be stricter about types.
        """
        self.strict = strict
        self.type_samples: Dict[str, Dict[str, int]] = {}

    def infer_from_json(
        self,
        path: Union[str, Path],
        model_name: str = "GeneratedModel",
        **kwargs
    ) -> Type[BaseModel]:
        """
        Infer a Pydantic model from a JSON file.
        
        Args:
            path: Path to the JSON file
            model_name: Name for the generated model
            **kwargs: Additional arguments passed to create_model
            
        Returns:
            A Pydantic model class
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        with open(path, 'r') as f:
            data = json.load(f)

        if isinstance(data, list):
            # Handle list of objects - infer from first item
            if not data:
                raise ValueError("Cannot infer schema from empty list")
            schema_dict = self._infer_schema_from_dict(data[0])
        elif isinstance(data, dict):
            schema_dict = self._infer_schema_from_dict(data)
        else:
            raise ValueError(f"Unsupported JSON root type: {type(data)}")

        return self._create_pydantic_model(schema_dict, model_name, **kwargs)

    def infer_from_csv(
        self,
        path: Union[str, Path],
        model_name: str = "GeneratedModel",
        sample_rows: Optional[int] = None,
        **kwargs
    ) -> Type[BaseModel]:
        """
        Infer a Pydantic model from a CSV file.
        
        Args:
            path: Path to the CSV file
            model_name: Name for the generated model
            sample_rows: Number of rows to sample for inference (None = all)
            **kwargs: Additional arguments passed to create_model
            
        Returns:
            A Pydantic model class
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        df = pd.read_csv(path, nrows=sample_rows)
        
        if df.empty:
            raise ValueError("Cannot infer schema from empty CSV")

        schema_dict = self._infer_schema_from_dataframe(df)
        return self._create_pydantic_model(schema_dict, model_name, **kwargs)

    def _infer_schema_from_dict(self, obj: Dict[str, Any]) -> Dict[str, tuple]:
        """Infer schema from a dictionary."""
        schema = {}
        for key, value in obj.items():
            python_type, default = self._infer_type(value)
            schema[key] = (Optional[python_type], default)
        return schema

    def _infer_schema_from_dataframe(self, df: pd.DataFrame) -> Dict[str, tuple]:
        """Infer schema from a pandas DataFrame."""
        schema = {}
        for column in df.columns:
            python_type, default = self._infer_type_from_series(df[column])
            schema[column] = (Optional[python_type], default)
        return schema

    def _infer_type(self, value: Any) -> tuple:
        """
        Infer Python type from a value.
        
        Returns:
            Tuple of (python_type, default_value)
        """
        if value is None:
            return str, Field(default=None)

        elif isinstance(value, bool):
            # Check bool before int since bool is subclass of int
            return bool, Field(default=None)

        elif isinstance(value, int):
            return int, Field(default=None)

        elif isinstance(value, float):
            return float, Field(default=None)

        elif isinstance(value, str):
            return str, Field(default=None)

        elif isinstance(value, list):
            if not value:
                return list, Field(default=None)
            # Infer list element type from first element
            element_type, _ = self._infer_type(value[0])
            return List[element_type], Field(default=None)

        elif isinstance(value, dict):
            # For nested dicts, return generic dict or nested model
            return dict, Field(default=None)

        else:
            # Default to string for unknown types
            return str, Field(default=None)

    def _infer_type_from_series(self, series: pd.Series) -> tuple:
        """Infer Python type from a pandas Series."""
        # Remove null values for type inference
        non_null = series.dropna()
        
        if non_null.empty:
            return str, Field(default=None)

        # Check for boolean (before int, as bool is subclass of int)
        if series.dtype == "bool" or set(non_null.unique()) <= {0, 1, True, False}:
            return bool, Field(default=None)

        # Check for integer
        if pd.api.types.is_integer_dtype(series):
            return int, Field(default=None)

        # Check for float
        if pd.api.types.is_float_dtype(series):
            return float, Field(default=None)

        # Check for string/object
        if pd.api.types.is_string_dtype(series) or pd.api.types.is_object_dtype(series):
            return str, Field(default=None)

        # Default to string
        return str, Field(default=None)

    def _create_pydantic_model(
        self,
        schema_dict: Dict[str, tuple],
        model_name: str,
        **kwargs
    ) -> Type[BaseModel]:
        """Create a Pydantic model from schema dictionary."""
        # Add docstring if not provided
        if "doc" not in kwargs:
            kwargs["__doc__"] = f"Auto-generated model from {model_name}"

        return create_model(model_name, **schema_dict, **kwargs)

    def to_code(self, model: Type[BaseModel], class_name: str = "Model") -> str:
        """
        Generate Python code for the model.
        
        Args:
            model: The Pydantic model class
            class_name: Name to use in generated code
            
        Returns:
            Python code string
        """
        lines = [
            "from pydantic import BaseModel",
            "from typing import Optional, List",
            "",
            f"class {class_name}(BaseModel):",
        ]

        # Get model fields
        for field_name, field_info in model.model_fields.items():
            annotation = self._type_to_string(field_info.annotation)
            # Only add default if it's actually set
            if field_info.default is not None:
                default_value = repr(field_info.default)
                lines.append(f"    {field_name}: {annotation} = {default_value}")
            else:
                lines.append(f"    {field_name}: {annotation}")

        return "\n".join(lines)

    def _type_to_string(self, annotation: Any) -> str:
        """Convert a type annotation to a string representation."""
        if annotation is None or annotation is type(None):
            return "None"
        
        if hasattr(annotation, "__origin__"):
            # Handle generic types like Optional, List, etc.
            origin = annotation.__origin__
            if origin is Union:
                args = annotation.__args__
                # Filter out NoneType and build the string properly
                non_none_args = [arg for arg in args if arg is not type(None)]
                none_args = [arg for arg in args if arg is type(None)]
                
                if not non_none_args:
                    return "None"
                
                if non_none_args and none_args:
                    # Optional type: int | None
                    if len(non_none_args) == 1:
                        return f"Optional[{self._type_to_string(non_none_args[0])}]"
                    else:
                        # Multiple non-none types, use union
                        return "Optional[" + " | ".join(self._type_to_string(arg) for arg in non_none_args) + "]"
                else:
                    # All non-none types
                    return " | ".join(self._type_to_string(arg) for arg in non_none_args)
            elif origin is list:
                inner_type = self._type_to_string(annotation.__args__[0])
                return f"List[{inner_type}]"
            elif origin is dict:
                return "dict"
        
        if hasattr(annotation, "__name__"):
            return annotation.__name__
        
        return str(annotation)
