"""Flexible data generation system for custom requirements."""

from typing import List, Dict, Any, Optional, Callable, Union
from dataclasses import dataclass, field
import random
import re
import json
from datetime import datetime, timedelta
from enum import Enum

from ..core.base import Example, Dataset
from ..utils.llm_clients import BaseLLMClient


class FieldType(Enum):
    """Supported field types for generation."""
    TEXT = "text"
    NUMBER = "number"
    DATE = "date"
    ENUM = "enum"
    JSON = "json"
    TEMPLATE = "template"
    FUNCTION = "function"
    REGEX = "regex"
    COMPOSITE = "composite"


@dataclass
class FieldDefinition:
    """Define how a field should be generated."""
    name: str
    field_type: FieldType
    description: str = ""
    
    # Type-specific parameters
    options: Optional[List[Any]] = None  # For ENUM
    min_value: Optional[Union[int, float]] = None  # For NUMBER
    max_value: Optional[Union[int, float]] = None  # For NUMBER
    pattern: Optional[str] = None  # For REGEX or TEMPLATE
    generator: Optional[Callable] = None  # For FUNCTION
    sub_fields: Optional[Dict[str, 'FieldDefinition']] = None  # For COMPOSITE
    format: Optional[str] = None  # For DATE or custom formatting
    constraints: Optional[List[str]] = None  # Validation rules
    dependencies: Optional[List[str]] = None  # Fields this depends on
    weight: Optional[Dict[Any, float]] = None  # Probability weights for ENUM


@dataclass
class GenerationTemplate:
    """Custom template for data generation."""
    name: str
    description: str
    input_fields: Dict[str, FieldDefinition]
    output_fields: Dict[str, FieldDefinition]
    
    # Optional components
    input_template: Optional[str] = None
    output_template: Optional[str] = None
    validation_rules: List[Callable] = field(default_factory=list)
    post_processors: List[Callable] = field(default_factory=list)
    examples: List[Dict[str, Any]] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)


class FlexibleDataGenerator:
    """Highly customizable data generation system."""
    
    def __init__(self, llm_client: Optional[BaseLLMClient] = None):
        self.llm_client = llm_client
        self.templates: Dict[str, GenerationTemplate] = {}
        self.custom_generators: Dict[str, Callable] = {}
        self._setup_default_generators()
    
    def _setup_default_generators(self):
        """Set up commonly used field generators."""
        self.custom_generators.update({
            "email": lambda: self._generate_email(),
            "phone": lambda: self._generate_phone(),
            "name": lambda: self._generate_name(),
            "company": lambda: self._generate_company(),
            "url": lambda: self._generate_url(),
            "address": lambda: self._generate_address(),
            "product_id": lambda: self._generate_product_id(),
            "transaction_id": lambda: self._generate_transaction_id(),
            "timestamp": lambda: self._generate_timestamp(),
            "price": lambda: self._generate_price(),
            "percentage": lambda: self._generate_percentage(),
            "code": lambda: self._generate_code(),
            "api_key": lambda: self._generate_api_key(),
            "uuid": lambda: self._generate_uuid()
        })
    
    def register_template(self, template: GenerationTemplate):
        """Register a custom generation template."""
        self.templates[template.name] = template
    
    def register_generator(self, name: str, generator: Callable):
        """Register a custom field generator function."""
        self.custom_generators[name] = generator
    
    def create_template(self, name: str, description: str,
                       input_spec: Dict[str, Union[str, Dict]], 
                       output_spec: Dict[str, Union[str, Dict]],
                       **kwargs) -> GenerationTemplate:
        """Create a template from simplified specifications."""
        input_fields = self._parse_field_specs(input_spec)
        output_fields = self._parse_field_specs(output_spec)
        
        return GenerationTemplate(
            name=name,
            description=description,
            input_fields=input_fields,
            output_fields=output_fields,
            **kwargs
        )
    
    def _parse_field_specs(self, specs: Dict[str, Union[str, Dict]]) -> Dict[str, FieldDefinition]:
        """Parse simplified field specifications into FieldDefinitions."""
        fields = {}
        
        for field_name, spec in specs.items():
            if isinstance(spec, str):
                # Simple type specification
                fields[field_name] = self._create_field_from_string(field_name, spec)
            elif isinstance(spec, dict):
                # Detailed specification
                fields[field_name] = self._create_field_from_dict(field_name, spec)
            else:
                raise ValueError(f"Invalid field specification for {field_name}")
        
        return fields
    
    def _create_field_from_string(self, name: str, type_str: str) -> FieldDefinition:
        """Create field definition from type string."""
        # Handle special formats like "enum:option1,option2,option3"
        if ":" in type_str:
            base_type, params = type_str.split(":", 1)
            
            if base_type == "enum":
                options = [opt.strip() for opt in params.split(",")]
                return FieldDefinition(
                    name=name,
                    field_type=FieldType.ENUM,
                    options=options
                )
            elif base_type == "regex":
                return FieldDefinition(
                    name=name,
                    field_type=FieldType.REGEX,
                    pattern=params
                )
            elif base_type == "template":
                return FieldDefinition(
                    name=name,
                    field_type=FieldType.TEMPLATE,
                    pattern=params
                )
            elif base_type == "generator":
                if params in self.custom_generators:
                    return FieldDefinition(
                        name=name,
                        field_type=FieldType.FUNCTION,
                        generator=self.custom_generators[params]
                    )
        
        # Map simple type strings to FieldType
        type_map = {
            "text": FieldType.TEXT,
            "number": FieldType.NUMBER,
            "date": FieldType.DATE,
            "json": FieldType.JSON
        }
        
        field_type = type_map.get(type_str, FieldType.TEXT)
        return FieldDefinition(name=name, field_type=field_type)
    
    def _create_field_from_dict(self, name: str, spec: Dict[str, Any]) -> FieldDefinition:
        """Create field definition from detailed specification."""
        field_type = FieldType(spec.get("type", "text"))
        
        return FieldDefinition(
            name=name,
            field_type=field_type,
            description=spec.get("description", ""),
            options=spec.get("options"),
            min_value=spec.get("min"),
            max_value=spec.get("max"),
            pattern=spec.get("pattern"),
            generator=spec.get("generator"),
            sub_fields=spec.get("sub_fields"),
            format=spec.get("format"),
            constraints=spec.get("constraints", []),
            dependencies=spec.get("dependencies", []),
            weight=spec.get("weight")
        )
    
    def generate(self, template_name: str, count: int,
                context: Optional[Dict[str, Any]] = None) -> Dataset:
        """Generate data using a registered template."""
        if template_name not in self.templates:
            raise ValueError(f"Template '{template_name}' not found")
        
        template = self.templates[template_name]
        examples = []
        
        for i in range(count):
            example = self._generate_example(template, context or {}, index=i)
            examples.append(example)
        
        return Dataset(
            examples=examples,
            name=f"{template_name}_synthetic",
            metadata={
                "synthetic": True,
                "template": template_name,
                "count": count,
                "context": context,
                "generated_at": datetime.now().isoformat()
            }
        )
    
    def _generate_example(self, template: GenerationTemplate,
                         context: Dict[str, Any], index: int = 0) -> Example:
        """Generate a single example from template."""
        # Generate input fields
        input_data = {}
        for field_name, field_def in template.input_fields.items():
            input_data[field_name] = self._generate_field_value(
                field_def, context, input_data
            )
        
        # Generate output fields (may depend on input)
        output_data = {}
        combined_context = {**context, **input_data}
        
        for field_name, field_def in template.output_fields.items():
            output_data[field_name] = self._generate_field_value(
                field_def, combined_context, output_data
            )
        
        # Format using templates if provided
        if template.input_template:
            input_text = self._format_template(template.input_template, input_data)
        else:
            input_text = json.dumps(input_data, indent=2)
        
        if template.output_template:
            output_text = self._format_template(template.output_template, output_data)
        else:
            output_text = json.dumps(output_data, indent=2)
        
        # Apply validation
        for validator in template.validation_rules:
            if not validator(input_data, output_data):
                # Retry generation if validation fails
                return self._generate_example(template, context, index)
        
        # Apply post-processing
        for processor in template.post_processors:
            input_text, output_text = processor(input_text, output_text)
        
        return Example(
            input=input_text,
            output=output_text,
            metadata={
                "template": template.name,
                "index": index,
                "input_data": input_data,
                "output_data": output_data
            }
        )
    
    def _generate_field_value(self, field_def: FieldDefinition,
                             context: Dict[str, Any],
                             current_data: Dict[str, Any]) -> Any:
        """Generate value for a specific field."""
        # Check dependencies
        if field_def.dependencies:
            for dep in field_def.dependencies:
                if dep not in context and dep not in current_data:
                    raise ValueError(f"Dependency '{dep}' not found for field '{field_def.name}'")
        
        # Generate based on field type
        if field_def.field_type == FieldType.TEXT:
            return self._generate_text(field_def, context)
        
        elif field_def.field_type == FieldType.NUMBER:
            return self._generate_number(field_def)
        
        elif field_def.field_type == FieldType.DATE:
            return self._generate_date(field_def)
        
        elif field_def.field_type == FieldType.ENUM:
            return self._generate_enum(field_def)
        
        elif field_def.field_type == FieldType.JSON:
            return self._generate_json(field_def, context)
        
        elif field_def.field_type == FieldType.TEMPLATE:
            return self._generate_from_template(field_def, {**context, **current_data})
        
        elif field_def.field_type == FieldType.FUNCTION:
            if field_def.generator:
                return field_def.generator(context, current_data)
            else:
                raise ValueError(f"No generator function for field '{field_def.name}'")
        
        elif field_def.field_type == FieldType.REGEX:
            return self._generate_from_regex(field_def)
        
        elif field_def.field_type == FieldType.COMPOSITE:
            return self._generate_composite(field_def, context)
        
        else:
            raise ValueError(f"Unsupported field type: {field_def.field_type}")
    
    def _generate_text(self, field_def: FieldDefinition, context: Dict[str, Any]) -> str:
        """Generate text value."""
        if self.llm_client and field_def.description:
            # Use LLM for more realistic text
            prompt = f"Generate a {field_def.description}. Context: {json.dumps(context, indent=2)}"
            response = self.llm_client.generate(prompt, max_tokens=100)
            return response.text.strip()
        else:
            # Fallback to simple generation
            return f"Generated {field_def.name}"
    
    def _generate_number(self, field_def: FieldDefinition) -> Union[int, float]:
        """Generate numeric value."""
        min_val = field_def.min_value or 0
        max_val = field_def.max_value or 1000
        
        if isinstance(min_val, int) and isinstance(max_val, int):
            return random.randint(min_val, max_val)
        else:
            return round(random.uniform(min_val, max_val), 2)
    
    def _generate_date(self, field_def: FieldDefinition) -> str:
        """Generate date value."""
        base = datetime.now()
        offset = random.randint(-365, 365)
        date = base + timedelta(days=offset)
        
        format_str = field_def.format or "%Y-%m-%d"
        return date.strftime(format_str)
    
    def _generate_enum(self, field_def: FieldDefinition) -> Any:
        """Generate enum value."""
        if not field_def.options:
            raise ValueError(f"No options provided for enum field '{field_def.name}'")
        
        if field_def.weight:
            # Weighted random selection
            choices = list(field_def.weight.keys())
            weights = list(field_def.weight.values())
            return random.choices(choices, weights=weights)[0]
        else:
            return random.choice(field_def.options)
    
    def _generate_json(self, field_def: FieldDefinition, context: Dict[str, Any]) -> Dict:
        """Generate JSON value."""
        if field_def.sub_fields:
            result = {}
            for sub_name, sub_def in field_def.sub_fields.items():
                result[sub_name] = self._generate_field_value(sub_def, context, result)
            return result
        else:
            return {"generated": field_def.name}
    
    def _generate_from_template(self, field_def: FieldDefinition, context: Dict[str, Any]) -> str:
        """Generate value from template pattern."""
        if not field_def.pattern:
            raise ValueError(f"No pattern provided for template field '{field_def.name}'")
        
        return self._format_template(field_def.pattern, context)
    
    def _generate_from_regex(self, field_def: FieldDefinition) -> str:
        """Generate value matching regex pattern."""
        # This is a simplified implementation
        # In production, you might use a library like rstr
        if not field_def.pattern:
            raise ValueError(f"No pattern provided for regex field '{field_def.name}'")
        
        # Handle common patterns
        pattern = field_def.pattern
        if pattern == r'\d+':
            return str(random.randint(1000, 9999))
        elif pattern == r'[A-Z]{3}\d{3}':
            letters = ''.join(random.choices('ABCDEFGHIJKLMNOPQRSTUVWXYZ', k=3))
            numbers = ''.join(random.choices('0123456789', k=3))
            return f"{letters}{numbers}"
        else:
            return f"REGEX_{field_def.name}"
    
    def _generate_composite(self, field_def: FieldDefinition, context: Dict[str, Any]) -> Dict:
        """Generate composite value with sub-fields."""
        if not field_def.sub_fields:
            raise ValueError(f"No sub-fields defined for composite field '{field_def.name}'")
        
        result = {}
        for sub_name, sub_def in field_def.sub_fields.items():
            result[sub_name] = self._generate_field_value(sub_def, context, result)
        
        return result
    
    def _format_template(self, template: str, data: Dict[str, Any]) -> str:
        """Format template string with data."""
        # Handle nested data access like {user.name}
        def replace_func(match):
            key = match.group(1)
            if '.' in key:
                parts = key.split('.')
                value = data
                for part in parts:
                    if isinstance(value, dict) and part in value:
                        value = value[part]
                    else:
                        return match.group(0)
                return str(value)
            else:
                return str(data.get(key, match.group(0)))
        
        return re.sub(r'\{([^}]+)\}', replace_func, template)
    
    # Default generators
    def _generate_email(self) -> str:
        first = random.choice(['john', 'jane', 'alex', 'sam', 'chris'])
        last = random.choice(['smith', 'jones', 'brown', 'davis', 'wilson'])
        domain = random.choice(['example.com', 'test.org', 'demo.net'])
        return f"{first}.{last}@{domain}"
    
    def _generate_phone(self) -> str:
        return f"+1-{random.randint(200,999)}-{random.randint(200,999)}-{random.randint(1000,9999)}"
    
    def _generate_name(self) -> str:
        first = random.choice(['John', 'Jane', 'Alex', 'Sam', 'Chris'])
        last = random.choice(['Smith', 'Jones', 'Brown', 'Davis', 'Wilson'])
        return f"{first} {last}"
    
    def _generate_company(self) -> str:
        prefix = random.choice(['Tech', 'Global', 'Digital', 'Smart', 'Next'])
        suffix = random.choice(['Corp', 'Solutions', 'Systems', 'Group', 'Inc'])
        return f"{prefix} {suffix}"
    
    def _generate_url(self) -> str:
        domain = random.choice(['example', 'test', 'demo'])
        tld = random.choice(['com', 'org', 'net'])
        path = random.choice(['', '/products', '/about', '/contact'])
        return f"https://{domain}.{tld}{path}"
    
    def _generate_address(self) -> str:
        number = random.randint(100, 9999)
        street = random.choice(['Main St', 'Oak Ave', 'Park Rd', 'First St'])
        city = random.choice(['New York', 'Los Angeles', 'Chicago', 'Houston'])
        state = random.choice(['NY', 'CA', 'IL', 'TX'])
        zip_code = random.randint(10000, 99999)
        return f"{number} {street}, {city}, {state} {zip_code}"
    
    def _generate_product_id(self) -> str:
        prefix = random.choice(['PRD', 'ITM', 'SKU'])
        return f"{prefix}-{random.randint(100000, 999999)}"
    
    def _generate_transaction_id(self) -> str:
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        random_suffix = random.randint(1000, 9999)
        return f"TXN-{timestamp}-{random_suffix}"
    
    def _generate_timestamp(self) -> str:
        offset_hours = random.randint(-720, 720)  # +/- 30 days
        timestamp = datetime.now() + timedelta(hours=offset_hours)
        return timestamp.isoformat()
    
    def _generate_price(self) -> float:
        return round(random.uniform(9.99, 999.99), 2)
    
    def _generate_percentage(self) -> float:
        return round(random.uniform(0, 100), 1)
    
    def _generate_code(self) -> str:
        chars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
        return ''.join(random.choices(chars, k=8))
    
    def _generate_api_key(self) -> str:
        chars = 'abcdefghijklmnopqrstuvwxyz0123456789'
        return 'sk-' + ''.join(random.choices(chars, k=32))
    
    def _generate_uuid(self) -> str:
        import uuid
        return str(uuid.uuid4())


# Builder pattern for easy template creation
class TemplateBuilder:
    """Fluent interface for building generation templates."""
    
    def __init__(self, name: str):
        self.name = name
        self.description = ""
        self.input_fields = {}
        self.output_fields = {}
        self.input_template = None
        self.output_template = None
        self.validation_rules = []
        self.post_processors = []
        self.examples = []
        self.metadata = {}
    
    def with_description(self, description: str) -> 'TemplateBuilder':
        self.description = description
        return self
    
    def add_input_field(self, name: str, field_type: str, **kwargs) -> 'TemplateBuilder':
        if isinstance(field_type, str):
            # Parse string specification
            if ":" in field_type:
                base_type, params = field_type.split(":", 1)
                if base_type == "enum":
                    options = [opt.strip() for opt in params.split(",")]
                    self.input_fields[name] = FieldDefinition(
                        name=name,
                        field_type=FieldType.ENUM,
                        options=options,
                        **kwargs
                    )
                elif base_type == "regex":
                    self.input_fields[name] = FieldDefinition(
                        name=name,
                        field_type=FieldType.REGEX,
                        pattern=params,
                        **kwargs
                    )
            else:
                type_map = {
                    "text": FieldType.TEXT,
                    "number": FieldType.NUMBER,
                    "date": FieldType.DATE,
                    "json": FieldType.JSON
                }
                self.input_fields[name] = FieldDefinition(
                    name=name,
                    field_type=type_map.get(field_type, FieldType.TEXT),
                    **kwargs
                )
        else:
            self.input_fields[name] = FieldDefinition(
                name=name,
                field_type=field_type,
                **kwargs
            )
        return self
    
    def add_output_field(self, name: str, field_type: str, **kwargs) -> 'TemplateBuilder':
        # Similar logic to add_input_field
        if isinstance(field_type, str):
            type_map = {
                "text": FieldType.TEXT,
                "number": FieldType.NUMBER,
                "date": FieldType.DATE,
                "json": FieldType.JSON
            }
            self.output_fields[name] = FieldDefinition(
                name=name,
                field_type=type_map.get(field_type, FieldType.TEXT),
                **kwargs
            )
        else:
            self.output_fields[name] = FieldDefinition(
                name=name,
                field_type=field_type,
                **kwargs
            )
        return self
    
    def with_input_template(self, template: str) -> 'TemplateBuilder':
        self.input_template = template
        return self
    
    def with_output_template(self, template: str) -> 'TemplateBuilder':
        self.output_template = template
        return self
    
    def add_validation(self, rule: Callable) -> 'TemplateBuilder':
        self.validation_rules.append(rule)
        return self
    
    def add_post_processor(self, processor: Callable) -> 'TemplateBuilder':
        self.post_processors.append(processor)
        return self
    
    def add_example(self, example: Dict[str, Any]) -> 'TemplateBuilder':
        self.examples.append(example)
        return self
    
    def with_metadata(self, **metadata) -> 'TemplateBuilder':
        self.metadata.update(metadata)
        return self
    
    def build(self) -> GenerationTemplate:
        return GenerationTemplate(
            name=self.name,
            description=self.description,
            input_fields=self.input_fields,
            output_fields=self.output_fields,
            input_template=self.input_template,
            output_template=self.output_template,
            validation_rules=self.validation_rules,
            post_processors=self.post_processors,
            examples=self.examples,
            metadata=self.metadata
        )