from typing import Optional, List, Union, Dict, Any, Pattern
from pydantic import BaseModel, field_validator, model_validator
from typing_extensions import Annotated
import re
from datetime import datetime
import json
from textwrap import dedent
import yaml
from dataclasses import dataclass


class TaggingField(BaseModel):
    """Represents a single tagging field configuration"""
    type: str
    required: bool = False
    description: Optional[str] = None
    allowed_values: Optional[List[Any]] = None  # for enum type
    min_value: Optional[Union[int, float]] = None  # for numeric types
    max_value: Optional[Union[int, float]] = None  # for numeric types

    @field_validator('type', mode='before')
    @classmethod
    def validate_type(cls, v: str) -> str:
        valid_types = ['string', 'integer', 'float', 'date', 'enum', 'color']
        if v not in valid_types:
            raise ValueError(f'type must be one of {valid_types}')
        return v

    @model_validator(mode='after')
    def validate_field_constraints(self) -> 'TaggingField':
        # Validate enum constraints
        if self.type == 'enum':
            if not self.allowed_values:
                raise ValueError('allowed_values must be provided for enum type')
        elif self.allowed_values is not None:
            raise ValueError('allowed_values only valid for enum type')

        # Validate numeric constraints
        if self.type not in ('integer', 'float'):
            if self.min_value is not None or self.max_value is not None:
                raise ValueError('min_value/max_value only valid for numeric types')
        else:
            if self.min_value is not None and self.max_value is not None and self.min_value >= self.max_value:
                raise ValueError('min_value must be less than max_value')

        return self


class TaggingFields(BaseModel):
    """Represents a collection of tagging fields, mapped by their names"""
    fields: Dict[str, TaggingField]

    @classmethod
    def from_dict(cls, data: Dict[str, Dict[str, Any]]) -> 'TaggingFields':
        return cls(fields={
            field_name: TaggingField(**field_config)
            for field_name, field_config in data.items()
        })

    def to_dict(self) -> Dict[str, Dict[str, Any]]:
        return {
            field_name: field.model_dump(exclude_none=True)
            for field_name, field in self.fields.items()
        }


class ChunkingPatternsField(BaseModel):
    """Represents a set of chunking patterns"""
    patterns: List[str]

    @field_validator('patterns')
    def validate_patterns(cls, patterns):
        for pattern in patterns:
            try:
                re.compile(pattern)
            except re.error as e:
                raise ValueError(f"Invalid regex pattern '{pattern}': {str(e)}")
        return patterns


class ArgumentConstraint(BaseModel):
    """Base class for all argument constraints"""
    description: Optional[str] = None
    error_message: Optional[str] = None


class NumericConstraint(ArgumentConstraint):
    """Constraints for numeric values (int/float)"""
    min_value: Optional[float] = None
    max_value: Optional[float] = None
    include_min: bool = True  # True for >= min_value, False for > min_value
    include_max: bool = True  # True for <= max_value, False for < max_value

    @model_validator(mode='after')
    def validate_ranges(self) -> 'NumericConstraint':
        if self.min_value is not None and self.max_value is not None:
            if self.min_value > self.max_value:
                raise ValueError("min_value must be less than or equal to max_value")
        return self

    def validate(self, value: Union[int, float]) -> bool:
        if self.min_value is not None:
            if self.include_min and value < self.min_value:
                return False
            if not self.include_min and value <= self.min_value:
                return False
        if self.max_value is not None:
            if self.include_max and value > self.max_value:
                return False
            if not self.include_max and value >= self.max_value:
                return False
        return True


class StringConstraint(ArgumentConstraint):
    """Constraints for string values"""
    min_length: Optional[int] = None
    max_length: Optional[int] = None
    patterns: Optional[List[str]] = None  # List of regex patterns to match
    pattern_match_all: bool = False  # If True, string must match all patterns
    forbidden_patterns: Optional[List[str]] = None  # List of regex patterns that must not match
    allow_empty: bool = False

    @field_validator('patterns', 'forbidden_patterns')
    @classmethod
    def validate_patterns(cls, v: Optional[List[str]]) -> Optional[List[str]]:
        if v is not None:
            # Validate each pattern compiles
            for pattern in v:
                try:
                    re.compile(pattern)
                except re.error as e:
                    raise ValueError(f"Invalid regex pattern '{pattern}': {str(e)}")
        return v

    def validate(self, value: str) -> bool:
        if not self.allow_empty and not value:
            return False

        if self.min_length is not None and len(value) < self.min_length:
            return False

        if self.max_length is not None and len(value) > self.max_length:
            return False

        if self.patterns:
            matches = [bool(re.search(pattern, value)) for pattern in self.patterns]
            if self.pattern_match_all and not all(matches):
                return False
            if not self.pattern_match_all and not any(matches):
                return False

        if self.forbidden_patterns:
            for pattern in self.forbidden_patterns:
                if re.search(pattern, value):
                    return False

        return True


class DateConstraint(ArgumentConstraint):
    """Constraints for date values"""
    min_date: Optional[datetime] = None
    max_date: Optional[datetime] = None
    include_min: bool = True
    include_max: bool = True
    allowed_formats: Optional[List[str]] = None  # List of allowed date formats

    @model_validator(mode='after')
    def validate_ranges(self) -> 'DateConstraint':
        if self.min_date and self.max_date and self.min_date > self.max_date:
            raise ValueError("min_date must be less than or equal to max_date")
        return self

    def validate(self, value: datetime) -> bool:
        if self.min_date is not None:
            if self.include_min and value < self.min_date:
                return False
            if not self.include_min and value <= self.min_date:
                return False

        if self.max_date is not None:
            if self.include_max and value > self.max_date:
                return False
            if not self.include_max and value >= self.max_date:
                return False

        return True


class EnumConstraint(ArgumentConstraint):
    """Constraints for enum values"""
    allowed_values: List[Any]
    case_sensitive: bool = True  # For string enums
    allow_multiple: bool = False  # If True, value can be a list of allowed values
    min_selections: Optional[int] = None  # When allow_multiple is True
    max_selections: Optional[int] = None  # When allow_multiple is True

    @model_validator(mode='after')
    def validate_selections(self) -> 'EnumConstraint':
        if self.allow_multiple:
            if self.min_selections is not None and self.max_selections is not None:
                if self.min_selections > self.max_selections:
                    raise ValueError("min_selections must be less than or equal to max_selections")
                if self.max_selections > len(self.allowed_values):
                    raise ValueError("max_selections cannot be greater than number of allowed values")
        return self

    def validate(self, value: Union[Any, List[Any]]) -> bool:
        if self.allow_multiple:
            if not isinstance(value, list):
                return False

            if self.min_selections is not None and len(value) < self.min_selections:
                return False

            if self.max_selections is not None and len(value) > self.max_selections:
                return False

            for v in value:
                if not self._validate_single_value(v):
                    return False
        else:
            return self._validate_single_value(value)

        return True

    def _validate_single_value(self, value: Any) -> bool:
        if isinstance(value, str) and not self.case_sensitive:
            return any(str(value).lower() == str(v).lower() for v in self.allowed_values)
        return value in self.allowed_values


class ArgumentDefinition(BaseModel):
    """Defines an argument with its type and constraints"""
    name: str
    type: str
    description: Optional[str] = None
    required: bool = False
    default: Optional[Any] = None
    constraints: Optional[Union[NumericConstraint, StringConstraint, DateConstraint, EnumConstraint]] = None

    @field_validator('type')
    @classmethod
    def validate_type(cls, v: str) -> str:
        valid_types = ['string', 'integer', 'float', 'date', 'enum', 'color']
        if v not in valid_types:
            raise ValueError(f'type must be one of {valid_types}')
        return v

    @model_validator(mode='after')
    def validate_constraints(self) -> 'ArgumentDefinition':
        if self.constraints:
            expected_constraint_types = {
                'string': StringConstraint,
                'integer': NumericConstraint,
                'float': NumericConstraint,
                'date': DateConstraint,
                'enum': EnumConstraint,
                'color': StringConstraint
            }

            expected_type = expected_constraint_types.get(self.type)
            if not isinstance(self.constraints, expected_type):
                raise ValueError(f'Constraints for type {self.type} must be of type {expected_type.__name__}')

            if self.default is not None:
                if not self.constraints.validate(self.default):
                    raise ValueError(f'Default value does not satisfy constraints for {self.name}')

        return self


class ArgumentDefinitions(BaseModel):
    """Collection of argument definitions"""
    arguments: Dict[str, ArgumentDefinition]

    @classmethod
    def from_dict(cls, data: Dict[str, Dict[str, Any]]) -> 'ArgumentDefinitions':
        return cls(arguments={
            arg_name: ArgumentDefinition(**arg_config)
            for arg_name, arg_config in data.items()
        })

    def to_dict(self) -> Dict[str, Dict[str, Any]]:
        return {
            arg_name: arg.model_dump(exclude_none=True)
            for arg_name, arg in self.arguments.items()
        }

    def validate_argument_values(self, values: Dict[str, Any]) -> Dict[str, str]:
        """
        Validate a set of argument values against their definitions
        Returns a dictionary of error messages for invalid arguments
        """
        errors = {}

        # Check for required arguments
        for name, arg_def in self.arguments.items():
            if arg_def.required and name not in values:
                errors[name] = "Required argument missing"
                continue

            if name in values:
                value = values[name]

                # Validate type
                try:
                    if arg_def.type == 'integer':
                        value = int(value)
                    elif arg_def.type == 'float':
                        value = float(value)
                    elif arg_def.type == 'date' and isinstance(value, str):
                        if arg_def.constraints and arg_def.constraints.allowed_formats:
                            for fmt in arg_def.constraints.allowed_formats:
                                try:
                                    value = datetime.strptime(value, fmt)
                                    break
                                except ValueError:
                                    continue
                            else:
                                errors[
                                    name] = f"Invalid date format. Allowed formats: {arg_def.constraints.allowed_formats}"
                                continue
                except (ValueError, TypeError):
                    errors[name] = f"Invalid type. Expected {arg_def.type}"
                    continue

                # Validate constraints
                if arg_def.constraints and not arg_def.constraints.validate(value):
                    errors[name] = arg_def.constraints.error_message or "Value does not satisfy constraints"

        return errors


@dataclass
class DocumentationFormat:
    """Constants for documentation formats"""
    MARKDOWN = "markdown"
    JSON = "json"
    YAML = "yaml"


@dataclass
class DocumentationVersion:
    """Constants for documentation versions"""
    BASIC = "basic"  # Original documentation without retriever info
    EXTENDED = "extended"  # Including retriever documentation


def _generate_argument_constraints(field_config: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Generate possible argument constraints based on field type"""
    constraints = []

    base_constraint = {
        "description": f"Constraint for {field_config.get('description', 'field')}",
        "error_message": "Optional custom error message"
    }

    if field_config["type"] == "integer" or field_config["type"] == "float":
        constraints.append({
            **base_constraint,
            "type": "NumericConstraint",
            "possible_constraints": {
                "min_value": "number",
                "max_value": "number",
                "include_min": "boolean",
                "include_max": "boolean"
            },
            "example": {
                "min_value": field_config.get("min_value", 0),
                "max_value": field_config.get("max_value", 100),
                "include_min": True,
                "include_max": True
            }
        })

    elif field_config["type"] == "string":
        constraints.append({
            **base_constraint,
            "type": "StringConstraint",
            "possible_constraints": {
                "min_length": "integer",
                "max_length": "integer",
                "patterns": "list[str]",
                "pattern_match_all": "boolean",
                "forbidden_patterns": "list[str]",
                "allow_empty": "boolean"
            },
            "example": {
                "min_length": 1,
                "max_length": 100,
                "patterns": ["^[A-Za-z0-9]+$"],
                "pattern_match_all": False,
                "forbidden_patterns": ["^test_", "_temp$"],
                "allow_empty": False
            }
        })

    elif field_config["type"] == "enum":
        constraints.append({
            **base_constraint,
            "type": "EnumConstraint",
            "possible_constraints": {
                "allowed_values": f"list[{field_config.get('allowed_values', ['value1', 'value2'])}]",
                "case_sensitive": "boolean",
                "allow_multiple": "boolean",
                "min_selections": "integer",
                "max_selections": "integer"
            },
            "example": {
                "allowed_values": field_config.get("allowed_values", ["value1", "value2"]),
                "case_sensitive": True,
                "allow_multiple": True,
                "min_selections": 1,
                "max_selections": 2
            }
        })

    elif field_config["type"] == "date":
        constraints.append({
            **base_constraint,
            "type": "DateConstraint",
            "possible_constraints": {
                "min_date": "datetime",
                "max_date": "datetime",
                "include_min": "boolean",
                "include_max": "boolean",
                "allowed_formats": "list[str]"
            },
            "example": {
                "min_date": "2024-01-01T00:00:00",
                "max_date": "2024-12-31T23:59:59",
                "include_min": True,
                "include_max": True,
                "allowed_formats": ["%Y-%m-%d", "%Y/%m/%d"]
            }
        })

    return constraints


def generate_field_documentation(
        tagging_fields: Dict[str, Any],
        format: str = "markdown",
        version: str = "basic"
) -> str:
    """
    Generate documentation for tagging fields configuration.

    Args:
        tagging_fields: Dictionary containing tagging fields configuration
        format: Output format ("markdown", "json", or "yaml")
        version: Documentation version ("basic" or "extended")

    Returns:
        str: Formatted documentation
    """
    if version not in [DocumentationVersion.BASIC, DocumentationVersion.EXTENDED]:
        raise ValueError(f"Unsupported documentation version: {version}")

    # Normalize fields configuration
    normalized_fields = {}

    for field_name, field_config in tagging_fields.items():
        field_doc = {
            "name": field_name,
            "type": field_config["type"],
            "required": field_config.get("required", False),
            "description": field_config.get("description", "No description provided"),
            "constraints": []
        }

        # Only include possible arguments in extended version
        if version == DocumentationVersion.EXTENDED:
            field_doc["possible_arguments"] = _generate_argument_constraints(field_config)

        # Add type-specific constraints
        if field_config["type"] == "integer" or field_config["type"] == "float":
            if "min_value" in field_config:
                field_doc["constraints"].append(
                    f"Minimum value: {field_config['min_value']}")
            if "max_value" in field_config:
                field_doc["constraints"].append(
                    f"Maximum value: {field_config['max_value']}")

        elif field_config["type"] == "string":
            if "min_length" in field_config:
                field_doc["constraints"].append(
                    f"Minimum length: {field_config['min_length']}")
            if "max_length" in field_config:
                field_doc["constraints"].append(
                    f"Maximum length: {field_config['max_length']}")
            if "patterns" in field_config:
                field_doc["constraints"].append(
                    f"Must match patterns: {', '.join(field_config['patterns'])}")

        elif field_config["type"] == "enum":
            if "allowed_values" in field_config:
                field_doc["constraints"].append(
                    f"Allowed values: {', '.join(str(v) for v in field_config['allowed_values'])}")

        elif field_config["type"] == "date":
            if "min_date" in field_config:
                field_doc["constraints"].append(
                    f"Minimum date: {field_config['min_date']}")
            if "max_date" in field_config:
                field_doc["constraints"].append(
                    f"Maximum date: {field_config['max_date']}")
            if "allowed_formats" in field_config:
                field_doc["constraints"].append(
                    f"Allowed formats: {', '.join(field_config['allowed_formats'])}")

        normalized_fields[field_name] = field_doc

    # Generate documentation in requested format
    if format == DocumentationFormat.MARKDOWN:
        return _generate_markdown_docs(normalized_fields, version)
    elif format == DocumentationFormat.JSON:
        return _generate_json_docs(normalized_fields, version)
    elif format == DocumentationFormat.YAML:
        return _generate_yaml_docs(normalized_fields, version)
    else:
        raise ValueError(f"Unsupported documentation format: {format}")


def _generate_markdown_docs(fields: Dict[str, Any], version: str) -> str:
    """Generate markdown documentation"""
    docs = ["# Tagging Fields Documentation\n"]

    # Add overview table
    docs.append("## Fields Overview\n")
    docs.append("| Field Name | Type | Required | Description |")
    docs.append("|------------|------|----------|-------------|")

    for field_name, field in fields.items():
        docs.append(
            f"| {field_name} | {field['type']} | "
            f"{'Yes' if field['required'] else 'No'} | {field['description']} |"
        )

    # Add detailed field specifications
    docs.append("\n## Detailed Field Specifications\n")

    for field_name, field in fields.items():
        docs.append(f"### {field_name}\n")
        docs.append(f"**Type:** {field['type']}")
        docs.append(f"**Required:** {'Yes' if field['required'] else 'No'}")
        docs.append(f"**Description:** {field['description']}\n")

        if field["constraints"]:
            docs.append("**Field Constraints:**")
            for constraint in field["constraints"]:
                docs.append(f"- {constraint}")
            docs.append("")

        # Add retriever argument documentation only in extended version
        if version == DocumentationVersion.EXTENDED and "possible_arguments" in field:
            docs.append("**Possible Retriever Arguments:**")
            for arg_constraint in field["possible_arguments"]:
                docs.append(f"\n*{arg_constraint['type']}*")
                docs.append(f"Description: {arg_constraint['description']}")
                docs.append("\nPossible constraints:")
                for const_name, const_type in arg_constraint["possible_constraints"].items():
                    docs.append(f"- `{const_name}`: {const_type}")

                docs.append("\nExample:")
                docs.append("```python")
                docs.append(json.dumps(arg_constraint["example"], indent=2))
                docs.append("```\n")

    # Add example retriever configuration only in extended version
    if version == DocumentationVersion.EXTENDED:
        docs.append("\n## Example Retriever Configuration\n")
        docs.append("```python")
        example_config = {
            "metadata_filters": {
                field_name: field["possible_arguments"][0]["example"]
                for field_name, field in fields.items()
                if "possible_arguments" in field
            }
        }
        docs.append(json.dumps(example_config, indent=2))
        docs.append("```")

    return "\n".join(docs)


def _generate_json_docs(fields: Dict[str, Any], version: str) -> str:
    """Generate JSON documentation"""
    doc = {
        "tagging_fields_documentation": {
            "version": version,
            "fields": fields
        }
    }

    if version == DocumentationVersion.EXTENDED:
        doc["tagging_fields_documentation"]["example_retriever_config"] = {
            "metadata_filters": {
                field_name: field["possible_arguments"][0]["example"]
                for field_name, field in fields.items()
                if "possible_arguments" in field
            }
        }

    return json.dumps(doc, indent=2)


def _generate_yaml_docs(fields: Dict[str, Any], version: str) -> str:
    """Generate YAML documentation"""
    doc = {
        "tagging_fields_documentation": {
            "version": version,
            "fields": fields
        }
    }

    if version == DocumentationVersion.EXTENDED:
        doc["tagging_fields_documentation"]["example_retriever_config"] = {
            "metadata_filters": {
                field_name: field["possible_arguments"][0]["example"]
                for field_name, field in fields.items()
                if "possible_arguments" in field
            }
        }

    return yaml.dump(doc, sort_keys=False, default_flow_style=False)


def patterns_to_json(text_area_content: str) -> str:
    """Convert line-based patterns to JSON"""
    text_area_content = text_area_content.strip()
    if len(text_area_content) == 0:
        return json.dumps([])
    # Split on newlines and remove empty lines
    patterns = [line.strip() for line in text_area_content.split('\n') if line.strip()]
    return json.dumps(patterns)


def json_to_patterns(json_content: str) -> str:
    """Convert JSON patterns list to text area content"""
    try:
        patterns = json.loads(json_content)
        if not isinstance(patterns, list):
            raise ValueError("JSON must contain a list of patterns")
        # Join with newlines
        return '\n'.join(patterns)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON format: {e}")


def json_to_pattern_list(json_content: str) -> list:
    """Convert JSON patterns list to text area content"""
    try:
        if json_content:
            patterns = json.loads(json_content)
            if not isinstance(patterns, list):
                raise ValueError("JSON must contain a list of patterns")
            # Unescape if needed
            patterns = [pattern.replace('\\\\', '\\') for pattern in patterns]
            return patterns
        else:
            return []
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON format: {e}")


def normalize_json_field(value: str | dict | None, field_name: str = "JSON field") -> dict:
    """
    Normalize a JSON field value to ensure it's a valid dictionary.

    Args:
        value: The input value which can be:
            - None (will return empty dict)
            - String (will be parsed as JSON)
            - Dict (will be validated and returned)
        field_name: Name of the field for error messages

    Returns:
        dict: The normalized JSON data as a Python dictionary

    Raises:
        ValueError: If the input string is not valid JSON or the input dict contains invalid types
    """
    # Handle None case
    if value is None:
        return {}

    # Handle dictionary case
    if isinstance(value, dict):
        try:
            # Validate all values are JSON serializable
            import json
            json.dumps(value)
            return value
        except TypeError as e:
            raise ValueError(f"{field_name} contains invalid types: {str(e)}")

    # Handle string case
    if isinstance(value, str):
        if not value.strip():
            return {}

        try:
            import json
            return json.loads(value)
        except json.JSONDecodeError as e:
            raise ValueError(f"{field_name} contains invalid JSON: {str(e)}")

    raise ValueError(f"{field_name} must be a string, dictionary, or None (got {type(value)})")
