"""
Schema Loader

Loads Python dataclasses from a schema/ folder.
Expects the folder to have an __init__.py that exports:
- DATACLASSES: List of dataclass types to generate
- ENUMS: List of Enum types to include
- GRPC_MESSAGES: (optional) List of gRPC message types
- GRPC_SERVICE: (optional) gRPC service definition dict
"""

import dataclasses as dc
import importlib.util
import sys
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Type, get_type_hints


@dataclass
class FieldDefinition:
    """Represents a model field."""

    name: str
    type_hint: Any
    default: Any = dc.MISSING
    optional: bool = False


@dataclass
class ModelDefinition:
    """Represents a model/dataclass."""

    name: str
    fields: List[FieldDefinition]
    docstring: Optional[str] = None


@dataclass
class EnumDefinition:
    """Represents an enum."""

    name: str
    values: List[tuple[str, str]]  # (name, value) pairs


@dataclass
class GrpcServiceDefinition:
    """Represents a gRPC service."""

    package: str
    name: str
    methods: List[Dict[str, Any]]


class SchemaLoader:
    """Loads model definitions from Python dataclasses in schema/ folder."""

    def __init__(self, schema_path: Path):
        self.schema_path = Path(schema_path)
        self.models: List[ModelDefinition] = []
        self.api_models: List[ModelDefinition] = []
        self.enums: List[EnumDefinition] = []
        self.grpc_messages: List[ModelDefinition] = []
        self.grpc_service: Optional[GrpcServiceDefinition] = None

    def load(self, include: Optional[set] = None) -> "SchemaLoader":
        """Load schema definitions from the schema folder.

        Args:
            include: Set of groups to load (dataclasses, enums, api, grpc).
                     None means load all groups.
        """
        init_path = self.schema_path / "__init__.py"

        if not init_path.exists():
            raise FileNotFoundError(f"Schema folder must have __init__.py: {init_path}")

        # Import the schema module
        module = self._import_module(init_path)

        load_all = include is None

        # Extract DATACLASSES
        if load_all or "dataclasses" in include:
            dataclasses = getattr(module, "DATACLASSES", [])
            for cls in dataclasses:
                self.models.append(self._parse_dataclass(cls))

        # Extract API_MODELS (TypeScript-only request/response types)
        if load_all or "api" in include:
            api_models = getattr(module, "API_MODELS", [])
            for cls in api_models:
                self.api_models.append(self._parse_dataclass(cls))

        # Extract ENUMS
        if load_all or "enums" in include:
            enums = getattr(module, "ENUMS", [])
            for enum_cls in enums:
                self.enums.append(self._parse_enum(enum_cls))

        # Extract GRPC_MESSAGES (optional)
        if load_all or "grpc" in include:
            grpc_messages = getattr(module, "GRPC_MESSAGES", [])
            for cls in grpc_messages:
                self.grpc_messages.append(self._parse_dataclass(cls))

        # Extract GRPC_SERVICE (optional)
        if load_all or "grpc" in include:
            grpc_service = getattr(module, "GRPC_SERVICE", None)
            if grpc_service:
                self.grpc_service = GrpcServiceDefinition(
                    package=grpc_service.get("package", "service"),
                    name=grpc_service.get("name", "Service"),
                    methods=grpc_service.get("methods", []),
                )

        return self

    def _import_module(self, path: Path):
        """Import a Python module from a file path."""
        spec = importlib.util.spec_from_file_location("schema", path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load module from {path}")

        module = importlib.util.module_from_spec(spec)
        sys.modules["schema"] = module
        spec.loader.exec_module(module)
        return module

    def _parse_dataclass(self, cls: Type) -> ModelDefinition:
        """Parse a dataclass into a ModelDefinition."""
        hints = get_type_hints(cls)
        fields_info = {f.name: f for f in dc.fields(cls)}

        fields = []
        for name, type_hint in hints.items():
            if name.startswith("_"):
                continue

            field_info = fields_info.get(name)
            default = dc.MISSING
            if field_info:
                if field_info.default is not dc.MISSING:
                    default = field_info.default
                elif field_info.default_factory is not dc.MISSING:
                    default = field_info.default_factory

            # Check if optional (Union with None)
            optional = self._is_optional(type_hint)

            fields.append(
                FieldDefinition(
                    name=name,
                    type_hint=type_hint,
                    default=default,
                    optional=optional,
                )
            )

        return ModelDefinition(
            name=cls.__name__,
            fields=fields,
            docstring=cls.__doc__,
        )

    def _parse_enum(self, enum_cls: Type[Enum]) -> EnumDefinition:
        """Parse an Enum into an EnumDefinition."""
        values = [(m.name, m.value) for m in enum_cls]
        return EnumDefinition(name=enum_cls.__name__, values=values)

    def _is_optional(self, type_hint: Any) -> bool:
        """Check if a type hint is Optional (Union with None)."""
        from typing import Union, get_args, get_origin

        origin = get_origin(type_hint)
        if origin is Union:
            args = get_args(type_hint)
            return type(None) in args
        return False


def load_schema(schema_path: str | Path, include: Optional[set] = None) -> SchemaLoader:
    """Load schema definitions from folder."""
    loader = SchemaLoader(schema_path)
    return loader.load(include=include)