# ======================================================================
# (A) FILE PATH & IMPORT PATH
# depths/cli/app.py  →  import path: depths.cli.app
# ======================================================================

# ======================================================================
# (B) FILE OVERVIEW (concept & significance in v0.1.0)
# FastAPI service that exposes the Depths OTLP/HTTP ingestion surface and a
# minimal query API. Responsibilities:
#   • Accept OTLP JSON or protobuf for traces/logs/metrics and forward to
#     DepthsLogger (which orchestrates Producer + Aggregator + Shipper).
#   • Provide health and simple read endpoints backed by DepthsLogger readers.
#   • Own the app lifecycle (start logger on boot; auto-flush on shutdown).
#
# This is the main “edge” for v0.1.0: it receives telemetry and stores it
# into per-day Delta tables locally (and eventually ships sealed days to S3).
# ======================================================================

# ======================================================================
# (C) IMPORTS & GLOBALS (what & why)
# FastAPI, Request/Response/Headers/Query  → HTTP server & validation
# polars as pl                              → LazyFrame → DataFrame collection for /api/* queries
# json, gzip, os                            → decoding, compression handling, instance env
# DepthsLogger, DepthsLoggerOptions, S3Config → core runtime & S3 wiring
# google.protobuf… (optional)               → OTLP protobuf (falls back to JSON if unavailable)
#
# Globals:
#   _PB_OK      → whether protobuf deps are importable
#   _LOGGER     → process-wide singleton DepthsLogger (created lazily)
#   app         → FastAPI instance (with lifespan manager)
# ======================================================================

from __future__ import annotations

import gzip
import json
import os
from typing import Optional, Tuple, Iterable, List

import polars as pl
from fastapi import FastAPI, Request, Response, Header, HTTPException, Query
from fastapi.responses import JSONResponse, PlainTextResponse
from contextlib import asynccontextmanager

from depths.core.config import S3Config, DepthsLoggerOptions
from depths.core.logger import DepthsLogger

# Protobuf support
try:
    from google.protobuf.json_format import MessageToDict
    from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
        ExportTraceServiceRequest,
        ExportTraceServiceResponse,
    )
    from opentelemetry.proto.collector.metrics.v1.metrics_service_pb2 import (
        ExportMetricsServiceRequest,
        ExportMetricsServiceResponse,
    )
    from opentelemetry.proto.collector.logs.v1.logs_service_pb2 import (
        ExportLogsServiceRequest,
        ExportLogsServiceResponse,
    )
    _PB_OK = True
except ImportError:
    _PB_OK = False
    
# ---- Singleton logger wiring -------------------------------------------------

_LOGGER: Optional[DepthsLogger] = None



def _get_instance_settings() -> Tuple[str, str]:
    """
    Resolve Depths instance identity and data directory from the environment.

    Overview (v0.1.0 role):
        Centralizes how the service decides where to stage/write data on disk.
        Used by the singleton logger constructor to ensure consistent layout.

    Returns:
        (instance_id, instance_dir_abs): instance_id from DEPTHS_INSTANCE_ID
        (default 'default'), and absolute instance_dir from DEPTHS_INSTANCE_DIR
        (default './depths_data' resolved to an absolute path).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Keep defaults stable; the CLI injects these env vars for background runs.
    # - Do not validate existence here; the logger ensures/creates directories.

    instance_id = os.environ.get("DEPTHS_INSTANCE_ID", "default")
    instance_dir = os.environ.get("DEPTHS_INSTANCE_DIR", os.path.abspath("./depths_data"))
    return instance_id, instance_dir


def _get_logger() -> DepthsLogger:
    """
    Return the process-wide DepthsLogger singleton (construct if missing).

    Overview (v0.1.0 role):
        Wires in S3 (if env is present) and standard DepthsLoggerOptions, then
        starts background services per options. All endpoints use this singleton.

    Returns:
        DepthsLogger instance.

    Raises:
        (None directly) — S3Config.from_env errors are swallowed to allow
        local-only operation when S3 isn't configured.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Lazily initialized to avoid side effects during module import.
    # - Options are currently defaulted; evolve later to read persisted configs.
    # - S3 optional; local-only mode if not configured

    global _LOGGER
    if _LOGGER is None:
        instance_id, instance_dir = _get_instance_settings()
        s3 = None
        try:
            s3 = S3Config.from_env()
        except Exception:
            s3 = None  
        _LOGGER = DepthsLogger(
            instance_id=instance_id,
            instance_dir=instance_dir,
            s3=s3,
            options=DepthsLoggerOptions(),
        )
    return _LOGGER

@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    FastAPI lifespan manager that starts the logger on startup and
    gracefully flushes buffers on shutdown.

    Overview (v0.1.0 role):
        Ensures the Aggregators are running before any traffic arrives and that
        pending batches are flushed (with a bounded wait) on process exit.

    Args:
        app: FastAPI instance (unused, required by FastAPI).

    Yields:
        None (control returns to FastAPI to serve requests).

    Exceptions:
        Any start/stop exceptions are caught to keep the server responsive.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Uses DepthsLogger.astop(flush="auto") which applies a small wait envelope
    #   to encourage a final quiet/age-triggered flush before returning.

    lg = _get_logger()
    try:
        lg.start()  
    except Exception:
        pass
    try:
        yield
    finally:
        try:
            await lg.astop(flush="auto")  
        except Exception:
            pass

app = FastAPI(title="depths OTLP/HTTP", version="0.1.0", lifespan=lifespan)

def _want_json(content_type: str) -> bool:
    """
    Predicate for OTLP JSON payloads.

    Args:
        content_type: Raw Content-Type header.

    Returns:
        True if the media type equals 'application/json' (ignoring parameters).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Keep strict; OTLP spec uses 'application/json' and 'application/x-protobuf'.

    return (content_type or "").split(";")[0].strip().lower() == "application/json"


def _want_protobuf(content_type: str) -> bool:
    """
    Predicate for OTLP protobuf payloads.

    Args:
        content_type: Raw Content-Type header.

    Returns:
        True if the media type equals 'application/x-protobuf'.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Gate protobuf processing on _PB_OK; when False we return a 415-compatible
    #   empty response from _empty_response_for(...) to keep clients happy.

    return (content_type or "").split(";")[0].strip().lower() == "application/x-protobuf"


async def _read_body(request: Request) -> bytes:
    """
    Read and (optionally) gunzip the request body.

    Overview (v0.1.0 role):
        OTLP senders often compress payloads. This function normalizes input
        for both JSON and protobuf ingestion paths.

    Args:
        request: FastAPI Request object.

    Returns:
        Decompressed raw bytes.

    Raises:
        HTTPException(400): When 'content-encoding: gzip' is set but decompression fails.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - We only support 'gzip' here; other codings can be added if needed.

    raw = await request.body()
    enc = request.headers.get("content-encoding", "").lower()
    if enc == "gzip":
        try:
            return gzip.decompress(raw)
        except Exception:
            raise HTTPException(status_code=400, detail="gzip_decompress_failed")
    return raw


def _project_from(request: Request, header_override: Optional[str]) -> Optional[str]:
    """
    Determine project_id for ingestion.

    Overview (v0.1.0 role):
        Allows multi-tenancy by checking a sequence of locations:
        explicit override header param, 'x-depths-project-id', 'x-otlp-project-id',
        and 'project_id' query param.

    Args:
        request: FastAPI Request (used for headers and query params).
        header_override: Value from dependency-injected header param.

    Returns:
        project_id string or None if not provided.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - DepthsLogger also has a mapper default; passing None falls back to that.

    return (
        header_override
        or request.headers.get("x-depths-project-id")
        or request.headers.get("x-otlp-project-id")
        or request.query_params.get("project_id")
    )


def _csv_list(x: Optional[str]) -> Optional[List[str]]:
    """
    Parse a comma-separated string into a list of non-empty, stripped tokens.

    Args:
        x: Raw comma-separated string or None.

    Returns:
        List of strings or None if input is falsy.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Used by /api/* 'select' query to build a column list.

    if not x:
        return None
    return [s for s in (t.strip() for t in x.split(",")) if s]


def _collect_json(lf: pl.LazyFrame) -> list[dict]:
    """
    Collect a LazyFrame and return JSON-friendly row dicts.

    Overview (v0.1.0 role):
        Shared collector for the query endpoints to keep error handling
        consistent and avoid duplicating Polars interaction.

    Args:
        lf: Polars LazyFrame.

    Returns:
        List of Python dictionaries (Polars converts columns to native types).

    Raises:
        HTTPException(500): if .collect() fails (e.g., I/O/scan errors).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Avoids returning Polars objects directly; keeps API payloads simple.

    try:
        df = lf.collect()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"query_collect_failed: {e}")
    return df.to_dicts()


def _empty_response_for(content_type: str, signal: str, rejected: int = 0, error_message: Optional[str] = None) -> Response:
    """
    Build an OTLP-compliant empty (or partial-success) response.

    Overview (v0.1.0 role):
        OTLP Export*Service responses are empty on full success; on partial
        success we populate the appropriate 'partialSuccess' fields for the
        given signal (traces/metrics/logs) in JSON or protobuf.

    Args:
        content_type: Request Content-Type to mirror (json/protobuf).
        signal: One of 'traces' | 'metrics' | 'logs'.
        rejected: Count of rejected records (if any).
        error_message: Text to include with partial success.

    Returns:
        FastAPI Response (JSONResponse, protobuf bytes Response, or 415 text).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - When protobuf libs are unavailable, we return a helpful 415-ish message
    #   via JSON/text to nudge users; JSON path is always available.

    if _want_json(content_type):
        if rejected and error_message is not None:
            if signal == "traces":
                body = {"partialSuccess": {"rejectedSpans": int(rejected), "errorMessage": error_message}}
            elif signal == "metrics":
                body = {"partialSuccess": {"rejectedDataPoints": int(rejected), "errorMessage": error_message}}
            else:
                body = {"partialSuccess": {"rejectedLogRecords": int(rejected), "errorMessage": error_message}}
            return JSONResponse(content=body, media_type="application/json")
        return JSONResponse(content={}, media_type="application/json")
    elif _want_protobuf(content_type) and _PB_OK:
        if signal == "traces":
            msg = ExportTraceServiceResponse()
            if rejected and error_message is not None:
                ps = msg.partial_success
                ps.rejected_spans = int(rejected)
                ps.error_message = error_message
        elif signal == "metrics":
            msg = ExportMetricsServiceResponse()
            if rejected and error_message is not None:
                ps = msg.partial_success
                ps.rejected_data_points = int(rejected)
                ps.error_message = error_message
        else:
            msg = ExportLogsServiceResponse()
            if rejected and error_message is not None:
                ps = msg.partial_success
                ps.rejected_log_records = int(rejected)
                ps.error_message = error_message
        return Response(content=msg.SerializeToString(), media_type="application/x-protobuf")
    else:
        return PlainTextResponse(
            "Unsupported Content-Type. Use application/json or install opentelemetry-proto for protobuf.",
            status_code=415,
        )

@app.get("/healthz")
async def health():
    """
    Liveness/quick diagnostic endpoint.

    Returns:
        JSON with {'ok': True, 'logger': <metrics dict>} reflecting current
        aggregator states and minimal DepthsLogger process info.
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Keep inexpensive; metrics are already aggregated in memory.

    lg = _get_logger()
    m = lg.metrics()
    return {"ok": True, "logger": m}


@app.post("/v1/traces")
async def post_traces(
    request: Request,
    content_type: Optional[str] = Header(default="application/json"),
    x_depths_project_id: Optional[str] = Header(default=None),
):
    """
    OTLP/HTTP traces ingestion endpoint (JSON or protobuf).

    Overview (v0.1.0 role):
        Decodes payload, passes to DepthsLogger.ingest_otlp_traces_json(...),
        and responds with an empty or partial-success ExportTraceServiceResponse.

    Args:
        request: FastAPI Request used to read the body (possibly gzip).
        content_type: 'application/json' or 'application/x-protobuf'.
        x_depths_project_id: Optional per-request project override.

    Returns:
        FastAPI Response appropriate for the content type (JSON/protobuf).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Rejects due to schema/date checks surface in partialSuccess (rejectedSpans).
    # - Non-JSON/protobuf content types return HTTP 415.

    lg = _get_logger()
    body = await _read_body(request)
    project_id = _project_from(request, x_depths_project_id)

    try:
        if _want_json(content_type):
            payload = json.loads(body.decode("utf-8") or "{}")
        elif _want_protobuf(content_type):
            if not _PB_OK:
                return _empty_response_for(content_type, "traces")  # 415 handled inside
            req = ExportTraceServiceRequest()
            req.ParseFromString(body)
            payload = MessageToDict(req, preserving_proto_field_name=True)
        else:
            raise HTTPException(status_code=415, detail="unsupported_content_type")
    except Exception:
        raise HTTPException(status_code=400, detail="bad_request_payload")

    res = lg.ingest_otlp_traces_json(payload, project_id=project_id)
    rejected = int(res.get("rejected", 0))
    if rejected > 0:
        return _empty_response_for(
            content_type, "traces", rejected=rejected, error_message="Some spans/events/links were rejected due to schema or date checks."
        )
    return _empty_response_for(content_type, "traces")


@app.post("/v1/logs")
async def post_logs(
    request: Request,
    content_type: Optional[str] = Header(default="application/json"),
    x_depths_project_id: Optional[str] = Header(default=None),
):
    """
    OTLP/HTTP logs ingestion endpoint (JSON or protobuf).

    Args:
        request: FastAPI Request used to read (maybe gzip) body.
        content_type: 'application/json' or 'application/x-protobuf'.
        x_depths_project_id: Optional per-request project override.

    Returns:
        FastAPI Response (empty or partial-success for rejectedLogRecords).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Body AnyValue is stringified by the mapper for LOGS.body.

    lg = _get_logger()
    body = await _read_body(request)
    project_id = _project_from(request, x_depths_project_id)

    try:
        if _want_json(content_type):
            payload = json.loads(body.decode("utf-8") or "{}")
        elif _want_protobuf(content_type):
            if not _PB_OK:
                return _empty_response_for(content_type, "logs")
            req = ExportLogsServiceRequest()
            req.ParseFromString(body)
            payload = MessageToDict(req, preserving_proto_field_name=True)
        else:
            raise HTTPException(status_code=415, detail="unsupported_content_type")
    except Exception:
        raise HTTPException(status_code=400, detail="bad_request_payload")

    res = lg.ingest_otlp_logs_json(payload, project_id=project_id)
    rejected = int(res.get("rejected", 0))
    if rejected > 0:
        return _empty_response_for(
            content_type, "logs", rejected=rejected, error_message="Some log records were rejected due to schema or date checks."
        )
    return _empty_response_for(content_type, "logs")


@app.post("/v1/metrics")
async def post_metrics(
    request: Request,
    content_type: Optional[str] = Header(default="application/json"),
    x_depths_project_id: Optional[str] = Header(default=None),
):
    """
    OTLP/HTTP metrics ingestion endpoint (JSON or protobuf).

    Args:
        request: FastAPI Request used to read (maybe gzip) body.
        content_type: 'application/json' or 'application/x-protobuf'.
        x_depths_project_id: Optional per-request project override.

    Returns:
        FastAPI Response (empty or partial-success for rejectedDataPoints).
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Both Metric Points and Hists are handled by the mapper; producer enforces schema.

    lg = _get_logger()
    body = await _read_body(request)
    project_id = _project_from(request, x_depths_project_id)

    try:
        if _want_json(content_type):
            payload = json.loads(body.decode("utf-8") or "{}")
        elif _want_protobuf(content_type):
            if not _PB_OK:
                return _empty_response_for(content_type, "metrics")
            req = ExportMetricsServiceRequest()
            req.ParseFromString(body)
            payload = MessageToDict(req, preserving_proto_field_name=True)
        else:
            raise HTTPException(status_code=415, detail="unsupported_content_type")
    except Exception:
        raise HTTPException(status_code=400, detail="bad_request_payload")

    res = lg.ingest_otlp_metrics_json(payload, project_id=project_id)
    rejected = int(res.get("rejected", 0))
    if rejected > 0:
        return _empty_response_for(
            content_type, "metrics", rejected=rejected, error_message="Some metric points were rejected due to schema or date checks."
        )
    return _empty_response_for(content_type, "metrics")


@app.get("/api/spans")
async def api_spans(
    date_from: Optional[str] = Query(default=None),
    date_to: Optional[str] = Query(default=None),
    project_id: Optional[str] = Query(default=None),
    service_name: Optional[str] = Query(default=None),
    trace_id: Optional[str] = Query(default=None),
    span_id: Optional[str] = Query(default=None),
    name: Optional[str] = Query(default=None),
    name_like: Optional[str] = Query(default=None),
    status_code: Optional[str] = Query(default=None),
    kind: Optional[str] = Query(default=None),
    time_ms_from: Optional[int] = Query(default=None),
    time_ms_to: Optional[int] = Query(default=None),
    select: Optional[str] = Query(default=None, description="Comma-separated columns"),
    limit: Optional[int] = Query(default=100),
    storage: Optional[str] = Query(default="auto", pattern="^(auto|local|s3)$"),
):
    """
    Query spans with common predicates (lazy scan + pushdown).

    Overview (v0.1.0 role):
        Thin wrapper over DepthsLogger.read_spans(...). Returns JSON rows
        collected from a LazyFrame.

    Returns:
        {'rows': [...], 'count': int}
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Errors during collect surface as HTTP 500 via _collect_json.

    lg = _get_logger()
    lf = lg.read_spans(
        date_from=date_from,
        date_to=date_to,
        project_id=project_id,
        service_name=service_name,
        trace_id=trace_id,
        span_id=span_id,
        name=name,
        name_like=name_like,
        status_code=status_code,
        kind=kind,
        time_ms_from=time_ms_from,
        time_ms_to=time_ms_to,
        select=_csv_list(select),
        limit=limit,
        storage=storage or "auto",
    )
    rows = _collect_json(lf)
    return {"rows": rows, "count": len(rows)}


@app.get("/api/logs")
async def api_logs(
    date_from: Optional[str] = Query(default=None),
    date_to: Optional[str] = Query(default=None),
    project_id: Optional[str] = Query(default=None),
    service_name: Optional[str] = Query(default=None),
    severity_ge: Optional[int] = Query(default=None),
    body_like: Optional[str] = Query(default=None),
    trace_id: Optional[str] = Query(default=None),
    span_id: Optional[str] = Query(default=None),
    time_ms_from: Optional[int] = Query(default=None),
    time_ms_to: Optional[int] = Query(default=None),
    select: Optional[str] = Query(default=None, description="Comma-separated columns"),
    limit: Optional[int] = Query(default=100),
    storage: Optional[str] = Query(default="auto", pattern="^(auto|local|s3)$"),
):
    """
    Query logs with common predicates (lazy scan + pushdown).

    Overview (v0.1.0 role):
        Wrapper over DepthsLogger.read_logs(...). Returns JSON rows
        collected from a LazyFrame.

    
    Returns:
        {'rows': [...], 'count': int}
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - severity_ge maps to 'severity_number >= ...' in the lazy filter.
    
    lg = _get_logger()
    lf = lg.read_logs(
        date_from=date_from,
        date_to=date_to,
        project_id=project_id,
        service_name=service_name,
        severity_ge=severity_ge,
        body_like=body_like,
        trace_id=trace_id,
        span_id=span_id,
        time_ms_from=time_ms_from,
        time_ms_to=time_ms_to,
        select=_csv_list(select),
        limit=limit,
        storage=storage or "auto",
    )
    rows = _collect_json(lf)
    return {"rows": rows, "count": len(rows)}


@app.get("/api/metrics/points")
async def api_metrics_points(
    date_from: Optional[str] = Query(default=None),
    date_to: Optional[str] = Query(default=None),
    project_id: Optional[str] = Query(default=None),
    service_name: Optional[str] = Query(default=None),
    instrument_name: Optional[str] = Query(default=None),
    instrument_type: Optional[str] = Query(default=None),
    time_ms_from: Optional[int] = Query(default=None),
    time_ms_to: Optional[int] = Query(default=None),
    select: Optional[str] = Query(default=None, description="Comma-separated columns"),
    limit: Optional[int] = Query(default=100),
    storage: Optional[str] = Query(default="auto", pattern="^(auto|local|s3)$"),
):
    """
    Query metric points (Gauge/Sum).

    Overview (v0.1.0 role):
        Wrapper over DepthsLogger.read_metrics_points(...).

    Returns:
        {'rows': [...], 'count': int}
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Keep payload sizes small: default limit=100.

    lg = _get_logger()
    lf = lg.read_metrics_points(
        date_from=date_from,
        date_to=date_to,
        project_id=project_id,
        service_name=service_name,
        instrument_name=instrument_name,
        instrument_type=instrument_type,
        time_ms_from=time_ms_from,
        time_ms_to=time_ms_to,
        select=_csv_list(select),
        limit=limit,
        storage=storage or "auto",
    )
    rows = _collect_json(lf)
    return {"rows": rows, "count": len(rows)}


@app.get("/api/metrics/hist")
async def api_metrics_hist(
    date_from: Optional[str] = Query(default=None),
    date_to: Optional[str] = Query(default=None),
    project_id: Optional[str] = Query(default=None),
    service_name: Optional[str] = Query(default=None),
    instrument_name: Optional[str] = Query(default=None),
    instrument_type: Optional[str] = Query(default=None),
    time_ms_from: Optional[int] = Query(default=None),
    time_ms_to: Optional[int] = Query(default=None),
    select: Optional[str] = Query(default=None, description="Comma-separated columns"),
    limit: Optional[int] = Query(default=100),
    storage: Optional[str] = Query(default="auto", pattern="^(auto|local|s3)$"),
):
    """
    Query histogram-like metrics (Histogram/ExpHistogram/Summary).

    Overview (v0.1.0 role):
        Wrapper over DepthsLogger.read_metrics_hist(...).

    Returns:
        {'rows': [...], 'count': int}
    """
    # --- DEVELOPER NOTES -------------------------------------------------
    # - Bounds/counts/quantiles are JSON strings in the table schema.

    lg = _get_logger()
    lf = lg.read_metrics_hist(
        date_from=date_from,
        date_to=date_to,
        project_id=project_id,
        service_name=service_name,
        instrument_name=instrument_name,
        instrument_type=instrument_type,
        time_ms_from=time_ms_from,
        time_ms_to=time_ms_to,
        select=_csv_list(select),
        limit=limit,
        storage=storage or "auto",
    )
    rows = _collect_json(lf)
    return {"rows": rows, "count": len(rows)}
