"""OTel instrumentation code generator.

Generates ready-to-commit OTel configuration files for the detected stack.
Output follows OTel SDK best practices for ASGI/FastAPI setups.
Switchable to Jaeger / Grafana Cloud / Honeycomb / stdout exporter.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

_OTEL_SETUP_PY = """\
# Generated by ace_instrument — add to your main.py startup.
# Required packages: opentelemetry-sdk opentelemetry-exporter-otlp
# pip install opentelemetry-sdk opentelemetry-exporter-otlp-proto-grpc
#             opentelemetry-instrumentation-fastapi opentelemetry-instrumentation-httpx
import os
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.resources import Resource, SERVICE_NAME
from opentelemetry.sdk.trace.export import BatchSpanProcessor
{exporter_import}
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor


def setup_otel(app=None):
    \"\"\"Initialize OTel tracing. Call this at app startup.\"\"\"
    resource = Resource(attributes={{SERVICE_NAME: "{service_name}"}})
    provider = TracerProvider(resource=resource)
    exporter = {exporter_class}(
        endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "{default_endpoint}"),
    )
    provider.add_span_processor(BatchSpanProcessor(exporter))
    trace.set_tracer_provider(provider)
    if app is not None:
        FastAPIInstrumentor.instrument_app(app)
    HTTPXClientInstrumentor().instrument()
    return app
"""

_LOGGING_CONFIG_PY = """\
# Generated by ace_instrument — structured logging with OTel trace correlation.
import json
import logging
from opentelemetry import trace


class OTelFormatter(logging.Formatter):
    \"\"\"Injects trace_id + span_id into every log line as JSON.\"\"\"

    def format(self, record):
        ctx = trace.get_current_span().get_span_context()
        record.trace_id = format(ctx.trace_id, "032x") if ctx.is_valid else "-"
        record.span_id = format(ctx.span_id, "016x") if ctx.is_valid else "-"
        return json.dumps({{
            "time": self.formatTime(record),
            "level": record.levelname,
            "message": record.getMessage(),
            "trace_id": record.trace_id,
            "span_id": record.span_id,
            "module": record.module,
        }})


def configure_logging(level: str = "INFO") -> None:
    handler = logging.StreamHandler()
    handler.setFormatter(OTelFormatter())
    logging.basicConfig(level=getattr(logging, level), handlers=[handler])
"""

_EXPORTER_CONFIGS = {
    "jaeger": (
        "from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter",
        "OTLPSpanExporter",
        "http://localhost:4317",
    ),
    "grafana": (
        "from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter",
        "OTLPSpanExporter",
        "https://otlp-gateway-prod-us-east-0.grafana.net/otlp",
    ),
    "honeycomb": (
        "from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter",
        "OTLPSpanExporter",
        "https://api.honeycomb.io",
    ),
    "stdout": (
        "from opentelemetry.sdk.trace.export import ConsoleSpanExporter as OTLPSpanExporter",
        "OTLPSpanExporter",
        "",
    ),
}

_MISSING_PACKAGES = {
    "jaeger": [
        "opentelemetry-sdk",
        "opentelemetry-exporter-otlp-proto-grpc",
        "opentelemetry-instrumentation-fastapi",
        "opentelemetry-instrumentation-httpx",
    ],
    "grafana": [
        "opentelemetry-sdk",
        "opentelemetry-exporter-otlp-proto-http",
        "opentelemetry-instrumentation-fastapi",
        "opentelemetry-instrumentation-httpx",
    ],
    "honeycomb": [
        "opentelemetry-sdk",
        "opentelemetry-exporter-otlp-proto-http",
        "opentelemetry-instrumentation-fastapi",
        "opentelemetry-instrumentation-httpx",
    ],
    "stdout": [
        "opentelemetry-sdk",
        "opentelemetry-instrumentation-fastapi",
    ],
}

_SETUP_INSTRUCTIONS = """\
1. Install required packages:
   pip install {packages}

2. Add to your main.py:
   from otel_setup import setup_otel
   app = setup_otel(app)

3. Set environment variable before running:
   export OTEL_EXPORTER_OTLP_ENDPOINT="<your-collector-url>"

4. (Optional) Replace your logging setup:
   from logging_config import configure_logging
   configure_logging()
"""


@dataclass
class OTelConfig:
    stack: list[str] = field(default_factory=list)
    service_name: str = "ace-service"
    exporter: str = "jaeger"
    sample_rate: float = 1.0
    trace_db: bool = True
    trace_http: bool = True
    correlation_ids: bool = True


async def generate_otel_config(
    stack: list[str] | None = None,
    product_id: str = "product:platform",
    exporter: str = "jaeger",
    output_dir: str = ".",
) -> dict:
    """Generate OTel configuration files for the detected stack.

    Returns: {files_written, setup_instructions, missing_packages}
    """
    from core.engine.scanner.hardening import _detect_stack_filesystem

    detected_stack = stack
    if not detected_stack:
        try:
            from core.engine.scanner.hardening import _detect_stack as _db_detect

            detected_stack = await _db_detect(".")
        except Exception:
            detected_stack = _detect_stack_filesystem(".")
    if not detected_stack:
        detected_stack = ["python"]

    service_name = product_id.split(":")[-1].replace("_", "-")
    exporter_key = exporter if exporter in _EXPORTER_CONFIGS else "jaeger"
    exp_import, exp_class, default_endpoint = _EXPORTER_CONFIGS[exporter_key]

    otel_setup = _OTEL_SETUP_PY.format(
        exporter_import=exp_import,
        exporter_class=exp_class,
        service_name=service_name,
        default_endpoint=default_endpoint,
    )
    logging_cfg = _LOGGING_CONFIG_PY

    files_written: list[dict] = []
    os.makedirs(output_dir, exist_ok=True)

    for filename, content, desc in [
        ("otel_setup.py", otel_setup, "OTel SDK initialization + FastAPI instrumentation"),
        ("logging_config.py", logging_cfg, "Structured JSON logging with trace correlation IDs"),
    ]:
        path = os.path.join(output_dir, filename)
        with open(path, "w") as f:
            f.write(content)
        files_written.append({"path": path, "description": desc})

    packages = _MISSING_PACKAGES.get(exporter_key, _MISSING_PACKAGES["jaeger"])
    instructions = _SETUP_INSTRUCTIONS.format(packages=" ".join(packages))

    return {
        "files_written": files_written,
        "setup_instructions": instructions,
        "missing_packages": packages,
        "exporter": exporter_key,
        "stack": detected_stack,
    }
