"""gRPC generator for Prism.

Generates proto files and Python gRPC service stubs.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

from prisme.generators.base import GeneratedFile, ModelGenerator, create_init_file
from prisme.spec.fields import FieldType
from prisme.spec.stack import FileStrategy
from prisme.utils.case_conversion import to_snake_case
from prisme.utils.template_engine import TemplateRenderer

if TYPE_CHECKING:
    from prisme.spec.fields import FieldSpec
    from prisme.spec.model import ModelSpec

# Map Prism FieldType to proto3 types
PROTO_TYPE_MAP: dict[FieldType, str] = {
    FieldType.STRING: "string",
    FieldType.TEXT: "string",
    FieldType.INTEGER: "int64",
    FieldType.FLOAT: "double",
    FieldType.DECIMAL: "string",  # Decimal as string for precision
    FieldType.BOOLEAN: "bool",
    FieldType.DATETIME: "google.protobuf.Timestamp",
    FieldType.DATE: "string",  # ISO date string
    FieldType.TIME: "string",  # ISO time string
    FieldType.UUID: "string",
    FieldType.JSON: "google.protobuf.Struct",
    FieldType.ENUM: "string",
    FieldType.FOREIGN_KEY: "int64",
}


class GRPCGenerator(ModelGenerator):
    """Generator for gRPC proto files and service implementations."""

    REQUIRED_TEMPLATES = [
        "backend/grpc/proto.jinja2",
        "backend/grpc/server.py.jinja2",
        "backend/grpc/service.py.jinja2",
    ]

    def __init__(self, *args, **kwargs) -> None:  # type: ignore[no-untyped-def]
        super().__init__(*args, **kwargs)
        backend_base = Path(self.generator_config.backend_output)
        package_name = self.get_package_name()
        package_base = backend_base / package_name
        self.grpc_path = package_base / self.generator_config.grpc_path
        self.generated_path = package_base / self.generator_config.grpc_generated_path
        self.renderer = TemplateRenderer()
        self.renderer.validate_templates_exist(self.REQUIRED_TEMPLATES)

    def generate_shared_files(self) -> list[GeneratedFile]:
        """Generate shared gRPC server setup."""
        return [
            self._generate_server(),
            self._generate_generated_init(),
        ]

    def generate_model_files(self, model: ModelSpec) -> list[GeneratedFile]:
        """Generate gRPC files for a single model."""
        if not model.expose:
            return []

        return [
            self._generate_proto(model),
            self._generate_service(model),
        ]

    def generate_index_files(self) -> list[GeneratedFile]:
        """Generate __init__.py for gRPC module."""
        imports = ["from ._generated.server import serve, GRPC_PORT, GRPC_REFLECTION"]
        exports = ["serve", "GRPC_PORT", "GRPC_REFLECTION"]

        for model in self.spec.models:
            if model.expose:
                snake_name = to_snake_case(model.name)
                imports.append(f"from ._generated.{snake_name}_service import {model.name}Servicer")
                exports.append(f"{model.name}Servicer")

        return [
            create_init_file(
                self.grpc_path,
                imports,
                exports,
                "gRPC service definitions.",
            ),
        ]

    def _generate_server(self) -> GeneratedFile:
        """Generate the main gRPC server module."""
        enabled_models = [m for m in self.spec.models if m.expose]
        package_name = self.get_package_name()

        grpc_config = self.exposure_config.grpc if self.exposure_config else None
        grpc_port = grpc_config.port if grpc_config else 50051
        grpc_reflection = grpc_config.reflection if grpc_config else True

        # Build registration block
        register_lines = []
        service_name_lines = []
        for model in enabled_models:
            snake_name = to_snake_case(model.name)
            service_name = f"{model.name}Service"
            register_lines.append(
                f"    # TODO: Register {model.name}Servicer with generated pb2_grpc"
            )
            register_lines.append(
                f"    # {snake_name}_pb2_grpc.add_{service_name}Servicer_to_server("
            )
            register_lines.append(f"    #     {model.name}Servicer(), server")
            register_lines.append("    # )")
            service_name_lines.append(f'            "{package_name}.{service_name}",')

        register_block = "\n".join(register_lines) if register_lines else "    pass"
        service_names_block = "\n".join(service_name_lines)

        content = self.renderer.render_file(
            "backend/grpc/server.py.jinja2",
            context={
                "project_title": self.spec.effective_title,
                "grpc_port": grpc_port,
                "grpc_reflection": grpc_reflection,
                "register_block": register_block,
                "service_names_block": service_names_block,
            },
        )
        return GeneratedFile(
            path=self.generated_path / "server.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="gRPC server setup",
        )

    def _generate_generated_init(self) -> GeneratedFile:
        """Generate __init__.py for _generated folder."""
        imports = ["from .server import serve, GRPC_PORT, GRPC_REFLECTION"]
        exports = ["serve", "GRPC_PORT", "GRPC_REFLECTION"]

        for model in self.spec.models:
            if model.expose:
                snake_name = to_snake_case(model.name)
                imports.append(f"from .{snake_name}_service import {model.name}Servicer")
                exports.append(f"{model.name}Servicer")

        return create_init_file(
            self.generated_path,
            imports,
            exports,
            "Generated gRPC service components.",
        )

    def _generate_proto(self, model: ModelSpec) -> GeneratedFile:
        """Generate proto file for a model."""
        snake_name = to_snake_case(model.name)
        package_name = self.get_package_name()
        service_name = f"{model.name}Service"

        # Build proto fields
        proto_fields = []
        field_number = 2  # Start at 2, id is 1
        for field in model.fields:
            proto_type = self._get_proto_type(field)
            proto_fields.append(
                {
                    "name": field.name,
                    "proto_type": proto_type,
                    "number": field_number,
                }
            )
            field_number += 1

        # Timestamps
        ts_created_number = field_number
        ts_updated_number = field_number + 1

        # Create fields (same as proto_fields, just re-numbered from 1)
        create_fields = []
        for i, field in enumerate(model.fields, start=1):
            proto_type = self._get_proto_type(field)
            if not field.required:
                proto_type = f"optional {proto_type}"
            create_fields.append(
                {
                    "name": field.name,
                    "proto_type": proto_type,
                    "number": i,
                }
            )

        # Update fields (all optional, numbered from 2 since id=1)
        update_fields = []
        for i, field in enumerate(model.fields, start=2):
            proto_type = self._get_proto_type(field)
            update_fields.append(
                {
                    "name": field.name,
                    "proto_type": f"optional {proto_type}",
                    "number": i,
                }
            )

        # gRPC exposure config
        streaming_methods = True
        bidirectional_streaming = False

        content = self.renderer.render_file(
            "backend/grpc/proto.jinja2",
            context={
                "model_name": model.name,
                "snake_name": snake_name,
                "proto_package": package_name.replace("_", "."),
                "python_package": package_name,
                "service_name": service_name,
                "proto_fields": proto_fields,
                "create_fields": create_fields,
                "update_fields": update_fields,
                "has_timestamps": model.timestamps,
                "ts_created_number": ts_created_number,
                "ts_updated_number": ts_updated_number,
                "ops_read": model.has_operation("read"),
                "ops_list": model.has_operation("list"),
                "ops_create": model.has_operation("create"),
                "ops_update": model.has_operation("update"),
                "ops_delete": model.has_operation("delete"),
                "streaming_methods": streaming_methods,
                "bidirectional_streaming": bidirectional_streaming,
            },
        )
        return GeneratedFile(
            path=self.generated_path / f"{snake_name}.proto",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description=f"gRPC proto for {model.name}",
        )

    def _generate_service(self, model: ModelSpec) -> GeneratedFile:
        """Generate gRPC service implementation for a model."""
        snake_name = to_snake_case(model.name)
        package_name = self.get_package_name()
        service_name = f"{model.name}Service"

        streaming_methods = True
        bidirectional_streaming = False

        content = self.renderer.render_file(
            "backend/grpc/service.py.jinja2",
            context={
                "model_name": model.name,
                "snake_name": snake_name,
                "package_name": package_name,
                "service_name": service_name,
                "ops_read": model.has_operation("read"),
                "ops_list": model.has_operation("list"),
                "ops_create": model.has_operation("create"),
                "ops_update": model.has_operation("update"),
                "ops_delete": model.has_operation("delete"),
                "streaming_methods": streaming_methods,
                "bidirectional_streaming": bidirectional_streaming,
            },
        )
        return GeneratedFile(
            path=self.generated_path / f"{snake_name}_service.py",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description=f"gRPC service for {model.name}",
        )

    def _get_proto_type(self, field: FieldSpec) -> str:
        """Get the proto3 type for a field."""
        if field.type == FieldType.JSON and field.json_item_type:
            item_map = {
                "str": "string",
                "string": "string",
                "int": "int64",
                "integer": "int64",
                "float": "double",
                "number": "double",
                "bool": "bool",
                "boolean": "bool",
            }
            item_type = item_map.get(field.json_item_type, "string")
            return f"repeated {item_type}"
        return PROTO_TYPE_MAP.get(field.type, "string")


__all__ = ["GRPCGenerator"]
