import os
import re
import shutil


class RuntimePackager:
    """ Packages the GnuCOBOL compilation script, Dockerfile, and ctypes shim. """

    COMPILE_SH = """#!/bin/bash
# Auto-generated by cobol-wrap
COBOL_SRC=$1
if [ -z "$COBOL_SRC" ]; then
    echo "Usage: $0 <source.cbl>"
    exit 1
fi
NAME=$(basename -s .cbl "$COBOL_SRC")

# Generate stub C functions for SQL and CICS bridge intercepts.
# GnuCOBOL translates hyphens in CALL names to underscores for C symbol resolution.
cat << 'EOF' > bridges.c
void SQL_BRIDGE() {}
void CICS_BRIDGE() {}
EOF

# -F = free source format (handles long PROCEDURE DIVISION USING lines)
# -m = build a dynamically loadable module
cobc -F -m "$COBOL_SRC" bridges.c
STATUS=$?
# GnuCOBOL outputs .dylib on macOS, .so on Linux
LIB=$(ls "${NAME}.dylib" "${NAME}.so" 2>/dev/null | head -1)
if [ $STATUS -eq 0 ] && [ -n "$LIB" ]; then
    echo "Compiled successfully: $LIB"
else
    echo "Compilation failed."
    exit 1
fi
"""

    DOCKERFILE = """FROM python:3.11-slim
RUN apt-get update && apt-get install -y gnucobol && rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
RUN bash runtime/compile.sh runtime/*.cbl
EXPOSE 8080
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8080"]
"""

    @classmethod
    def _docker_compose(cls, kafka: bool = False, telemetry: bool = False) -> str:
        lines = [
            "version: '3.8'",
            "services:",
            "  api:",
            "    build: .",
            "    ports:",
            '      - "8080:8080"',
            "    volumes:",
            "      - ./runtime:/app/runtime",
        ]
        if kafka or telemetry:
            lines.append("    depends_on:")
            if kafka:
                lines.append("      - kafka")
            if telemetry:
                lines.append("      - otel-collector")
            lines.append("    environment:")
            if kafka:
                lines.append("      - KAFKA_BOOTSTRAP_SERVERS=kafka:9092")
            if telemetry:
                lines.append("      - OTEL_EXPORTER_OTLP_ENDPOINT=http://otel-collector:4318")
        if kafka:
            lines.extend([
                "",
                "  zookeeper:",
                "    image: confluentinc/cp-zookeeper:7.5.0",
                "    environment:",
                "      ZOOKEEPER_CLIENT_PORT: 2181",
                "",
                "  kafka:",
                "    image: confluentinc/cp-kafka:7.5.0",
                "    depends_on:",
                "      - zookeeper",
                "    ports:",
                '      - "9092:9092"',
                "    environment:",
                "      KAFKA_BROKER_ID: 1",
                "      KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181",
                "      KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:9092",
                "      KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1",
            ])
        if telemetry:
            lines.extend([
                "",
                "  otel-collector:",
                "    image: otel/opentelemetry-collector:latest",
                "    ports:",
                '      - "4318:4318"',
                '      - "55679:55679"',
            ])
        lines.append("")
        return "\n".join(lines)

    @classmethod
    def emit(cls, output_dir: str, program, docker: bool = True, source_path: str = None,
             graphql: bool = False, kafka: bool = False, telemetry: bool = False):
        rt_dir = os.path.join(output_dir, "runtime")
        os.makedirs(rt_dir, exist_ok=True)

        # Copy the COBOL source file into the runtime directory so compile.sh can find it
        if source_path and os.path.exists(source_path):
            dest = os.path.join(rt_dir, f"{program.program_id.lower()}.cbl")
            shutil.copy2(source_path, dest)

        # Write compile.sh
        compile_sh = os.path.join(rt_dir, "compile.sh")
        with open(compile_sh, "w") as f:
            f.write(cls.COMPILE_SH)
        os.chmod(compile_sh, 0o755)

        # Always generate requirements.txt (needed for pip install, not just Docker)
        with open(os.path.join(output_dir, "requirements.txt"), "w") as f:
            deps = ["fastapi", "uvicorn", "pydantic"]
            if graphql:
                deps.append("strawberry-graphql")
            if kafka:
                deps.append("faststream[kafka]")
            if telemetry:
                deps.extend([
                    "opentelemetry-api",
                    "opentelemetry-sdk",
                    "opentelemetry-instrumentation-fastapi",
                    "opentelemetry-exporter-otlp-proto-http",
                ])
            f.write("\n".join(deps) + "\n")

        if docker:
            with open(os.path.join(output_dir, "Dockerfile"), "w") as f:
                f.write(cls.DOCKERFILE)
            with open(os.path.join(output_dir, "docker-compose.yml"), "w") as f:
                f.write(cls._docker_compose(kafka=kafka, telemetry=telemetry))

        # Generate ctypes shim
        shim_code = cls._generate_shim(program)
        with open(os.path.join(rt_dir, "shim.py"), "w") as f:
            f.write(shim_code)

    @classmethod
    def _get_ctype(cls, field) -> str:
        """Map a DataField PIC clause to a ctypes C type string."""
        if getattr(field, "is_pointer", False):
            return "ctypes.c_void_p"
        if not field.pic_clause:
            return "ctypes.c_char * 100"

        pic = field.pic_clause.upper().replace(" ", "")

        if 'X' in pic or 'A' in pic:
            match = re.search(r'[XA]\((\d+)\)', pic)
            if match:
                length = int(match.group(1))
            else:
                # Bare repeated chars: PIC XX = 2, PIC XXX = 3, etc.
                length = max(pic.count('X'), pic.count('A'))
            return f"ctypes.c_char * {length}"

        if '9' in pic:
            # Count total digits (integer + decimal)
            int_match = re.search(r'9\((\d+)\)', pic)
            digits = int(int_match.group(1)) if int_match else pic.count('9')
            v_match = re.search(r'V(?:9\((\d+)\)|(9+))', pic)
            if v_match:
                if v_match.group(1):
                    digits += int(v_match.group(1))
                else:
                    digits += len(v_match.group(2))

            if field.is_comp3:
                # COMP-3 (packed BCD): each digit = half byte, plus sign nibble
                byte_len = (digits + 2) // 2
                return f"ctypes.c_char * {byte_len}"

            if field.is_comp:
                # COMP/BINARY: GnuCOBOL size depends on digit count
                is_signed = 'S' in pic
                if digits <= 4:
                    return "ctypes.c_int16" if is_signed else "ctypes.c_uint16"
                elif digits <= 9:
                    return "ctypes.c_int32" if is_signed else "ctypes.c_uint32"
                else:
                    return "ctypes.c_int64" if is_signed else "ctypes.c_uint64"

            return f"ctypes.c_char * {digits}"

        return "ctypes.c_char * 10"

    @classmethod
    def _generate_shim(cls, program) -> str:
        """
        Translates the COBOL AST Working-Storage into Python ctypes.Structure classes.
        """
        from .ast import DataField

        lines = [
            "import ctypes",
            "import ctypes.util",
            "import os",
            "import sys",
            "import json",
            "",
            "# --- GnuCOBOL runtime init (must happen before loading any COBOL module) ---",
            "_libcob_name = ctypes.util.find_library('cob')",
            "if _libcob_name:",
            "    _libcob = ctypes.CDLL(_libcob_name, mode=ctypes.RTLD_GLOBAL)",
            "    _libcob.cob_init(0, None)",
            "",
            "# Platform-specific shared library extension",
            "_LIB_EXT = '.dylib' if sys.platform == 'darwin' else '.so'",
            "",
            "# --- COBOL Working Storage C-Type Maps ---",
        ]

        def _traverse_structure(field: DataField, parent_name: str):
            struct_name = f"{parent_name}_{field.name.replace('-', '_')}_Struct"
            if field.children:
                for child in field.children:
                    _traverse_structure(child, struct_name)

                lines.append(f"class {struct_name}(ctypes.Structure):")
                lines.append("    _fields_ = [")
                for child in field.children:
                    if child.children:
                        c_type = f"{struct_name}_{child.name.replace('-', '_')}_Struct"
                    else:
                        c_type = cls._get_ctype(child)
                    # OCCURS N TIMES → ctypes array (Type * N)
                    if getattr(child, 'occurs', None) and child.occurs > 1:
                        c_type = f"{c_type} * {child.occurs}"
                    lines.append(f"        ('{child.name.replace('-', '_')}', {c_type}),")
                lines.append("    ]")
                lines.append("")

        for section in [program.working_storage, program.linkage_section]:
            for root_field in section:
                if getattr(root_field, 'level', 0) in (1, 77):
                    _traverse_structure(root_field, "Root")
                    if not root_field.children:
                        struct_name = f"Root_{root_field.name.replace('-', '_')}_Struct"
                        lines.append(f"class {struct_name}(ctypes.Structure):")
                        lines.append("    _fields_ = [")
                        lines.append(f"        ('{root_field.name.replace('-', '_')}', {cls._get_ctype(root_field)})")
                        lines.append("    ]")
                        lines.append("")

        lines.extend([
            "class CobolRuntime:",
            "    def __init__(self, so_path: str):",
            "        # Accept .so path but also try platform-native extension (.dylib on macOS)",
            "        if not os.path.exists(so_path):",
            "            alt = os.path.splitext(so_path)[0] + _LIB_EXT",
            "            if os.path.exists(alt):",
            "                so_path = alt",
            "        if not os.path.exists(so_path):",
            "            raise RuntimeError(f'Shared library not found: {so_path}. Run compile.sh first.')",
            "        self.lib = ctypes.CDLL(so_path)",
            "",
            "    def _dict_to_struct(self, d: dict, struct_obj):",
            "        for field_name, field_type in struct_obj._fields_:",
            "            matched_val = None",
            "            for k, v in d.items():",
            "                if k.replace('_', '').lower() == field_name.replace('_', '').lower():",
            "                    matched_val = v; break",
            "            if matched_val is not None:",
            "                if isinstance(matched_val, dict):",
            "                    self._dict_to_struct(matched_val, getattr(struct_obj, field_name))",
            "                else:",
            "                    try:",
            "                        if hasattr(field_type, '_length_'):",
            "                            val_str = str(matched_val) if not isinstance(matched_val, str) else matched_val",
            "                            # Convert decimal values to COBOL display format (remove decimal point)",
            "                            if '.' in val_str:",
            "                                parts = val_str.split('.')",
            "                                val_str = parts[0] + parts[1]",
            "                            setattr(struct_obj, field_name, val_str.rjust(field_type._length_, '0').encode('utf-8'))",
            "                        elif field_type in (ctypes.c_int16, ctypes.c_uint16, ctypes.c_int32, ctypes.c_uint32, ctypes.c_int64, ctypes.c_uint64):",
            "                            setattr(struct_obj, field_name, int(matched_val))",
            "                        else:",
            "                            setattr(struct_obj, field_name, matched_val)",
            "                    except (ValueError, TypeError, AttributeError) as exc:",
            "                        import logging",
            "                        logging.getLogger('cobol_wrap.shim').warning(",
            "                            'Could not map field %s: %s', field_name, exc)",
            "",
            "    def _struct_to_dict(self, struct_obj) -> dict:",
            "        result = {}",
            "        for field_name, field_type in struct_obj._fields_:",
            "            val = getattr(struct_obj, field_name)",
            "            if hasattr(val, '_fields_'):",
            "                result[field_name.lower()] = self._struct_to_dict(val)",
            "            elif isinstance(val, bytes):",
            "                result[field_name.lower()] = val.decode('utf-8', errors='ignore').strip()",
            "            else:",
            "                result[field_name.lower()] = val",
            "        return result",
            "",
            "    def call(self, entry_point: str, payload: dict) -> dict:",
            "        c_func = getattr(self.lib, entry_point.lower().replace('-', '_'), None)",
            "        if not c_func:",
            "             c_func = getattr(self.lib, entry_point, None)",
            "        if not c_func:",
            "             raise ValueError(f'Entry point {entry_point} not found in binary.')",
            "        result = {}"
        ])

        for ep in program.procedure_division:
            lines.append(f"        if entry_point.upper() == '{ep.name.upper()}':")
            arg_names = []
            for param in ep.params:
                struct_name = f"Root_{param.replace('-', '_')}_Struct"
                var_name = f"arg_{param.replace('-', '_')}"
                lines.append(f"            {var_name} = {struct_name}()")
                lines.append("            if payload:")
                lines.append(f"                self._dict_to_struct(payload, {var_name})")
                arg_names.append(f"ctypes.byref({var_name})")

            lines.append(f"            c_func({', '.join(arg_names) if arg_names else ''})")
            for param in ep.params:
                var_name = f"arg_{param.replace('-', '_')}"
                lines.append(f"            result.update(self._struct_to_dict({var_name}))")

            lines.append("            return result")

        lines.append("        return {'status': 'success', 'note': f'Executed {entry_point} with generic payload but no mapped parameters.'}")
        return "\n".join(lines)
