from __future__ import annotations

import random
from typing import TYPE_CHECKING, Any

from portkey_ai import (
    PORTKEY_GATEWAY_URL,
    AsyncPortkey,
)

from pipelex import log
from pipelex.cogt.extract.extract_job import ExtractJob
from pipelex.cogt.img_gen.img_gen_job import ImgGenJob
from pipelex.cogt.inference.inference_constants import InferenceOutputType
from pipelex.cogt.llm.llm_job import LLMJob
from pipelex.hub import get_telemetry_manager
from pipelex.plugins.gateway.gateway_exceptions import GatewayCredentialsError
from pipelex.plugins.gateway.gateway_protocols import GatewayExtractProtocol
from pipelex.plugins.gateway.gateway_schemas import GatewayExtractRequestParams
from pipelex.plugins.google.google_img_gen_factory import GoogleImgGenFactory
from pipelex.plugins.portkey.portkey_constants import PortkeyHeaderKey
from pipelex.system.telemetry.otel_constants import OTelConstants

if TYPE_CHECKING:
    from pipelex.cogt.inference.inference_job_abstract import InferenceJobAbstract
    from pipelex.cogt.model_backends.backend import InferenceBackend
    from pipelex.cogt.model_backends.model_spec import InferenceModelSpec


class GatewayFactory:
    @classmethod
    def is_debug_enabled(cls, backend: InferenceBackend) -> bool:
        is_debug_configured = backend.extra_config.get("debug", False)
        return get_telemetry_manager().is_pipelex_gateway_portkey_logging_enabled(is_debug_configured=is_debug_configured)

    @classmethod
    def get_endpoint(cls, backend: InferenceBackend) -> str:
        return backend.endpoint or PORTKEY_GATEWAY_URL

    @classmethod
    def get_api_key(cls, backend: InferenceBackend) -> str:
        if not backend.api_key:
            msg = "Portkey API key is not set"
            raise GatewayCredentialsError(msg)
        return backend.api_key

    @classmethod
    def make_portkey_client(
        cls,
        backend: InferenceBackend,
    ) -> AsyncPortkey:
        is_debug_enabled = cls.is_debug_enabled(backend=backend)
        endpoint = cls.get_endpoint(backend=backend)
        api_key = cls.get_api_key(backend=backend)
        log.verbose(f"Making Portkey client with endpoint: {endpoint}, debug: {is_debug_enabled}")

        return AsyncPortkey(
            base_url=endpoint,
            api_key=api_key,
            debug=is_debug_enabled,
        )

    @classmethod
    def make_extras(
        cls, inference_model: InferenceModelSpec, inference_job: InferenceJobAbstract, output_desc: str
    ) -> tuple[dict[str, str], dict[str, Any]]:
        extra_headers: dict[str, str] = {}
        extra_body: dict[str, Any] = {}
        if inference_model.extra_headers:
            extra_headers.update(inference_model.extra_headers)
        if not extra_headers.get(PortkeyHeaderKey.CONFIG) and not extra_headers.get(PortkeyHeaderKey.PROVIDER):
            extra_headers[PortkeyHeaderKey.PROVIDER] = inference_model.backend_name

        if isinstance(inference_job, ExtractJob):
            # Derive boolean from max_nb_images: None/positive = True, 0 = False
            max_nb_images = inference_job.job_params.max_nb_images
            should_include_images = max_nb_images is None or max_nb_images > 0
            extract_protocol = GatewayExtractProtocol.make_from_model_handle(model_handle=inference_model.name)
            match extract_protocol:
                case GatewayExtractProtocol.MISTRAL_DOC_AI:
                    extra_body["include_image_base64"] = should_include_images
                case GatewayExtractProtocol.AZURE_DOC_INTEL:
                    request_params = GatewayExtractRequestParams(should_include_images=should_include_images)
                    messages_azure: list[dict[str, str]] = [{"role": "user", "content": request_params.model_dump_json()}]
                    extra_body["messages"] = messages_azure
                case GatewayExtractProtocol.DEEPSEEK_OCR:
                    messages_deepseek: list[dict[str, str]] = [{"role": "user", "content": "Convert the document to markdown."}]
                    extra_body["messages"] = messages_deepseek
                case GatewayExtractProtocol.LINKUP_FETCH:
                    pass  # Fetch params are built directly in GatewayExtractWorker._extract_web_fetch
        elif isinstance(inference_job, LLMJob) and inference_model.model_id.lower().startswith("mistral-") and inference_job.job_params.seed is None:
            # Mistral models really want non-null seed
            extra_body["seed"] = random.randint(0, 1000000)
        elif isinstance(inference_job, ImgGenJob) and inference_model.model_id.startswith("gemini"):
            aspect_ratio_str = GoogleImgGenFactory.aspect_ratio_literal(inference_job.job_params.aspect_ratio)
            extra_body["image_config"] = {
                "aspect_ratio": aspect_ratio_str,
                # "image_size": "2K",
            }
        # OTel-correlated Portkey tracing (only when enabled and OTel context available)
        if get_telemetry_manager().is_pipelex_gateway_portkey_tracing_enabled() and (otel_context := inference_job.job_metadata.otel_context):
            # Use OTel trace_id and span_id for correlation
            extra_headers[PortkeyHeaderKey.TRACE_ID] = f"{otel_context.trace_id:032x}"
            extra_headers[PortkeyHeaderKey.SPAN_ID] = f"{otel_context.span_id:016x}"

            # Build span name with redacted output class name (consistent with Pipelex telemetry policy)
            # Pipelex services always redact sensitive data to protect user privacy
            unit_job_id = inference_job.job_metadata.unit_job_id or "unknown"
            display_output = output_desc if output_desc == InferenceOutputType.TEXT else OTelConstants.OUTPUT_CLASS_REDACTED
            extra_headers[PortkeyHeaderKey.SPAN_NAME] = f"{unit_job_id} -> {display_output}"

        return extra_headers, extra_body
