"""
Test fixtures for FHEnom AI testing.
"""

import tempfile
import json
import os
from pathlib import Path
from typing import Dict, Any, List, Optional
import logging

logger = logging.getLogger(__name__)


def create_test_model(
    model_name: str = "test-model",
    model_type: str = "llama",
    vocab_size: int = 32000,
    hidden_size: int = 4096
) -> str:
    """
    Create a minimal test model directory.
    
    Args:
        model_name: Name for the model
        model_type: Model architecture type
        vocab_size: Vocabulary size
        hidden_size: Hidden layer size
    
    Returns:
        Path to test model directory
    
    Example:
        ```python
        from fhenomai.testing import create_test_model
        
        model_path = create_test_model()
        # Use for testing...
        ```
    """
    # Create temporary directory
    temp_dir = tempfile.mkdtemp(prefix=f"fhenomai_test_{model_name}_")
    
    # Create config.json
    config = {
        "architectures": [model_type],
        "model_type": model_type,
        "vocab_size": vocab_size,
        "hidden_size": hidden_size,
        "num_hidden_layers": 32,
        "num_attention_heads": 32,
        "intermediate_size": 11008,
        "hidden_act": "silu",
        "max_position_embeddings": 2048,
        "initializer_range": 0.02,
        "rms_norm_eps": 1e-6,
        "use_cache": True,
        "pad_token_id": 0,
        "bos_token_id": 1,
        "eos_token_id": 2,
        "tie_word_embeddings": False
    }
    
    with open(os.path.join(temp_dir, "config.json"), 'w') as f:
        json.dump(config, f, indent=2)
    
    # Create generation_config.json
    gen_config = {
        "bos_token_id": 1,
        "eos_token_id": 2,
        "max_length": 2048,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.9
    }
    
    with open(os.path.join(temp_dir, "generation_config.json"), 'w') as f:
        json.dump(gen_config, f, indent=2)
    
    # Create dummy model file
    model_file = os.path.join(temp_dir, "model.safetensors")
    with open(model_file, 'wb') as f:
        # Write some dummy binary data
        f.write(b'\x00' * 1024)  # 1KB dummy file
    
    # Create tokenizer files
    tokenizer_config = {
        "model_max_length": 2048,
        "padding_side": "right",
        "tokenizer_class": "LlamaTokenizer"
    }
    
    with open(os.path.join(temp_dir, "tokenizer_config.json"), 'w') as f:
        json.dump(tokenizer_config, f, indent=2)
    
    special_tokens = {
        "bos_token": "<s>",
        "eos_token": "</s>",
        "unk_token": "<unk>",
        "pad_token": "<pad>"
    }
    
    with open(os.path.join(temp_dir, "special_tokens_map.json"), 'w') as f:
        json.dump(special_tokens, f, indent=2)
    
    # Create minimal tokenizer.json
    tokenizer = {
        "version": "1.0",
        "truncation": None,
        "padding": None,
        "added_tokens": [],
        "normalizer": None,
        "pre_tokenizer": None,
        "post_processor": None,
        "decoder": None,
        "model": {
            "type": "BPE",
            "vocab": {},
            "merges": []
        }
    }
    
    with open(os.path.join(temp_dir, "tokenizer.json"), 'w') as f:
        json.dump(tokenizer, f, indent=2)
    
    logger.info(f"Created test model at {temp_dir}")
    return temp_dir


def create_test_dataset(
    dataset_name: str = "test-dataset",
    num_rows: int = 100,
    text_fields: Optional[List[str]] = None
) -> str:
    """
    Create a minimal test dataset (Parquet format).
    
    Args:
        dataset_name: Name for the dataset
        num_rows: Number of rows to generate
        text_fields: List of text field names (default: ["text"])
    
    Returns:
        Path to test dataset file
    
    Example:
        ```python
        from fhenomai.testing import create_test_dataset
        
        dataset_path = create_test_dataset(num_rows=50)
        # Use for testing...
        ```
    """
    try:
        import pandas as pd
        import pyarrow.parquet as pq
    except ImportError:
        raise ImportError("pandas and pyarrow are required for dataset creation. "
                         "Install with: pip install pandas pyarrow")
    
    if text_fields is None:
        text_fields = ["text"]
    
    # Generate sample data
    data = {}
    
    for field in text_fields:
        data[field] = [
            f"Sample {field} content for row {i}. This is test data." 
            for i in range(num_rows)
        ]
    
    # Add some metadata fields
    data["id"] = list(range(num_rows))
    data["label"] = [i % 5 for i in range(num_rows)]
    
    # Create DataFrame
    df = pd.DataFrame(data)
    
    # Save as Parquet
    temp_file = tempfile.NamedTemporaryFile(
        prefix=f"fhenomai_test_{dataset_name}_",
        suffix=".parquet",
        delete=False
    )
    temp_file.close()
    
    df.to_parquet(temp_file.name, engine='pyarrow', index=False)
    
    logger.info(f"Created test dataset at {temp_file.name} ({num_rows} rows)")
    return temp_file.name


def cleanup_test_resources(
    model_paths: Optional[List[str]] = None,
    dataset_paths: Optional[List[str]] = None
):
    """
    Clean up test resources (models and datasets).
    
    Args:
        model_paths: List of model directory paths to delete
        dataset_paths: List of dataset file paths to delete
    
    Example:
        ```python
        from fhenomai.testing import cleanup_test_resources
        
        cleanup_test_resources(
            model_paths=[model_path],
            dataset_paths=[dataset_path]
        )
        ```
    """
    import shutil
    
    if model_paths:
        for path in model_paths:
            if os.path.exists(path):
                try:
                    shutil.rmtree(path)
                    logger.info(f"Deleted test model: {path}")
                except Exception as e:
                    logger.error(f"Failed to delete {path}: {e}")
    
    if dataset_paths:
        for path in dataset_paths:
            if os.path.exists(path):
                try:
                    os.unlink(path)
                    logger.info(f"Deleted test dataset: {path}")
                except Exception as e:
                    logger.error(f"Failed to delete {path}: {e}")


def create_test_config(
    admin_host: str = "localhost",
    admin_port: int = 8000,
    sftp_host: str = "localhost",
    sftp_port: int = 22,
    output_file: Optional[str] = None
) -> str:
    """
    Create a test configuration file.
    
    Args:
        admin_host: Admin API host
        admin_port: Admin API port
        sftp_host: SFTP host
        sftp_port: SFTP port
        output_file: Optional output file path (creates temp file if not provided)
    
    Returns:
        Path to configuration file
    
    Example:
        ```python
        from fhenomai.testing import create_test_config
        
        config_path = create_test_config()
        # Use for testing...
        ```
    """
    import yaml
    
    config_data = {
        "admin": {
            "host": admin_host,
            "port": admin_port,
            "api_key": "test-api-key"
        },
        "sftp": {
            "host": sftp_host,
            "port": sftp_port,
            "username": "test-user",
            "password": "test-password"
        },
        "defaults": {
            "timeout": 30,
            "max_retries": 3
        }
    }
    
    if output_file is None:
        temp_file = tempfile.NamedTemporaryFile(
            prefix="fhenomai_test_config_",
            suffix=".yaml",
            mode='w',
            delete=False
        )
        output_file = temp_file.name
        temp_file.close()
    
    with open(output_file, 'w') as f:
        yaml.dump(config_data, f, default_flow_style=False)
    
    logger.info(f"Created test config at {output_file}")
    return output_file
