from __future__ import annotations

import contextlib
import contextvars
import json
from typing import Any

from .trace_config import TraceConfig

_FIELD_ORDER = [
    "tick",
    "lang",
    "scenario",
    "region",
    "layer_index",
    "layer_name",
    "source_layer_index",
    "dest_layer_index",
    "neuron_id",
    "source_neuron_id",
    "target_neuron_id",
    "slot_id",
    "signal_id",
    "input_port",
    "input_value",
    "output_value",
    "fallback",
    "created_slot",
    "created_neuron",
    "created_layer",
    "reason",
    "hop",
    "desired_slot",
    "chosen_slot",
    "slots_before",
    "slots_after",
    "slot_limit",
    "neuron_count",
    "strength_before",
    "strength_after",
    "modulation",
    "hit_count",
    "theta_before",
    "theta_after",
    "fired",
    "delivered_events",
    "total_slots",
    "total_synapses",
    "step",
    "focus_policy",
    "candidate_count",
    "selected_rank",
    "focus_point",
    "focus_value",
    "previous_focus_point",
    "anchor_before",
    "anchor_before_value",
    "anchor_after",
    "anchor_after_value",
    "distance_from_anchor",
    "candidate_energy",
    "candidate_novelty",
    "candidate_familiarity",
    "inhibited",
]

_config: TraceConfig = TraceConfig()
_writer = None
_tick_context: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar("grownet_tick_context", default={})
_tick_counter = 0


def configure(trace_config: TraceConfig, writer) -> None:
    global _config, _writer, _tick_counter
    _config = trace_config
    _writer = writer
    _tick_counter = 0
    _tick_context.set({})


def reset() -> None:
    global _tick_counter
    _tick_counter = 0
    _tick_context.set({})


def is_enabled() -> bool:
    return bool(_config.trace_enabled)


def begin_tick(region_name: str, input_port: str | None, input_value: float | None, use_default_signal_id: bool) -> int:
    global _tick_counter
    _tick_counter += 1
    signal_id = format_tick_signal_id(_tick_counter) if use_default_signal_id else None
    _tick_context.set(
        {
            "tick": _tick_counter,
            "region": region_name,
            "signal_id": signal_id,
            "hop": 0,
        }
    )
    event_fields: dict[str, Any] = {}
    if input_port is not None:
        event_fields["input_port"] = input_port
    if input_value is not None:
        event_fields["input_value"] = input_value
    emit("tick.start", **event_fields)
    return _tick_counter


def end_tick(metrics) -> None:
    if metrics is None:
        _tick_context.set({})
        return
    emit(
        "tick.end",
        delivered_events=int(getattr(metrics, "delivered_events", 0)),
        total_slots=int(getattr(metrics, "total_slots", 0)),
        total_synapses=int(getattr(metrics, "total_synapses", 0)),
    )
    _tick_context.set({})


@contextlib.contextmanager
def run_with_signal(signal_id: str):
    context_state = dict(_tick_context.get() or {})
    context_state["signal_id"] = signal_id
    context_state["hop"] = 0
    token = _tick_context.set(context_state)
    try:
        yield
    finally:
        _tick_context.reset(token)


@contextlib.contextmanager
def run_with_hop_increment():
    context_state = dict(_tick_context.get() or {})
    context_state["hop"] = int(context_state.get("hop", 0)) + 1
    token = _tick_context.set(context_state)
    try:
        yield
    finally:
        _tick_context.reset(token)


def current_tick() -> int:
    return int((_tick_context.get() or {}).get("tick", 0))


def current_signal_id() -> str | None:
    return (_tick_context.get() or {}).get("signal_id")


def current_hop() -> int:
    return int((_tick_context.get() or {}).get("hop", 0))


def format_tick_signal_id(tick_number: int) -> str:
    return f"tick-{tick_number:06d}"


def format_pixel_signal_id(tick_number: int, row_index: int, col_index: int) -> str:
    return f"tick-{tick_number:06d}:px-{row_index:04d}-{col_index:04d}"


def add_layer_fields(event_fields: dict[str, Any], layer_obj) -> None:
    region_obj = getattr(layer_obj, "region", None)
    if event_fields is None or layer_obj is None or region_obj is None:
        return
    try:
        layer_index = region_obj.get_layers().index(layer_obj)
    except Exception:
        return
    event_fields["layer_index"] = layer_index
    event_fields["layer_name"] = f"layer_{layer_index}"


def add_neuron_fields(event_fields: dict[str, Any], neuron_obj) -> None:
    if event_fields is None or neuron_obj is None:
        return
    event_fields["neuron_id"] = getattr(neuron_obj, "id", None)
    add_layer_fields(event_fields, getattr(neuron_obj, "owner", None))


def emit_bus_decay(layer_obj) -> None:
    if layer_obj is None:
        return
    event_fields: dict[str, Any] = {}
    add_layer_fields(event_fields, layer_obj)
    try:
        event_fields["step"] = int(layer_obj.get_bus().get_current_step())
    except Exception:
        event_fields["step"] = 0
    emit("bus.decay", **event_fields)


def emit_layer_forward(layer_obj, input_value: float) -> None:
    if layer_obj is None:
        return
    event_fields: dict[str, Any] = {}
    add_layer_fields(event_fields, layer_obj)
    event_fields["neuron_count"] = len(layer_obj.get_neurons())
    event_fields["input_value"] = float(input_value)
    emit("layer.forward", **event_fields)


def emit(event_name: str, **raw_fields: Any) -> None:
    if not is_enabled():
        return
    context_state = dict(_tick_context.get() or {})
    if int(context_state.get("tick", 0)) <= 0 or _writer is None:
        return
    event_fields: dict[str, Any] = {
        "event": event_name,
        "tick": int(context_state.get("tick", 0)),
        "lang": _config.lang,
        "scenario": _config.scenario,
        "region": context_state.get("region"),
    }
    for field_name, field_value in raw_fields.items():
        if field_value is not None:
            event_fields[field_name] = field_value
    if "signal_id" not in event_fields and context_state.get("signal_id") is not None:
        event_fields["signal_id"] = context_state.get("signal_id")
    if "hop" not in event_fields and context_state.get("signal_id") is not None:
        event_fields["hop"] = int(context_state.get("hop", 0))
    if not _passes_filters(event_name, event_fields):
        return
    try:
        if _config.trace_format == "jsonl":
            line = _format_json_line(event_fields)
        else:
            line = _format_text_line(event_fields)
        _writer.write(line + "\n")
        _writer.flush()
    except Exception:
        return


def _passes_filters(event_name: str, event_fields: dict[str, Any]) -> bool:
    if _config.trace_families:
        if not _matches_event_filter(event_name, _config.trace_families):
            return False
    if _config.trace_neuron is not None:
        neuron_id = _first_present(event_fields, ["neuron_id", "source_neuron_id", "target_neuron_id"])
        if neuron_id is None or str(neuron_id) != _config.trace_neuron:
            return False
    if _config.trace_layer is not None:
        layer_index = _first_present(event_fields, ["layer_index", "source_layer_index", "dest_layer_index"])
        if layer_index is None or int(layer_index) != int(_config.trace_layer):
            return False
    return True


def _matches_event_filter(event_name: str, filter_tokens: set[str]) -> bool:
    normalized_event_name = str(event_name).lower()
    family_name = normalized_event_name.split(".", 1)[0]
    for filter_token in filter_tokens:
        normalized_filter_token = str(filter_token).lower()
        if normalized_filter_token == normalized_event_name or normalized_filter_token == family_name:
            return True
    return False


def _first_present(event_fields: dict[str, Any], field_names: list[str]) -> Any:
    for field_name in field_names:
        if event_fields.get(field_name) is not None:
            return event_fields.get(field_name)
    return None


def _ordered_keys(event_fields: dict[str, Any]) -> list[str]:
    ordered = [field_name for field_name in _FIELD_ORDER if field_name in event_fields]
    additional = sorted(field_name for field_name in event_fields if field_name not in _FIELD_ORDER and field_name != "event")
    return ["event"] + ordered + additional


def _format_text_line(event_fields: dict[str, Any]) -> str:
    parts = [f"[{event_fields['event']}]"]
    for field_name in _ordered_keys(event_fields):
        if field_name == "event":
            continue
        field_value = event_fields.get(field_name)
        if field_value is None:
            continue
        parts.append(f"{field_name}={_format_scalar(field_value)}")
    return " ".join(parts)


def _format_json_line(event_fields: dict[str, Any]) -> str:
    ordered_payload = {
        field_name: event_fields[field_name]
        for field_name in _ordered_keys(event_fields)
        if event_fields.get(field_name) is not None
    }
    return json.dumps(ordered_payload, separators=(",", ":"))


def _format_scalar(field_value: Any) -> str:
    if isinstance(field_value, float):
        return f"{field_value:.3f}"
    return str(field_value)
