import ast
from datetime import datetime
from enum import Enum
from typing import Any, List, Literal, Optional

from kajson.kajson_manager import KajsonManager
from pydantic import Field

from pipelex.core.concepts.concept_structure_blueprint import ConceptStructureBlueprint, ConceptStructureBlueprintFieldType
from pipelex.core.concepts.helpers import extract_concept_code_from_concept_ref_or_code
from pipelex.core.concepts.native.concept_native import NativeConceptCode
from pipelex.core.concepts.structure_generation.exceptions import ConceptStructureGeneratorError, ConceptStructureValidationError, SyntaxErrorData
from pipelex.core.stuffs.structured_content import StructuredContent
from pipelex.core.stuffs.stuff_content import StuffContent


class ConceptClassInfo:
    """Information about a concept's structure class for import generation."""

    def __init__(self, class_name: str, module_path: str | None = None):
        """Initialize ConceptClassInfo.

        Args:
            class_name: The name of the structure class (e.g., "Skill")
            module_path: Optional module path where the class is defined (e.g., "myapp.structures.skill").
                         If None, a forward reference (string) will be used.
        """
        self.class_name = class_name
        self.module_path = module_path


class StructureGenerator:
    """Generate Pydantic BaseModel classes from concept structure blueprints."""

    # TODO: use StrEnum instead of Enum
    def __init__(
        self,
        concept_ref_to_class_info: dict[str, ConceptClassInfo] | None = None,
    ):
        """Initialize the StructureGenerator.

        Args:
            concept_ref_to_class_info: Optional mapping from concept refs to ConceptClassInfo.
                When module_path is provided, proper import statements will be generated.
                When module_path is None, forward references (strings) will be used.
        """
        self.imports = {
            "from typing import Optional, List, Dict, Any, Literal",
            "from enum import Enum",
            "from pipelex.core.stuffs.structured_content import StructuredContent",
            "from pydantic import Field",
        }
        self.enum_definitions: dict[str, dict[str, Any]] = {}  # Store enum definitions
        self.concept_ref_to_class_info = concept_ref_to_class_info or {}
        # Track concept classes that need to be mocked during validation (for file generation)
        self._concept_classes_to_mock: set[str] = set()

    def generate_from_structure_blueprint(
        self,
        class_name: str,
        structure_blueprint: dict[str, ConceptStructureBlueprint],
        base_class_name: str | None = None,
        description: str = "",
    ) -> tuple[str, type]:
        """Generate Python module content from structure blueprint.

        Args:
            class_name: Name of the class to generate
            structure_blueprint: Dictionary mapping field names to their ConceptStructureBlueprint definitions
            base_class_name: Optional base class name to inherit from (defaults to StructuredContent)
            description: Description for the class docstring (uses class_name if empty)

        Returns:
            Generated Python module content, the generated class

        Raises:
            ConceptStructureGeneratorError: If the generated code is not syntactically correct or does not inherit from the required base class

        """
        # Generate the class
        class_code = self._generate_class_source_code_from_blueprint(class_name, structure_blueprint, base_class_name, description)

        # Generate the complete module with a header comment
        imports_section = "\n".join(sorted(self.imports))

        autogen_header = (
            '"""\n'
            "AUTOGENERATED CODE - DO NOT EDIT\n"
            "\n"
            "If you want to customize this structure:\n"
            "  1. Copy this file to your own module\n"
            "  2. Remove the 'structure' or 'refines' declaration from the concept in the MTHDS file\n"
            "     and declare it in inline mode (see https://docs.pipelex.com/home/6-build-reliable-ai-workflows/concepts/define_your_concepts/#basic-concept-definition)\n"
            "  3. Make sure your custom class is importable and registered\n"
            "\n"
            "To regenerate: pipelex build structures <target_directory>\n"
            '"""\n\n'
        )

        generated_code = f"{autogen_header}{imports_section}\n\n\n{class_code}\n"

        # Validate the generated code
        try:
            the_class = self.validate_generated_code(
                python_code=generated_code,
                expected_class_name=class_name,
                base_class_name=base_class_name,
            )
        except SyntaxError as syntax_error:
            msg = f"Error validating generated code: {syntax_error}"
            syntax_error_data = SyntaxErrorData.from_syntax_error(syntax_error)
            raise ConceptStructureGeneratorError(
                msg, structure_class_python_code=generated_code, syntax_error_data=syntax_error_data
            ) from syntax_error
        except (ConceptStructureValidationError, ValueError, ImportError, Exception) as exc:
            msg = f"Error validating generated code: {exc}\nGenerated code:\n```python\n{generated_code}\n```"
            raise ConceptStructureGeneratorError(msg, structure_class_python_code=generated_code) from exc

        return generated_code, the_class

    def validate_generated_code(
        self,
        python_code: str,
        expected_class_name: str,
        base_class_name: str | None = None,
    ) -> type:
        """Validate that the generated Python code is syntactically correct and executable.

        Args:
            python_code: The generated Python code to validate
            expected_class_name: The name of the class that should be created
            base_class_name: Optional name of a custom base class to use

        """
        ast.parse(python_code)

        compile(python_code, "<generated>", "exec")

        return self._validate_execution(
            python_code=python_code,
            expected_class_name=expected_class_name,
            base_class_name=base_class_name,
        )

    ############################################################
    # Private methods
    ############################################################

    def _escape_string_for_python(self, value: str) -> str:
        """Escape a string value for safe inclusion in generated Python code.

        This method ensures that strings containing special characters like quotes,
        backslashes, newlines, etc. are properly escaped when inserted into generated
        Python source code.

        Args:
            value: The string to escape

        Returns:
            A properly escaped string representation suitable for Python code with double quotes
        """
        # Use repr() which handles all escaping correctly (backslashes, newlines, quotes, etc.)
        # repr() returns a string with surrounding quotes included
        escaped = repr(value)

        # If repr used single quotes, convert to double quotes
        # This is for consistency with existing code expectations
        if escaped.startswith("'") and escaped.endswith("'") and not escaped.startswith("'''"):
            # Single-line string with single quotes
            # Remove the single quotes and re-add as double quotes
            inner = escaped[1:-1]
            # Unescape single quotes that repr escaped (since we're switching to double quotes)
            inner = inner.replace("\\'", "'")
            # Escape any double quotes
            inner = inner.replace('"', '\\"')
            return f'"{inner}"'

        # If repr already used double quotes (because the string contains single quotes), return as is
        return escaped

    def _format_default_value(self, value: Any) -> str:
        """Format default value for Python code, ensuring strings use double quotes."""
        if isinstance(value, str):
            return self._escape_string_for_python(value)
        return repr(value)

    def _generate_class_source_code_from_blueprint(
        self,
        class_name: str,
        structure_blueprint: dict[str, ConceptStructureBlueprint],
        base_class_name: str | None = None,
        description: str = "",
    ) -> str:
        """Generate a class definition from ConceptStructureBlueprint.

        Args:
            class_name: Name of the class
            structure_blueprint: Dictionary mapping field names to their ConceptStructureBlueprint definitions
            base_class_name: Optional base class name to inherit from (defaults to StructuredContent)
            description: Description for the class docstring (uses class_name if empty)

        Returns:
            Generated class code with markers for regeneration

        """
        # Determine base class
        base_class = base_class_name or "StructuredContent"

        # Add import for the base class if it's a native content class (other than StructuredContent which is already imported)
        if NativeConceptCode.is_native_structure_class(base_class):
            cls = NativeConceptCode.get_native_structure_class(base_class)
            if not cls:
                msg = f"Native structure class '{base_class}' not found"
                raise ConceptStructureGeneratorError(msg)
            self.imports.add(f"from {cls.__module__} import {cls.__name__}")
        elif base_class != "StructuredContent":
            # Check if we have module info for this custom base class in concept_ref_to_class_info
            for class_info in self.concept_ref_to_class_info.values():
                if class_info.class_name == base_class and class_info.module_path:
                    self.imports.add(f"from {class_info.module_path} import {class_info.class_name}")
                    break

        # Generate class header with docstring (use class name if no description provided)
        docstring = description or f"Generated {class_name} class"
        class_header = f'class {class_name}({base_class}):\n    """{docstring}"""\n'

        # Generate fields
        field_definitions: list[str] = []
        for field_name, field_blueprint in structure_blueprint.items():
            field_code = self._generate_field_from_blueprint(field_name, field_blueprint)
            field_definitions.append(field_code)

        if not field_definitions:
            # Empty class - docstring is sufficient, no need for pass
            return class_header

        fields_code = "\n".join(field_definitions)
        return class_header + "\n" + fields_code

    def _generate_field_from_blueprint(self, field_name: str, field_blueprint: ConceptStructureBlueprint) -> str:
        """Generate a field definition from ConceptStructureBlueprint.

        Args:
            field_name: Name of the field
            field_blueprint: ConceptStructureBlueprint instance

        Returns:
            Generated field code

        """
        # Determine Python type
        if field_blueprint.choices:
            # Inline choices - use Literal type
            python_type = f"Literal[{', '.join(repr(c) for c in field_blueprint.choices)}]"
        else:
            # Handle complex types
            python_type = self._get_python_type_from_blueprint(field_blueprint)

        # Make optional if not required
        if not field_blueprint.required:
            python_type = f"Optional[{python_type}]"

        # Generate Field parameters
        field_params = [f"description={self._escape_string_for_python(field_blueprint.description)}"]

        if field_blueprint.required:
            if field_blueprint.default_value is not None:
                field_params.insert(0, f"default={self._format_default_value(field_blueprint.default_value)}")
            else:
                field_params.insert(0, "...")
        elif field_blueprint.default_value is not None:
            field_params.insert(0, f"default={self._format_default_value(field_blueprint.default_value)}")
        else:
            field_params.insert(0, "default=None")

        field_call = f"Field({', '.join(field_params)})"

        return f"    {field_name}: {python_type} = {field_call}"

    def _get_python_type_from_blueprint(self, field_blueprint: ConceptStructureBlueprint) -> str:
        """Convert ConceptStructureBlueprint to Python type annotation.

        Args:
            field_blueprint: ConceptStructureBlueprint instance

        Returns:
            Python type annotation string

        """
        if field_blueprint.type is None:
            # This should not happen based on validation, but handle gracefully
            return "str"

        # Use match/case for type handling
        match field_blueprint.type:
            case ConceptStructureBlueprintFieldType.TEXT:
                return "str"
            case ConceptStructureBlueprintFieldType.NUMBER:
                return "float"
            case ConceptStructureBlueprintFieldType.INTEGER:
                return "int"
            case ConceptStructureBlueprintFieldType.BOOLEAN:
                return "bool"
            case ConceptStructureBlueprintFieldType.DATE:
                self.imports.add("from datetime import datetime")
                return "datetime"
            case ConceptStructureBlueprintFieldType.CONCEPT:
                return self._resolve_concept_ref_to_type(field_blueprint.concept_ref)
            case ConceptStructureBlueprintFieldType.LIST:
                item_type = field_blueprint.item_type or "Any"
                # Handle list of concepts
                if item_type == "concept" and field_blueprint.item_concept_ref:
                    resolved_type = self._resolve_concept_ref_to_type(field_blueprint.item_concept_ref)
                    return f"List[{resolved_type}]"
                # Recursively handle item types if they're FieldType enums
                try:
                    item_type_enum = ConceptStructureBlueprintFieldType(item_type)
                    match item_type_enum:
                        case ConceptStructureBlueprintFieldType.DICT:
                            item_blueprint = ConceptStructureBlueprint(
                                description="lorem ipsum", type=item_type_enum, key_type="str", value_type="Any"
                            )
                            item_type = self._get_python_type_from_blueprint(item_blueprint)
                        case ConceptStructureBlueprintFieldType.CONCEPT:
                            # This case is handled above with item_concept_ref
                            pass
                        case (
                            ConceptStructureBlueprintFieldType.TEXT
                            | ConceptStructureBlueprintFieldType.NUMBER
                            | ConceptStructureBlueprintFieldType.INTEGER
                            | ConceptStructureBlueprintFieldType.BOOLEAN
                            | ConceptStructureBlueprintFieldType.DATE
                            | ConceptStructureBlueprintFieldType.LIST
                        ):
                            # Create a temporary blueprint for the item type
                            item_blueprint = ConceptStructureBlueprint(description="lorem ipsum", type=item_type_enum)
                            item_type = self._get_python_type_from_blueprint(item_blueprint)
                except ValueError:
                    # Keep as string if not a known FieldType
                    pass
                return f"List[{item_type}]"
            case ConceptStructureBlueprintFieldType.DICT:
                key_type = "str"
                value_type = field_blueprint.value_type or "Any"
                try:
                    value_type_enum = ConceptStructureBlueprintFieldType(value_type)
                    item_blueprint = ConceptStructureBlueprint(description="lorem ipsum", type=value_type_enum)
                    value_type = self._get_python_type_from_blueprint(item_blueprint)
                except ValueError:
                    pass
                return f"Dict[{key_type}, {value_type}]"

    def _resolve_concept_ref_to_type(self, concept_ref: str | None) -> str:
        """Resolve a concept reference to a Python type annotation.

        Args:
            concept_ref: The concept reference (e.g., "myapp.Customer" or "native.Html")

        Returns:
            The Python type annotation (structure class name).
            - For native concepts: generates import and returns class name (e.g., HtmlContent)
            - If module_path is available in concept_ref_to_class_info: generates an import,
              tracks for mocking, returns class name directly (for file generation)
            - Otherwise: returns a forward reference string (for runtime - resolved via model_rebuild)
        """
        if not concept_ref:
            return "Any"

        # Handle native concepts (e.g., "native.Html" -> HtmlContent)
        if NativeConceptCode.is_native_concept_ref_or_code(concept_ref):
            # Extract the concept code (e.g., "Html" from "native.Html")
            concept_code = extract_concept_code_from_concept_ref_or_code(concept_ref)
            try:
                native_code = NativeConceptCode(concept_code)
                structure_class = native_code.structure_class
                if structure_class:
                    self.imports.add(f"from {structure_class.__module__} import {structure_class.__name__}")
                    return structure_class.__name__
            except ValueError:
                pass

        # Check concept_ref_to_class_info for the concept reference
        if concept_ref in self.concept_ref_to_class_info:
            class_info = self.concept_ref_to_class_info[concept_ref]
            if class_info.module_path:
                # File generation: generate import and mock during validation
                self.imports.add(f"from {class_info.module_path} import {class_info.class_name}")
                self._concept_classes_to_mock.add(class_info.class_name)
                return class_info.class_name
            # No module path: use forward reference (resolved via model_rebuild at runtime)
            return f'"{class_info.class_name}"'

        # Default: extract concept code and use as forward reference
        # e.g., "myapp.Customer" -> '"Customer"'
        concept_code = extract_concept_code_from_concept_ref_or_code(concept_ref)
        return f'"{concept_code}"'

    def _generate_field(self, field_name: str, field_def: dict[str, Any] | str) -> str:
        """Generate a single field definition.

        Args:
            field_name: Name of the field
            field_def: Field definition (dict or string for simple types)

        Returns:
            Generated field code

        """
        # Handle simple string definitions (just the definition text)
        if isinstance(field_def, str):
            field_def = {"type": ConceptStructureBlueprintFieldType.TEXT, "description": field_def}

        field_type = field_def.get("type", ConceptStructureBlueprintFieldType.TEXT)
        description = field_def.get("description", f"{field_name} field")
        required = field_def.get("required", False)
        default_value = field_def.get("default")
        choices = field_def.get("choices")  # For inline enum-like choices

        # Determine Python type
        if choices:
            # Inline choices - use Literal type
            python_type = f"Literal[{', '.join(repr(c) for c in choices)}]"
        else:
            # Handle complex types or enum references
            python_type = self._get_python_type(field_type, field_def)

        # Make optional if not required
        if not required:
            python_type = f"Optional[{python_type}]"

        # Generate Field parameters
        field_params = [f"description={self._escape_string_for_python(description)}"]

        if required:
            if default_value is not None:
                field_params.insert(0, f"default={self._format_default_value(default_value)}")
            else:
                field_params.insert(0, "...")
        elif default_value is not None:
            field_params.insert(0, f"default={self._format_default_value(default_value)}")
        else:
            field_params.insert(0, "default=None")

        field_call = f"Field({', '.join(field_params)})"

        return f"    {field_name}: {python_type} = {field_call}"

    def _get_python_type(self, field_type: Any, field_def: dict[str, Any]) -> str:
        """Convert high-level type to Python type annotation.

        Args:
            field_type: High-level type name or FieldType enum
            field_def: Complete field definition

        Returns:
            Python type annotation string

        """
        # Check if it's a reference to a defined enum
        if isinstance(field_type, str) and field_type in self.enum_definitions:
            return field_type

        # Convert string to FieldType if needed
        if isinstance(field_type, str):
            try:
                field_type_enum = ConceptStructureBlueprintFieldType(field_type)
            except ValueError:
                # Unknown type, assume it's a custom type or class reference
                return field_type
            field_type = field_type_enum

        # Use match/case for type handling
        # Note: field_type is already a ConceptStructureBlueprintFieldType at this point
        # (converted from string above if needed)
        if not isinstance(field_type, ConceptStructureBlueprintFieldType):
            return str(field_type)

        match field_type:
            case ConceptStructureBlueprintFieldType.TEXT:
                return "str"
            case ConceptStructureBlueprintFieldType.NUMBER:
                return "float"
            case ConceptStructureBlueprintFieldType.INTEGER:
                return "int"
            case ConceptStructureBlueprintFieldType.BOOLEAN:
                return "bool"
            case ConceptStructureBlueprintFieldType.DATE:
                self.imports.add("from datetime import datetime")
                return "datetime"
            case ConceptStructureBlueprintFieldType.CONCEPT:
                # Handle concept type - resolve using concept_ref from field_def
                concept_ref = field_def.get("concept_ref")
                return self._resolve_concept_ref_to_type(concept_ref)
            case ConceptStructureBlueprintFieldType.LIST:
                item_type = field_def.get("item_type", "Any")
                # Handle list of concepts
                if item_type == "concept":
                    item_concept_ref = field_def.get("item_concept_ref")
                    if item_concept_ref:
                        resolved_type = self._resolve_concept_ref_to_type(item_concept_ref)
                        return f"List[{resolved_type}]"
                # Check if item_type is an enum reference
                if isinstance(item_type, str) and item_type in self.enum_definitions:
                    return f"List[{item_type}]"
                # Recursively handle item types
                if isinstance(item_type, str):
                    try:
                        item_type_enum = ConceptStructureBlueprintFieldType(item_type)
                        item_type = self._get_python_type(item_type_enum, {})
                    except ValueError:
                        # Keep as string if not a known FieldType
                        pass
                return f"List[{item_type}]"
            case ConceptStructureBlueprintFieldType.DICT:
                key_type = field_def.get("key_type", "str")
                value_type = field_def.get("value_type", "Any")
                # Recursively handle key and value types
                if isinstance(key_type, str):
                    try:
                        key_type_enum = ConceptStructureBlueprintFieldType(key_type)
                        key_type = self._get_python_type(key_type_enum, {})
                    except ValueError:
                        pass
                if isinstance(value_type, str):
                    try:
                        value_type_enum = ConceptStructureBlueprintFieldType(value_type)
                        value_type = self._get_python_type(value_type_enum, {})
                    except ValueError:
                        pass
                return f"Dict[{key_type}, {value_type}]"

        # Fallback for unknown types
        return str(field_type)

    def _validate_execution(
        self,
        python_code: str,
        expected_class_name: str,
        base_class_name: str | None = None,
    ) -> type:
        """Validate that the code executes and creates the expected class."""
        # exec_globals with basic types but let native class imports execute naturally
        # This ensures we use the same class objects for inheritance checks
        exec_globals: dict[str, Any] = {
            "__builtins__": __builtins__,
            # Provide these directly as they're used in generated code but not imported from pipelex
            "datetime": datetime,
            "Enum": Enum,
            "Optional": Optional,
            "List": List,  # noqa: UP006
            "Dict": dict,
            "Any": Any,
            "Literal": Literal,
            "Field": Field,
        }

        # Add all native content classes to exec_globals so they're available during validation
        # This is needed because exec() with separate globals/locals puts imports in locals,
        # but class annotations look in globals
        for native_code in NativeConceptCode:
            structure_class = native_code.structure_class
            if structure_class:
                exec_globals[structure_class.__name__] = structure_class

        validation_code = python_code

        # For file generation: create mock classes for concept references that would fail import
        # These classes don't exist yet or their modules aren't in the path
        for mock_class_name in self._concept_classes_to_mock:
            # Create a mock class that inherits from StructuredContent
            mock_class = type(mock_class_name, (StructuredContent,), {})
            exec_globals[mock_class_name] = mock_class
            # Remove the import line for this class from the validation code
            # (the actual generated code will still have it)
            lines = validation_code.split("\n")
            filtered_lines = [
                line for line in lines if not (line.strip().startswith("from ") and f"import {mock_class_name}" in line and "pipelex" not in line)
            ]
            validation_code = "\n".join(filtered_lines)

        # If a custom base class is specified that's not a native class, get it from the registry
        if base_class_name:
            if not NativeConceptCode.is_native_structure_class(base_class_name):
                # Not a native class, provide it from registry
                custom_base_class = KajsonManager.get_class_registry().get_class(name=base_class_name)
                if custom_base_class is None:
                    msg = f"Base class '{base_class_name}' not found in native classes or class registry"
                    raise ConceptStructureValidationError(msg)
                exec_globals[base_class_name] = custom_base_class  # type: ignore[assignment]

        exec_locals: dict[str, Any] = {}
        exec(validation_code, exec_globals, exec_locals)

        # Verify the expected class was created
        if expected_class_name not in exec_locals:
            msg = f"Expected class '{expected_class_name}' not found in generated code"
            raise ConceptStructureValidationError(msg)

        the_class = exec_locals[expected_class_name]

        # Verify it's actually a class
        if not isinstance(the_class, type):
            msg = f"'{expected_class_name}' is not a class"
            raise ConceptStructureValidationError(msg)

        # Verify it inherits from the required base class
        # Check the MRO (Method Resolution Order) by class name to handle cases
        # where the class object references might differ due to exec() context
        mro_class_names = [cls.__name__ for cls in the_class.__mro__]

        if StuffContent.__name__ not in mro_class_names:
            msg = f"'{expected_class_name}' does not inherit from StuffContent. MRO: {mro_class_names}"
            raise ConceptStructureValidationError(msg)

        return the_class
