import json
import pandas as pd
from typing import Dict, Any, Optional
from .utils import calculate_checksum

def infer_schema(df: pd.DataFrame) -> Dict[str, str]:
    """Infer PFX schema from pandas DataFrame"""
    schema = {}

    for column, dtype in df.dtypes.items():
        if pd.api.types.is_integer_dtype(dtype):
            if dtype == 'int32':
                schema[column] = 'int32'
            else:
                schema[column] = 'int64'
        elif pd.api.types.is_float_dtype(dtype):
            schema[column] = 'float64'
        elif pd.api.types.is_bool_dtype(dtype):
            schema[column] = 'bool'
        elif pd.api.types.is_string_dtype(dtype):
            schema[column] = 'str'
        elif pd.api.types.is_list_like(dtype):
            # Try to infer list type from first non-null value
            non_null = df[column].dropna()
            if len(non_null) > 0:
                first = non_null.iloc[0]
                if isinstance(first, list) and len(first) > 0:
                    sample_type = type(first[0]).__name__
                    if sample_type == 'str':
                        schema[column] = 'list[str]'
                    elif sample_type == 'int':
                        schema[column] = 'list[int32]'
                    elif sample_type == 'float':
                        schema[column] = 'list[float64]'
                    else:
                        schema[column] = f'list[{sample_type}]'
                else:
                    schema[column] = 'list[str]'  # default
            else:
                schema[column] = 'list[str]'  # default
        elif pd.api.types.is_dict_like(dtype):
            schema[column] = 'dict'
        else:
            schema[column] = 'str'  # fallback

    return schema

def format_value(value: Any) -> str:
    """Format a value for PFX output"""
    if value is None:
        return 'null'
    elif isinstance(value, bool):
        return 'true' if value else 'false'
    elif isinstance(value, (int, float)):
        return str(value)
    elif isinstance(value, str):
        return f'"{value.replace("\\", "\\\\").replace("\"", "\\\"")}"'
    elif isinstance(value, list):
        return json.dumps(value)
    elif isinstance(value, dict):
        return json.dumps(value)
    else:
        return str(value)

def to_pfx(df: pd.DataFrame, file_path: str, metadata: Optional[Dict] = None,
           compression: str = 'none', description: str = ""):
    """
    Write a pandas DataFrame to a PandaFlex file

    Args:
        df: pandas DataFrame to save
        file_path: Path to save the PFX file
        metadata: Optional metadata dictionary
        compression: Compression type (currently only 'none' supported)
        description: Description of the dataset
    """
    # Infer schema
    schema = infer_schema(df)

    # Prepare header
    header = [
        "# PandaFlex v1.0",
        f'# Metadata: {json.dumps({"description": description} if description else {})}'
    ]

    # Prepare schema section
    schema_lines = ["[Schema]"] + [f"{col}: {dtype}" for col, dtype in schema.items()]

    # Prepare data section
    data_lines = ["[Data]"]
    for _, row in df.iterrows():
        row_values = [format_value(row[col]) for col in df.columns]
        data_lines.append(", ".join(row_values))

    # Prepare footer
    footer = [
        "[Footer]",
        f"row_count: {len(df)}",
        f"compression: {compression}",
        f"checksum: \"{calculate_checksum('\n'.join(data_lines))}\""
    ]

    # Combine all sections
    content = "\n".join(header + [""] + schema_lines + [""] + data_lines + [""] + footer)

    # Write to file
    with open(file_path, 'w') as f:
        f.write(content)

    # Add metadata if provided
    if metadata is not None:
        with open(file_path, 'a') as f:
            f.write(f"\n\n[Metadata]\n{json.dumps(metadata, indent=2)}")