"""TEI ingest helpers extracted from `hiveos.observe` monolith."""

from datetime import datetime, timezone

from hiveos import config as hive_config
from hiveos.intelligence.tracing import build_trace_event, emit_trace
from hiveos.tei import TEI_VERSION, validate_tei_batch, validate_tei_event


def _iso_from_epoch_ms(epoch_ms):
    return datetime.fromtimestamp(float(epoch_ms) / 1000.0, tz=timezone.utc).isoformat()


def _ensure_tei_run_record(
    run_id,
    *,
    occurred_at_ms,
    source,
    read_run_record,
    persist_run_record,
    active_trace_log_path,
    now_ms,
):
    rid = str(run_id or "").strip()
    if not rid:
        return None
    existing = read_run_record(rid)
    if existing:
        return existing
    started = int(occurred_at_ms or now_ms())
    framework = str((source or {}).get("framework") or "").strip() or "unknown"
    emitter = str((source or {}).get("emitter") or "").strip()
    run_name = f"tei:{framework}" if not emitter else f"tei:{framework}:{emitter}"
    record = {
        "run_id": rid,
        "run_name": run_name,
        "command": "",
        "argv": [],
        "cwd": "",
        "trace_log_path": active_trace_log_path(),
        "archived": False,
        "archived_at_ms": None,
        "archive_reason": None,
        "archive_actor": None,
        "started_at_ms": started,
        "last_update_ms": started,
        "last_output_ms": None,
        "status": "running",
        "exit_code": None,
        "duration_ms": None,
        "source": "tei",
    }
    persist_run_record(record)
    return record


def _to_int_or_none(value):
    try:
        if value is None or isinstance(value, bool):
            return None
        return int(value)
    except Exception:
        return None


def _derive_run_record_updates(*, record, canonical_event_type, normalized_status, normalized_outcome, payload, occurred_at_ms, now_ms):
    updates = {"last_update_ms": int(now_ms())}
    command = str((payload or {}).get("command") or "").strip()
    command_argv_raw = (payload or {}).get("command_argv")
    command_argv = [str(part) for part in command_argv_raw if str(part).strip()] if isinstance(command_argv_raw, list) else []
    existing_command = str((record or {}).get("command") or "").strip()
    existing_argv = list((record or {}).get("argv") or []) if isinstance((record or {}).get("argv"), list) else []
    if (not existing_command) and command:
        updates["command"] = command
    if (not existing_argv) and command_argv:
        updates["argv"] = command_argv
    terminal = str(canonical_event_type or "").strip().lower() in {"turn.finish", "error"}
    if not terminal:
        return updates
    exit_code = _to_int_or_none((payload or {}).get("exit_code"))
    status = str(normalized_status or "").strip().lower()
    outcome = str(normalized_outcome or "").strip().lower()
    if str(canonical_event_type or "").strip().lower() == "error":
        run_status = "failed"
    elif status == "failed" or outcome == "failed":
        run_status = "failed"
    elif exit_code is not None and int(exit_code) != 0:
        run_status = "failed"
    else:
        run_status = "success"
    started_at_ms = int((record or {}).get("started_at_ms") or occurred_at_ms or updates["last_update_ms"])
    duration_ms = _to_int_or_none((payload or {}).get("duration_ms"))
    if duration_ms is None:
        duration_ms = max(0, int(occurred_at_ms) - started_at_ms)
    updates.update(
        {
            "status": run_status,
            "finished_at_ms": int(occurred_at_ms),
            "duration_ms": max(0, int(duration_ms)),
            "exit_code": int(exit_code) if exit_code is not None else (0 if run_status == "success" else 1),
        }
    )
    return updates


def ingest_tei_event(
    event,
    *,
    agent_id="observer.tei",
    persist_db=True,
    persist_file=True,
    strict_version=None,
    now_ms,
    read_run_record,
    persist_run_record,
    touch_run_record,
    active_trace_log_path,
):
    """Validate and ingest one TEI event into the trace stream."""

    if not bool(getattr(hive_config, "TRACE_TEI_INGEST_ENABLED", False)):
        return {
            "ok": False,
            "reason": "tei_ingest_disabled",
            "hint": "Set HIVE_TRACE_TEI_INGEST_ENABLED=true to enable TEI ingestion.",
        }

    strict = bool(getattr(hive_config, "TRACE_TEI_STRICT_VERSION", True))
    if strict_version is not None:
        strict = bool(strict_version)
    validated = validate_tei_event(event, strict_version=strict)
    if not bool(validated.get("ok")):
        return {
            "ok": False,
            "reason": "tei_validation_failed",
            "errors": list(validated.get("errors") or []),
            "strict_version": strict,
        }

    normalized = dict(validated.get("normalized") or {})
    run_id = str(normalized.get("run_id") or "").strip()
    occurred_at_ms = int(normalized.get("occurred_at_ms") or now_ms())
    source = dict(normalized.get("source") or {})
    raw_event_type = str(normalized.get("event_type") or "").strip()
    canonical_event_type = str(normalized.get("canonical_event_type") or "").strip() or raw_event_type
    record = _ensure_tei_run_record(
        run_id,
        occurred_at_ms=occurred_at_ms,
        source=source,
        read_run_record=read_run_record,
        persist_run_record=persist_run_record,
        active_trace_log_path=active_trace_log_path,
        now_ms=now_ms,
    )
    payload = dict(normalized.get("payload") or {})
    payload.setdefault("source", "tei")
    run_updates = _derive_run_record_updates(
        record=record,
        canonical_event_type=canonical_event_type,
        normalized_status=str(normalized.get("status") or "unknown").strip() or "unknown",
        normalized_outcome=str(normalized.get("outcome") or "").strip(),
        payload=payload,
        occurred_at_ms=occurred_at_ms,
        now_ms=now_ms,
    )
    touch_run_record(run_id, **run_updates)

    metadata = dict(normalized.get("metadata") or {})
    tags = list(normalized.get("tags") or [])
    status = str(normalized.get("status") or "unknown").strip() or "unknown"
    snapshot = {
        "input_ref": str(payload.get("input_ref") or "").strip() or None,
        "output_ref": str(payload.get("output_ref") or "").strip() or None,
        "state_ref": str(payload.get("state_ref") or "").strip() or None,
        "payload_hash": str(payload.get("payload_hash") or "").strip() or None,
    }
    snapshot = {key: value for key, value in snapshot.items() if value}
    anchor = None
    checkpoint_id = str(normalized.get("checkpoint_id") or "").strip()
    step_id = str(normalized.get("step_id") or "").strip()
    if checkpoint_id:
        anchor = {"anchor_id": checkpoint_id, "anchor_type": "checkpoint.state"}
    elif canonical_event_type in {"tool.request", "tool.result", "decision", "response.commit"} and step_id:
        anchor_type = {
            "tool.request": "tool.request",
            "tool.result": "tool.result",
            "decision": "decision.branch",
            "response.commit": "output.commit",
        }.get(canonical_event_type, "step.generic")
        anchor = {"anchor_id": step_id, "anchor_type": anchor_type}
    source_block = {
        "runtime": str(source.get("framework") or "").strip() or "unknown",
        "shim": str(source.get("emitter") or "").strip() or None,
        "shim_version": str(metadata.get("shim_version") or "").strip() or None,
    }
    source_block = {key: value for key, value in source_block.items() if value}
    contract_v1 = {
        "contract_version": "1.0",
        "event_id": str(normalized.get("event_id") or ""),
        "event_type": canonical_event_type,
        "timestamp": _iso_from_epoch_ms(occurred_at_ms),
        "run_id": run_id,
        "step_id": step_id or None,
        "parent_step_id": str(normalized.get("parent_step_id") or "").strip() or None,
        "status": status,
        "source": source_block,
    }
    if anchor:
        contract_v1["anchor"] = anchor
    if snapshot:
        contract_v1["snapshot"] = snapshot
    contract_v1["meta"] = {
        "provider": str(metadata.get("provider") or "").strip() or None,
        "model": str(metadata.get("model") or "").strip() or None,
        "tags": tags,
        "ext": {
            "tei_raw_event_type": raw_event_type,
            "tei_event_type_is_canonical": bool(normalized.get("event_type_is_canonical")),
        },
    }
    contract_v1["meta"] = {
        key: value
        for key, value in contract_v1["meta"].items()
        if value not in (None, "", [], {})
    }
    if not contract_v1.get("meta"):
        contract_v1.pop("meta", None)
    payload["tei"] = {
        "tei_version": str(normalized.get("tei_version") or TEI_VERSION),
        "event_id": str(normalized.get("event_id") or ""),
        "occurred_at_ms": occurred_at_ms,
        "source": source,
        "metadata": metadata,
        "tags": tags,
        "canonical_event_type": canonical_event_type,
        "status": status,
    }
    payload["contract_v1"] = contract_v1

    ingested = build_trace_event(
        event_type=canonical_event_type or "tei.event",
        run_id=run_id,
        outcome=str(normalized.get("outcome") or "success"),
        payload=payload,
        agent_id=str(agent_id or "observer.tei"),
        flow_id=str(normalized.get("flow_id") or "") or None,
        step_id=str(normalized.get("step_id") or "") or None,
        parent_step_id=str(normalized.get("parent_step_id") or "") or None,
        tool_call_id=str(normalized.get("tool_call_id") or "") or None,
        attempt=normalized.get("attempt"),
        step_type=str(normalized.get("step_type") or "") or None,
        timestamp=_iso_from_epoch_ms(occurred_at_ms),
    )
    emit_trace(ingested, persist_db=bool(persist_db), persist_file=bool(persist_file))
    return {
        "ok": True,
        "run_id": run_id,
        "event_id": str(normalized.get("event_id") or ""),
        "event_type": canonical_event_type,
        "strict_version": strict,
    }


def ingest_tei_events(
    events,
    *,
    agent_id="observer.tei",
    persist_db=True,
    persist_file=True,
    strict_version=None,
    now_ms,
    read_run_record,
    persist_run_record,
    touch_run_record,
    active_trace_log_path,
):
    """Validate and ingest a TEI batch. Returns aggregate diagnostics."""

    strict = bool(getattr(hive_config, "TRACE_TEI_STRICT_VERSION", True))
    if strict_version is not None:
        strict = bool(strict_version)
    batch = validate_tei_batch(events, strict_version=strict)
    if not bool(getattr(hive_config, "TRACE_TEI_INGEST_ENABLED", False)):
        return {
            "ok": False,
            "reason": "tei_ingest_disabled",
            "total": int(batch.get("total") or 0),
            "valid": int(batch.get("valid") or 0),
            "invalid": int(batch.get("invalid") or 0),
            "items": list(batch.get("items") or []),
        }

    results = []
    for item in list(batch.get("items") or []):
        index = int(item.get("index") or 0)
        raw_event = events[index] if isinstance(events, list) and 0 <= index < len(events) else None
        if not bool(item.get("ok")):
            results.append({"index": index, "ok": False, "errors": list(item.get("errors") or [])})
            continue
        ingested = ingest_tei_event(
            raw_event,
            agent_id=agent_id,
            persist_db=persist_db,
            persist_file=persist_file,
            strict_version=strict,
            now_ms=now_ms,
            read_run_record=read_run_record,
            persist_run_record=persist_run_record,
            touch_run_record=touch_run_record,
            active_trace_log_path=active_trace_log_path,
        )
        results.append({"index": index, **ingested})
    success_count = sum(1 for item in results if bool(item.get("ok")))
    return {
        "ok": success_count == len(results),
        "total": len(results),
        "ingested": success_count,
        "failed": max(0, len(results) - success_count),
        "items": results,
        "strict_version": strict,
    }

