#!/usr/bin/env python3
"""
CI Trace Report Script - Reports CI Trace data to AgentLoop from Gitea Actions workflows.

Usage in CI workflow jobs:
  - At start: python3 scripts/ci/ci_trace_report.py --status running
  - At end: python3 scripts/ci/ci_trace_report.py --status ok --start-time $CI_TRACE_START_TIME

Environment variables (built-in Gitea Actions):
  GITEA_REPOSITORY / GITHUB_REPOSITORY - repository (owner/repo)
  GITEA_WORKFLOW / GITHUB_WORKFLOW     - workflow name
  GITEA_JOB / GITHUB_JOB               - job ID
  GITEA_SHA / GITHUB_SHA               - commit SHA
  GITEA_REF_NAME / GITHUB_REF_NAME     - branch name
  GITEA_RUN_ID / GITHUB_RUN_ID         - run ID
  GITEA_ACTOR / GITHUB_ACTOR           - trigger actor
  GITEA_EVENT_NAME / GITHUB_EVENT_NAME - event type
  PR_NUMBER / GITEA_PR_NUMBER          - PR number (if PR triggered)

AgentLoop configuration (injected via Secrets):
  AGENTLOOP_LICENSE_KEY - LicenseKey (required)
  AGENTLOOP_ENDPOINT    - Trace endpoint (optional, has default)
  AGENTLOOP_PROJECT     - SLS Project name (optional)
  AGENTLOOP_WORKSPACE   - CMS Workspace name (optional)
"""

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
import uuid

# ========== Default Configuration ==========
DEFAULT_ENDPOINT = "https://proj-xtrace-495e81719a1fd9a2c5fd671eefafbe-cn-hangzhou.cn-hangzhou.log.aliyuncs.com/apm/trace/opentelemetry/v1/traces"
DEFAULT_PROJECT = "proj-xtrace-495e81719a1fd9a2c5fd671eefafbe-cn-hangzhou"
DEFAULT_WORKSPACE = "agentloop-13b8d6efb7fde6e9b193eb982ade68e2"


# ========== OTLP Protobuf Manual Encoding ==========


def _encode_varint(value):
    result = bytearray()
    while value > 0x7F:
        result.append((value & 0x7F) | 0x80)
        value >>= 7
    result.append(value & 0x7F)
    return bytes(result)


def _encode_tag(field_number, wire_type):
    return _encode_varint((field_number << 3) | wire_type)


def _encode_string_field(field_number, value):
    value_bytes = value.encode("utf-8")
    return _encode_tag(field_number, 2) + _encode_varint(len(value_bytes)) + value_bytes


def _encode_bytes_field(field_number, value_bytes):
    return _encode_tag(field_number, 2) + _encode_varint(len(value_bytes)) + value_bytes


def _encode_int_field(field_number, value):
    return _encode_tag(field_number, 0) + _encode_varint(value & 0xFFFFFFFFFFFFFFFF)


def _encode_message_field(field_number, message_bytes):
    return _encode_tag(field_number, 2) + _encode_varint(len(message_bytes)) + message_bytes


def _encode_key_value(key, value_str):
    any_value = _encode_string_field(1, value_str)
    return _encode_string_field(1, key) + _encode_message_field(2, any_value)


def _encode_status(status_code, status_msg=""):
    data = _encode_int_field(1, status_code)
    if status_msg:
        data += _encode_string_field(2, status_msg)
    return data


def _encode_span(
    trace_id_bytes,
    span_id_bytes,
    parent_span_id_bytes,
    name,
    start_time_unix_nano,
    end_time_unix_nano,
    span_kind,
    attributes,
    status_code,
    status_msg="",
):
    data = b""
    data += _encode_bytes_field(1, trace_id_bytes)
    data += _encode_bytes_field(2, span_id_bytes)
    if parent_span_id_bytes:
        data += _encode_bytes_field(3, parent_span_id_bytes)
    data += _encode_string_field(4, name)
    data += _encode_int_field(5, span_kind)
    data += _encode_int_field(6, start_time_unix_nano)
    data += _encode_int_field(7, end_time_unix_nano)
    for key, value in attributes.items():
        kv = _encode_key_value(key, str(value))
        data += _encode_message_field(9, kv)
    status = _encode_status(status_code, status_msg)
    data += _encode_message_field(12, status)
    return data


def _encode_resource_spans(service_name, scope_spans_bytes):
    svc_kv = _encode_key_value("service.name", service_name)
    resource = _encode_message_field(1, svc_kv)
    data = _encode_message_field(1, resource)
    data += _encode_message_field(2, scope_spans_bytes)
    return data


def _encode_scope_spans(scope_name, spans_bytes_list):
    scope = _encode_string_field(1, scope_name)
    data = _encode_message_field(1, scope)
    for span_bytes in spans_bytes_list:
        data += _encode_message_field(2, span_bytes)
    return data


def _encode_traces_data(resource_spans_bytes_list):
    data = b""
    for rs_bytes in resource_spans_bytes_list:
        data += _encode_message_field(1, rs_bytes)
    return data


# ========== Helper Functions ==========


def _gen_trace_id():
    return uuid.uuid4().bytes


def _gen_span_id():
    return uuid.uuid4().bytes[:8]


def _env(name, default=""):
    """Get env var with GITEA_/GITHUB_ prefix fallback."""
    val = os.getenv(name, "")
    if val:
        return val
    if name.startswith("GITEA_"):
        alt = "GITHUB_" + name[6:]
        return os.getenv(alt, default)
    if name.startswith("GITHUB_"):
        alt = "GITEA_" + name[7:]
        return os.getenv(alt, default)
    return default


def _get_pr_number():
    """Get PR number from environment or event file."""
    pr = os.getenv("PR_NUMBER", "") or os.getenv("GITEA_PR_NUMBER", "")
    if pr:
        return pr

    event_path = os.getenv("GITHUB_EVENT_PATH", "") or os.getenv("GITEA_EVENT_PATH", "")
    if event_path and os.path.isfile(event_path):
        try:
            with open(event_path, "r") as f:
                event = json.load(f)
            if "pull_request" in event and "number" in event["pull_request"]:
                return str(event["pull_request"]["number"])
        except Exception:
            pass

    return ""


def _get_ci_attributes():
    """Collect attributes from CI environment variables."""
    attrs = {
        "ci.repo": _env("GITEA_REPOSITORY") or _env("GITHUB_REPOSITORY") or "unknown",
        "ci.workflow": _env("GITEA_WORKFLOW") or _env("GITHUB_WORKFLOW") or "unknown",
        "ci.job": _env("GITEA_JOB") or _env("GITHUB_JOB") or "unknown",
        "ci.commit_sha": _env("GITEA_SHA") or _env("GITHUB_SHA") or "unknown",
        "ci.branch": _env("GITEA_REF_NAME") or _env("GITHUB_REF_NAME") or "unknown",
        "ci.run_id": _env("GITEA_RUN_ID") or _env("GITHUB_RUN_ID") or "unknown",
        "ci.actor": _env("GITEA_ACTOR") or _env("GITHUB_ACTOR") or "unknown",
        "ci.event": _env("GITEA_EVENT_NAME") or _env("GITHUB_EVENT_NAME") or "unknown",
    }
    pr = _get_pr_number()
    if pr:
        attrs["ci.pr_number"] = pr
    return attrs


# ========== Trace Building & Reporting ==========


def build_trace(service_name, trace_name, status, duration_ms, attributes=None):
    """Build an OTLP trace payload (protobuf bytes). No external dependencies."""
    trace_id = _gen_trace_id()
    end_time = int(time.time() * 1e9)
    start_time = end_time - int(duration_ms * 1e6)
    status_code = 1 if status in ("ok", "running") else 2
    status_msg = "" if status in ("ok", "running") else "Job failed"

    main_attrs = {
        "agent.trace_name": trace_name,
        "agent.service": service_name,
        "ci.trace_status": status,
    }
    if attributes:
        main_attrs.update(attributes)

    main_span = _encode_span(
        trace_id_bytes=trace_id,
        span_id_bytes=_gen_span_id(),
        parent_span_id_bytes=b"",
        name=trace_name,
        start_time_unix_nano=start_time,
        end_time_unix_nano=end_time,
        span_kind=1,
        attributes=main_attrs,
        status_code=status_code,
        status_msg=status_msg,
    )

    scope_spans = _encode_scope_spans("ci-trace", [main_span])
    resource_spans = _encode_resource_spans(service_name, scope_spans)
    return _encode_traces_data([resource_spans])


def report_ci_trace(
    service_name,
    trace_name,
    status="ok",
    duration_ms=1000,
    endpoint=None,
    license_key=None,
    project=None,
    workspace=None,
    extra_attributes=None,
):
    """
    Report CI Trace data. Returns (success: bool, message: str).
    Never raises exceptions; returns False on failure.
    """
    try:
        endpoint = endpoint or os.getenv("AGENTLOOP_ENDPOINT", DEFAULT_ENDPOINT)
        license_key = license_key or os.getenv("AGENTLOOP_LICENSE_KEY", "")
        project = project or os.getenv("AGENTLOOP_PROJECT", DEFAULT_PROJECT)
        workspace = workspace or os.getenv("AGENTLOOP_WORKSPACE", DEFAULT_WORKSPACE)

        if not license_key:
            return False, "[Trace] skipped: AGENTLOOP_LICENSE_KEY not configured"

        attrs = _get_ci_attributes()
        if extra_attributes:
            attrs.update(extra_attributes)

        payload = build_trace(
            service_name=service_name,
            trace_name=trace_name,
            status=status,
            duration_ms=duration_ms,
            attributes=attrs,
        )

        headers = {
            "Content-Type": "application/x-protobuf",
            "x-arms-license-key": license_key,
            "x-arms-project": project,
            "x-cms-workspace": workspace,
        }

        req = urllib.request.Request(endpoint, data=payload, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                status_code = resp.status
                resp_body = resp.read().decode("utf-8", errors="replace")
        except urllib.error.HTTPError as e:
            status_code = e.code
            resp_body = e.read().decode("utf-8", errors="replace")

        if status_code in (200, 202):
            return True, (f"[Trace] success: {service_name} / {trace_name} " f"({status}, {duration_ms}ms)")
        else:
            return False, (f"[Trace] failed: HTTP {status_code} - {resp_body[:200]}")
    except Exception as e:
        return False, f"[Trace] error: {type(e).__name__}: {str(e)}"


def main():
    parser = argparse.ArgumentParser(description="CI AgentLoop Trace Reporter")
    parser.add_argument(
        "--service",
        dest="service_name",
        default=os.getenv("TRACE_SERVICE", ""),
        help="Service name (also via TRACE_SERVICE env)",
    )
    parser.add_argument(
        "--name",
        dest="trace_name",
        default=os.getenv("TRACE_NAME", ""),
        help="Trace name (also via TRACE_NAME env)",
    )
    parser.add_argument(
        "--status",
        default=os.getenv("TRACE_STATUS", "ok"),
        choices=["ok", "error", "running"],
        help="Status: ok / error / running (default ok)",
    )
    parser.add_argument(
        "--start-time",
        dest="start_time",
        default=os.getenv("TRACE_START_TIME", ""),
        help="Start timestamp (seconds) for duration calculation",
    )
    parser.add_argument(
        "--duration-ms",
        dest="duration_ms",
        type=int,
        default=0,
        help="Direct duration in ms; takes precedence over --start-time",
    )
    parser.add_argument("--attrs", default="", help="Extra attributes (JSON string)")

    args = parser.parse_args()

    if not args.service_name:
        print("[Trace] skipped: no service specified (--service or TRACE_SERVICE)")
        sys.exit(0)

    duration_ms = args.duration_ms
    if duration_ms <= 0 and args.start_time:
        try:
            start_ts = float(args.start_time)
            duration_ms = int((time.time() - start_ts) * 1000)
        except (ValueError, TypeError):
            duration_ms = 1000
    if duration_ms <= 0:
        duration_ms = 1000

    extra_attrs = {}
    if args.attrs:
        try:
            extra_attrs = json.loads(args.attrs)
        except json.JSONDecodeError:
            pass

    trace_name = args.trace_name
    if not trace_name:
        wf = _env("GITEA_WORKFLOW") or _env("GITHUB_WORKFLOW") or "CI"
        job = _env("GITEA_JOB") or _env("GITHUB_JOB") or "job"
        trace_name = f"{wf} / {job}"

    success, msg = report_ci_trace(
        service_name=args.service_name,
        trace_name=trace_name,
        status=args.status,
        duration_ms=duration_ms,
        extra_attributes=extra_attrs,
    )

    print(msg)
    sys.exit(0)


if __name__ == "__main__":
    main()