"""
HTTP utilities for the Unify SDK.

Provides a requests session with retry logic and optional trace-aware logging.

Logging is controlled by two environment variables:
- UNIFY_LOG: Enable/disable logging (default: true)
- UNIFY_LOG_DIR: Directory for log files (default: console only)

When UNIFY_LOG_DIR is set, structured JSON files are written:
- Before request: {timestamp}_{method}_{route}_PENDING_{trace_id}.json
- After response: {timestamp}_{method}_{route}_{duration}ms_{status}_{trace_id}.json

The trace_id suffix enables correlation with pytest logs and Orchestra traces.

OpenTelemetry tracing is controlled by:
- UNIFY_OTEL: Enable/disable OTel tracing (default: false)
- UNIFY_OTEL_ENDPOINT: OTLP endpoint for trace export (optional)
- UNIFY_OTEL_LOG_DIR: Directory for file-based span export (optional)

When UNIFY_OTEL is enabled, HTTP requests create OTel spans that can be
correlated with parent spans (from Unity) and child spans (in Orchestra).

File-based span export:
When UNIFY_OTEL_LOG_DIR is set, spans are written to JSONL files keyed
by trace_id: {UNIFY_OTEL_LOG_DIR}/{trace_id}.jsonl
"""

import json
import logging
import os
import re
import threading
import time
from datetime import datetime, timezone
from functools import wraps
from pathlib import Path
from typing import Any, Callable, Optional
from urllib.parse import urlparse

import requests
from requests.adapters import HTTPAdapter
from urllib3 import Retry

# ---------------------------------------------------------------------------
# Console logging setup
# ---------------------------------------------------------------------------

_LOGGER = logging.getLogger("unify")
_LOG_ENABLED = os.getenv("UNIFY_LOG", "true").lower() in ("true", "1")
_LOGGER.setLevel(logging.DEBUG if _LOG_ENABLED else logging.WARNING)

# ---------------------------------------------------------------------------
# OpenTelemetry setup
# ---------------------------------------------------------------------------

_OTEL_ENABLED = os.getenv("UNIFY_OTEL", "false").lower() in ("true", "1")
_OTEL_ENDPOINT = os.getenv("UNIFY_OTEL_ENDPOINT", "").strip()
_OTEL_LOG_DIR = os.getenv("UNIFY_OTEL_LOG_DIR", "").strip()
_OTEL_INITIALIZED = False
_TRACER = None


# ---------------------------------------------------------------------------
# File-based Span Exporter
# ---------------------------------------------------------------------------


class FileSpanExporter:
    """Exports spans to JSONL files, one file per trace_id.

    This enables standalone trace logging without a parent TracerProvider
    or external collector (Tempo/Jaeger).
    """

    def __init__(self, log_dir: Path, service_name: str = "unify"):
        self.log_dir = log_dir
        self.service_name = service_name
        self._lock = threading.Lock()
        self.log_dir.mkdir(parents=True, exist_ok=True)

    def export(self, spans) -> int:
        """Export a batch of spans to files."""
        try:
            for span in spans:
                self._write_span(span)
            return 0  # SUCCESS
        except Exception as e:
            _LOGGER.warning(f"FileSpanExporter failed to export spans: {e}")
            return 1  # FAILURE

    def _write_span(self, span) -> None:
        """Write a single span to its trace file."""
        ctx = span.get_span_context()
        if ctx is None or not ctx.is_valid:
            return

        trace_id = f"{ctx.trace_id:032x}"
        span_id = f"{ctx.span_id:016x}"

        parent_span_id = None
        if span.parent is not None:
            parent_span_id = f"{span.parent.span_id:016x}"

        span_data: dict[str, Any] = {
            "trace_id": trace_id,
            "span_id": span_id,
            "parent_span_id": parent_span_id,
            "name": span.name,
            "service": self.service_name,
            "start_time": (
                datetime.fromtimestamp(
                    span.start_time / 1e9,
                    tz=timezone.utc,
                ).isoformat()
                if span.start_time
                else None
            ),
            "end_time": (
                datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc).isoformat()
                if span.end_time
                else None
            ),
            "duration_ms": (
                (span.end_time - span.start_time) / 1e6
                if span.end_time and span.start_time
                else None
            ),
            "status": span.status.status_code.name if span.status else None,
            "attributes": dict(span.attributes) if span.attributes else {},
        }

        trace_file = self.log_dir / f"{trace_id}.jsonl"
        with self._lock:
            with open(trace_file, "a", encoding="utf-8") as f:
                f.write(json.dumps(span_data, default=str) + "\n")

    def shutdown(self) -> None:
        """Shutdown the exporter (no-op for file exporter)."""

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush any buffered spans (no-op for file exporter)."""
        return True


def _setup_otel() -> None:
    """Initialize OpenTelemetry if enabled and not already configured.

    When UNIFY_OTEL_LOG_DIR is set, spans are written to JSONL files
    keyed by trace_id for standalone trace logging.
    """
    global _OTEL_INITIALIZED, _TRACER

    if _OTEL_INITIALIZED or not _OTEL_ENABLED:
        return

    _OTEL_INITIALIZED = True

    try:
        from opentelemetry import trace
        from opentelemetry.sdk.resources import SERVICE_NAME, Resource
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import SimpleSpanProcessor

        # Check if a TracerProvider already exists (parent set it up)
        existing = trace.get_tracer_provider()
        if existing and not isinstance(
            existing,
            (trace.NoOpTracerProvider, trace.ProxyTracerProvider),
        ):
            # Parent (e.g., Unity) already configured OTel - use theirs
            _TRACER = trace.get_tracer("unify")
            _LOGGER.debug("Using existing OTel TracerProvider from parent")
            return

        # We're the outermost layer - set up our own provider
        resource = Resource.create({SERVICE_NAME: "unify"})
        provider = TracerProvider(resource=resource)

        # Add OTLP exporter if endpoint configured
        if _OTEL_ENDPOINT:
            try:
                from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
                    OTLPSpanExporter,
                )
                from opentelemetry.sdk.trace.export import BatchSpanProcessor

                exporter = OTLPSpanExporter(endpoint=_OTEL_ENDPOINT, insecure=True)
                provider.add_span_processor(BatchSpanProcessor(exporter))
                _LOGGER.debug(f"Configured OTLP exporter at {_OTEL_ENDPOINT}")
            except ImportError:
                _LOGGER.warning(
                    "OTLP exporter not available - install opentelemetry-exporter-otlp",
                )
            except Exception as e:
                _LOGGER.warning(f"Failed to configure OTLP exporter: {e}")

        # Add file exporter if log directory configured
        if _OTEL_LOG_DIR:
            try:
                log_dir = Path(_OTEL_LOG_DIR)
                file_exporter = FileSpanExporter(log_dir, service_name="unify")
                provider.add_span_processor(SimpleSpanProcessor(file_exporter))
                _LOGGER.debug(f"Configured file span exporter at {_OTEL_LOG_DIR}")
            except Exception as e:
                _LOGGER.warning(f"Failed to configure file span exporter: {e}")

        trace.set_tracer_provider(provider)
        _TRACER = trace.get_tracer("unify")
        _LOGGER.debug("Initialized OTel TracerProvider for unify")

    except ImportError:
        _LOGGER.debug("OpenTelemetry not available - tracing disabled")
    except Exception as e:
        _LOGGER.warning(f"Failed to initialize OpenTelemetry: {e}")


def get_tracer():
    """Get the OpenTelemetry tracer, initializing if needed."""
    global _TRACER
    if _TRACER is None and _OTEL_ENABLED:
        _setup_otel()
    return _TRACER


def is_otel_enabled() -> bool:
    """Check if OpenTelemetry tracing is enabled."""
    return _OTEL_ENABLED


def get_otel_log_dir() -> Optional[Path]:
    """Get the OTel span log directory, if configured."""
    if not _OTEL_LOG_DIR:
        return None
    return Path(_OTEL_LOG_DIR)


# ---------------------------------------------------------------------------
# File-based trace logging
# ---------------------------------------------------------------------------

_LOG_DIR: Optional[Path] = None
_LOG_DIR_CHECKED = False


def configure_log_dir(log_dir: Optional[str] = None) -> Optional[Path]:
    """Configure or reconfigure the log directory for file-based logging.

    Call this after setting UNIFY_LOG_DIR if the env var was set
    after this module was imported.

    Args:
        log_dir: Explicit log directory path. If None, reads from
                 UNIFY_LOG_DIR env var.

    Returns:
        The configured log directory Path, or None if disabled.
    """
    global _LOG_DIR, _LOG_DIR_CHECKED

    _LOG_DIR_CHECKED = False
    _LOG_DIR = None

    if log_dir is not None:
        os.environ["UNIFY_LOG_DIR"] = log_dir

    return _get_log_dir()


def _get_log_dir() -> Optional[Path]:
    """Get the log directory from UNIFY_LOG_DIR env var.

    Returns None if not set or directory creation fails.
    The directory is created on first access.
    """
    global _LOG_DIR, _LOG_DIR_CHECKED

    if _LOG_DIR_CHECKED:
        return _LOG_DIR

    _LOG_DIR_CHECKED = True
    log_dir_str = os.getenv("UNIFY_LOG_DIR", "").strip()
    if not log_dir_str:
        return None

    try:
        log_dir = Path(log_dir_str)
        log_dir.mkdir(parents=True, exist_ok=True)
        _LOG_DIR = log_dir
        _LOGGER.debug(f"HTTP trace logging enabled: {log_dir}")
    except Exception as e:
        _LOGGER.warning(f"Failed to create HTTP log directory {log_dir_str}: {e}")
        _LOG_DIR = None

    return _LOG_DIR


def _get_current_trace_id() -> Optional[str]:
    """Extract the current OpenTelemetry trace_id if available.

    Returns a 32-character hex string, or None if no active span.
    """
    try:
        from opentelemetry import trace

        span = trace.get_current_span()
        if span is None:
            return None
        ctx = span.get_span_context()
        if ctx is not None and ctx.is_valid:
            return f"{ctx.trace_id:032x}"
    except Exception:
        pass
    return None


def _inject_trace_context(headers: Optional[dict]) -> dict:
    """Inject W3C Trace Context headers into outgoing request headers.

    This enables distributed tracing across service boundaries (e.g., to Orchestra).
    The traceparent header format is: {version}-{trace-id}-{parent-id}-{trace-flags}

    Args:
        headers: Existing headers dict (may be None)

    Returns:
        Headers dict with traceparent injected (if OTel span is active)
    """
    result = dict(headers) if headers else {}

    if not _OTEL_ENABLED:
        return result

    try:
        from opentelemetry import trace
        from opentelemetry.trace.propagation.tracecontext import (
            TraceContextTextMapPropagator,
        )

        span = trace.get_current_span()
        if span is None:
            return result

        ctx = span.get_span_context()
        if ctx is None or not ctx.is_valid:
            return result

        # Use the standard W3C propagator to inject headers
        propagator = TraceContextTextMapPropagator()
        propagator.inject(result)

        _LOGGER.debug(f"Injected traceparent: {result.get('traceparent', 'N/A')}")

    except ImportError:
        pass
    except Exception as e:
        _LOGGER.debug(f"Failed to inject trace context: {e}")

    return result


def get_traceparent() -> Optional[str]:
    """Get the current traceparent header value, if an OTel span is active.

    Returns:
        The traceparent header string, or None if no active span.
    """
    if not _OTEL_ENABLED:
        return None

    try:
        from opentelemetry import trace
        from opentelemetry.trace.propagation.tracecontext import (
            TraceContextTextMapPropagator,
        )

        span = trace.get_current_span()
        if span is None:
            return None

        ctx = span.get_span_context()
        if ctx is None or not ctx.is_valid:
            return None

        # Inject into a temp dict to get the header value
        carrier: dict = {}
        propagator = TraceContextTextMapPropagator()
        propagator.inject(carrier)

        return carrier.get("traceparent")

    except Exception:
        return None


def _extract_route(url: str) -> str:
    """Extract a safe route identifier from a URL for use in filenames.

    Examples:
        https://api.unify.ai/v0/logs -> logs
        https://api.unify.ai/v0/logs/derived -> logs-derived
        https://api.unify.ai/v0/project/foo/contexts -> project-contexts
    """
    try:
        parsed = urlparse(url)
        path = parsed.path.strip("/")
        # Remove version prefix (v0/, v1/, etc.)
        path = re.sub(r"^v\d+/", "", path)
        # Replace path separators and collapse multiple dashes
        route = re.sub(r"[/]+", "-", path)
        route = re.sub(r"-+", "-", route)
        # Limit length and remove unsafe chars
        route = re.sub(r"[^a-zA-Z0-9_-]", "", route)[:30]
        return route or "unknown"
    except Exception:
        return "unknown"


def _mask_headers(headers: Optional[dict]) -> Optional[dict]:
    """Return a copy of headers with Authorization masked."""
    if not headers:
        return headers
    masked = dict(headers)
    if "Authorization" in masked:
        masked["Authorization"] = "***"
    if "authorization" in masked:
        masked["authorization"] = "***"
    return masked


def _write_pending_trace(
    method: str,
    url: str,
    request_kwargs: dict,
) -> Optional[Path]:
    """Write a pending trace file before the request is made.

    Returns the file path for later finalization, or None if logging disabled.
    """
    log_dir = _get_log_dir()
    if log_dir is None:
        return None

    try:
        trace_id = _get_current_trace_id()
        trace_suffix = trace_id[-8:] if trace_id else "no-trace"
        route = _extract_route(url)

        now = datetime.now(timezone.utc)
        timestamp = now.strftime("%H-%M-%S.") + f"{now.microsecond // 1000:03d}"

        filename = f"{timestamp}_{method}_{route}_PENDING_{trace_suffix}.json"
        filepath = log_dir / filename

        # Build the trace record
        record = {
            "timestamp": now.isoformat(),
            "trace_id": trace_id,
            "method": method,
            "url": url,
            "route": route,
            "status": "pending",
            "request": {
                "headers": _mask_headers(request_kwargs.get("headers")),
                "params": request_kwargs.get("params"),
                "json": request_kwargs.get("json"),
                "data": (
                    str(request_kwargs.get("data"))[:1000]
                    if request_kwargs.get("data")
                    else None
                ),
            },
        }

        with filepath.open("w", encoding="utf-8") as f:
            json.dump(record, f, indent=2, default=str)

        return filepath
    except Exception as e:
        _LOGGER.debug(f"Failed to write pending trace: {e}")
        return None


def _finalize_trace(
    pending_path: Optional[Path],
    response: requests.Response,
    duration_ms: int,
) -> None:
    """Finalize a pending trace file with response data.

    Renames the file from _PENDING_ to _{duration}ms_{status}_.
    """
    if pending_path is None or not pending_path.exists():
        return

    try:
        # Read existing record
        with pending_path.open("r", encoding="utf-8") as f:
            record = json.load(f)

        # Add response data
        record["status"] = "complete"
        record["duration_ms"] = duration_ms
        record["response"] = {
            "status_code": response.status_code,
            "headers": dict(response.headers),
        }

        # Try to include response body (truncated for large responses)
        try:
            body = response.json()
            # Truncate large responses
            body_str = json.dumps(body)
            if len(body_str) > 10000:
                record["response"]["body"] = "(truncated, >10KB)"
                record["response"]["body_preview"] = body_str[:1000]
            else:
                record["response"]["body"] = body
        except Exception:
            text = response.text
            if len(text) > 1000:
                record["response"]["body"] = text[:1000] + "...(truncated)"
            else:
                record["response"]["body"] = text

        # Write updated record
        with pending_path.open("w", encoding="utf-8") as f:
            json.dump(record, f, indent=2, default=str)

        # Rename file to reflect completion
        new_name = pending_path.name.replace(
            "_PENDING_",
            f"_{duration_ms}ms_{response.status_code}_",
        )
        new_path = pending_path.parent / new_name
        pending_path.rename(new_path)

    except Exception as e:
        _LOGGER.debug(f"Failed to finalize trace: {e}")


def _mark_trace_failed(
    pending_path: Optional[Path],
    error: Exception,
    duration_ms: int,
) -> None:
    """Mark a pending trace as failed due to an exception."""
    if pending_path is None or not pending_path.exists():
        return

    try:
        with pending_path.open("r", encoding="utf-8") as f:
            record = json.load(f)

        record["status"] = "failed"
        record["duration_ms"] = duration_ms
        record["error"] = {
            "type": type(error).__name__,
            "message": str(error),
        }

        with pending_path.open("w", encoding="utf-8") as f:
            json.dump(record, f, indent=2, default=str)

        # Rename to indicate failure
        new_name = pending_path.name.replace("_PENDING_", f"_{duration_ms}ms_FAILED_")
        new_path = pending_path.parent / new_name
        pending_path.rename(new_path)

    except Exception as e:
        _LOGGER.debug(f"Failed to mark trace as failed: {e}")


# ---------------------------------------------------------------------------
# Session configuration
# ---------------------------------------------------------------------------

_SESSION = requests.Session()
_RETRIES = Retry(
    total=5,
    connect=3,
    read=2,
    backoff_factor=0.1,
    status_forcelist=[500, 502, 503, 504],
    allowed_methods=None,  # Retry all methods including POST
)
_ADAPTER = HTTPAdapter(max_retries=_RETRIES, pool_connections=20, pool_maxsize=20)
_SESSION.mount("http://", _ADAPTER)
_SESSION.mount("https://", _ADAPTER)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


class RequestError(Exception):
    def __init__(self, url: str, r_type: str, response: requests.Response, /, **kwargs):
        super().__init__(
            f"{r_type}:{url} with {kwargs} failed with status code "
            f"{response.status_code}: {response.text}",
        )
        self.response = response


def _log_to_console(log_type: str, url: str, mask_key: bool = True, /, **kwargs):
    """Log request/response details to the console logger."""
    kwargs_str = ""
    safe_kwargs = dict(kwargs)

    if mask_key and "headers" in safe_kwargs:
        safe_kwargs["headers"] = _mask_headers(safe_kwargs["headers"])

    for k, v in safe_kwargs.items():
        if isinstance(v, dict):
            kwargs_str += f"{k}:{json.dumps(v, indent=2)},\n"
        else:
            kwargs_str += f"{k}:{v},\n"

    trace_id = _get_current_trace_id()
    trace_prefix = f"[{trace_id[-8:]}] " if trace_id else ""

    log_msg = f"""
{trace_prefix}====== {log_type} =======
url:{url}
{kwargs_str}
"""
    _LOGGER.debug(log_msg)


def _mask_auth_key(kwargs: dict) -> dict:
    """Return a sanitized copy of request kwargs suitable for error messages.

    IMPORTANT: This must NOT mutate caller-owned objects (e.g. the headers dict),
    because higher-level callers may reuse those objects across multiple HTTP
    calls (e.g. in loops).
    """
    if "headers" not in kwargs:
        return kwargs

    headers = kwargs.get("headers")
    if not isinstance(headers, dict):
        return kwargs

    safe_kwargs = dict(kwargs)
    safe_kwargs["headers"] = _mask_headers(headers)
    return safe_kwargs


def _log_request_if_enabled(fn: Callable) -> Callable:
    """Decorator that adds console logging, file-based logging, and OTel tracing."""
    if not _LOG_ENABLED and not _OTEL_ENABLED:
        return fn

    @wraps(fn)
    def inner(method: str, url: str, **kwargs) -> requests.Response:
        # Console log: request (if logging enabled)
        if _LOG_ENABLED:
            _log_to_console(f"{method}", url, True, **kwargs)

        # File trace: pending (if logging enabled)
        pending_path = (
            _write_pending_trace(method, url, kwargs) if _LOG_ENABLED else None
        )

        # Get tracer for OTel spans
        tracer = get_tracer()
        route = _extract_route(url)
        span_name = f"{method} {route}"

        # Create OTel span context manager (or dummy if not enabled)
        if tracer is not None:
            from opentelemetry.trace import SpanKind, Status, StatusCode

            span_cm = tracer.start_as_current_span(
                span_name,
                kind=SpanKind.CLIENT,
            )
        else:
            from contextlib import nullcontext

            span_cm = nullcontext()

        start_time = time.monotonic()
        with span_cm as span:
            try:
                # Set span attributes before request
                if span is not None:
                    span.set_attribute("http.method", method)
                    span.set_attribute("http.url", url)
                    span.set_attribute("http.route", route)
                    if kwargs.get("params"):
                        span.set_attribute(
                            "http.request.params",
                            json.dumps(kwargs["params"]),
                        )

                # Inject trace context headers for distributed tracing
                if _OTEL_ENABLED:
                    kwargs = dict(kwargs)  # Don't mutate caller's kwargs
                    kwargs["headers"] = _inject_trace_context(kwargs.get("headers"))

                res: requests.Response = fn(method, url, **kwargs)
                duration_ms = int((time.monotonic() - start_time) * 1000)

                # Set span attributes after response
                if span is not None:
                    span.set_attribute("http.status_code", res.status_code)
                    span.set_attribute("http.duration_ms", duration_ms)
                    if res.status_code >= 400:
                        span.set_status(Status(StatusCode.ERROR))
                    else:
                        span.set_status(Status(StatusCode.OK))

                # Console log: response (if logging enabled)
                if _LOG_ENABLED:
                    try:
                        _log_to_console(
                            f"{method} response:{res.status_code}",
                            url,
                            response=res.json(),
                        )
                    except requests.exceptions.JSONDecodeError:
                        _log_to_console(
                            f"{method} response:{res.status_code}",
                            url,
                            response=res.text[:500] if res.text else "(empty)",
                        )

                # File trace: finalize (if logging enabled)
                if _LOG_ENABLED:
                    _finalize_trace(pending_path, res, duration_ms)

                return res

            except Exception as e:
                duration_ms = int((time.monotonic() - start_time) * 1000)

                # Set span error status
                if span is not None:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.set_attribute("error.type", type(e).__name__)
                    span.set_attribute("error.message", str(e))
                    span.set_attribute("http.duration_ms", duration_ms)

                # File trace: mark failed (if logging enabled)
                if _LOG_ENABLED:
                    _mark_trace_failed(pending_path, e, duration_ms)

                raise

    return inner


@_log_request_if_enabled
def request(method, url, raise_for_status=True, **kwargs) -> requests.Response:
    try:
        res = _SESSION.request(method, url, **kwargs)
        if raise_for_status:
            res.raise_for_status()
        return res
    except requests.exceptions.HTTPError as e:
        kwargs = _mask_auth_key(kwargs)
        raise RequestError(url, method, e.response, **kwargs)


def get(url, params=None, **kwargs):
    return request("GET", url, params=params, **kwargs)


def options(url, **kwargs):
    return request("OPTIONS", url, **kwargs)


def head(url, **kwargs):
    return request("HEAD", url, **kwargs)


def post(url, data=None, json=None, **kwargs):
    return request("POST", url, data=data, json=json, **kwargs)


def put(url, data=None, **kwargs):
    return request("PUT", url, data=data, **kwargs)


def patch(url, data=None, **kwargs):
    return request("PATCH", url, data=data, **kwargs)


def delete(url, **kwargs):
    return request("DELETE", url, **kwargs)
