"""SpanExporter that writes human-readable Markdown files, incrementally.

Each span is appended to the .md file as it arrives so the file can be
followed live (``tail -f``).  On ``force_flush()`` / ``shutdown()`` the
file is rewritten with a summary header (cost table, totals) prepended.
"""

import base64
import os
import re
from collections import defaultdict
from datetime import datetime, timezone
from typing import Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult

from minitrail._pricing import compute_cost as _cost


# ---------------------------------------------------------------------------
# Attribute helpers — support both OpenInference (llm.*) and GenAI (gen_ai.*)
# ---------------------------------------------------------------------------

def _get_model(attrs: dict) -> str | None:
    """Extract model name from OpenInference or GenAI attributes."""
    return attrs.get("llm.model_name") or attrs.get("gen_ai.request.model") or None


def _get_tokens(attrs: dict) -> tuple[int, int]:
    """Extract (input, output) token counts from OpenInference or GenAI attributes."""
    inp = int(
        attrs.get("llm.token_count.prompt")
        or attrs.get("gen_ai.usage.input_tokens")
        or attrs.get("gen_ai.usage.prompt_tokens")
        or 0
    )
    out = int(
        attrs.get("llm.token_count.completion")
        or attrs.get("gen_ai.usage.output_tokens")
        or attrs.get("gen_ai.usage.completion_tokens")
        or 0
    )
    return inp, out


def _is_llm_span(attrs: dict, name: str) -> bool:
    """Detect LLM spans from either convention."""
    kind = attrs.get("openinference.span.kind", "")
    if kind == "LLM":
        return True
    op = attrs.get("gen_ai.operation.name", "")
    if op == "chat":
        return True
    low = name.lower()
    return "llm" in low or "chat" in low


def _get_genai_messages(span) -> tuple[list[dict], list[dict]]:
    """Extract input/output messages from GenAI semantic convention events."""
    input_msgs: list[dict] = []
    output_msgs: list[dict] = []

    for event in (span.events or []):
        event_name = event.name if hasattr(event, "name") else ""
        event_attrs = {}
        if hasattr(event, "attributes") and event.attributes:
            event_attrs = dict(event.attributes)

        content_str = event_attrs.get("content", "")

        if event_name == "gen_ai.user.message":
            text = _parse_genai_content(content_str)
            if text:
                input_msgs.append({"role": "user", "content": [{"type": "text", "text": text}]})
        elif event_name == "gen_ai.system.message":
            text = _parse_genai_content(content_str)
            if text:
                input_msgs.append({"role": "system", "content": [{"type": "text", "text": text}]})
        elif event_name == "gen_ai.choice":
            text = event_attrs.get("message", "")
            text = _parse_genai_content(text) or str(text)
            if text:
                output_msgs.append({"role": "assistant", "content": [{"type": "text", "text": text}]})

    return input_msgs, output_msgs


def _parse_genai_content(raw: str) -> str:
    """Parse GenAI event content — may be a JSON array of {text: ...} blocks."""
    if not raw:
        return ""
    try:
        import json
        parsed = json.loads(raw)
        if isinstance(parsed, list):
            parts = []
            for item in parsed:
                if isinstance(item, dict) and "text" in item:
                    parts.append(item["text"])
                elif isinstance(item, str):
                    parts.append(item)
            return "\n".join(parts) if parts else raw
    except (json.JSONDecodeError, TypeError):
        pass
    return raw


def _fmt_cost(usd: float) -> str:
    return f"${usd:.4f}" if usd < 0.01 else f"${usd:.2f}"


def _fmt_dur(ns: int) -> str:
    """Nanosecond delta → human string."""
    ms = ns / 1_000_000
    if ms < 1000:
        return f"{ms:.0f}ms"
    s = ms / 1000
    if s < 60:
        return f"{s:.1f}s"
    return f"{s / 60:.1f}m"


# ---------------------------------------------------------------------------
# Reassemble OpenInference flat attributes into chat messages
# ---------------------------------------------------------------------------

def _reassemble_messages(attrs: dict, prefix: str) -> list[dict]:
    """Return [{role, content: str | [{type, text?, url?}]}]."""
    dot = prefix + "."
    entries = [(k[len(dot):], v) for k, v in attrs.items() if k.startswith(dot)]
    if not entries:
        return []

    by_idx: dict[int, dict[str, object]] = defaultdict(dict)
    for suffix, val in entries:
        parts = suffix.split(".", 1)
        if len(parts) < 2:
            continue
        try:
            idx = int(parts[0])
        except ValueError:
            continue
        by_idx[idx][parts[1]] = val

    messages: list[dict] = []
    for idx in sorted(by_idx):
        fields = by_idx[idx]
        role = str(fields.get("message.role", "unknown"))
        blocks: list[dict] = []

        simple = fields.get("message.content")
        if isinstance(simple, str) and simple:
            blocks.append({"type": "text", "text": simple})

        # Multimodal content blocks
        content_fields: dict[int, dict[str, object]] = defaultdict(dict)
        for k, v in fields.items():
            m = re.match(r"^message\.contents\.(\d+)\.message_content\.(.+)$", k)
            if m:
                content_fields[int(m.group(1))][m.group(2)] = v

        for cidx in sorted(content_fields):
            cf = content_fields[cidx]
            ctype = cf.get("type")
            if ctype == "text":
                t = cf.get("text")
                if isinstance(t, str):
                    blocks.append({"type": "text", "text": t})
            elif ctype == "image":
                url = cf.get("image.image.url")
                if isinstance(url, str):
                    blocks.append({"type": "image", "url": url})

        messages.append({"role": role, "content": blocks})

    return messages


_MIME_TO_EXT = {
    "image/png": "png",
    "image/jpeg": "jpg",
    "image/jpg": "jpg",
    "image/gif": "gif",
    "image/webp": "webp",
    "image/jp2": "jp2",
}


def _save_image(data_uri: str, images_dir: str, prefix: str, index: int) -> str | None:
    """Decode a ``data:image/…;base64,…`` URI, save to *images_dir*, return relative path."""
    m = re.match(r"^data:(image/[a-z0-9+]+);base64,(.+)$", data_uri, re.IGNORECASE | re.DOTALL)
    if not m:
        return None
    mime, b64 = m.group(1), m.group(2)
    ext = _MIME_TO_EXT.get(mime.lower(), "png")
    os.makedirs(images_dir, exist_ok=True)
    filename = f"{prefix}_{index}.{ext}"
    filepath = os.path.join(images_dir, filename)
    try:
        with open(filepath, "wb") as f:
            f.write(base64.b64decode(b64))
    except Exception:
        return None
    # Return path relative to the markdown file (which sits in logs_dir)
    return f"images/{filename}"


def _render_messages(messages: list[dict], images_dir: str, img_prefix: str, img_counter: list[int]) -> str:
    """Render a list of chat messages as markdown, saving images to disk."""
    parts: list[str] = []
    for msg in messages:
        role = msg["role"].upper()
        parts.append(f"**{role}**\n")
        for block in msg.get("content", []):
            if block["type"] == "text":
                lines = block["text"].splitlines()
                quoted = "\n".join(f"> {l}" for l in lines)
                parts.append(quoted + "\n")
            elif block["type"] == "image":
                rel = _save_image(block["url"], images_dir, img_prefix, img_counter[0])
                img_counter[0] += 1
                if rel:
                    parts.append(f"![image]({rel})\n")
                else:
                    parts.append("> *(image — could not decode)*\n")
        parts.append("")
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# smolagents helper — parse input.value JSON with ChatMessage repr strings
# ---------------------------------------------------------------------------

def _parse_smolagents_input_value(raw: str) -> list[dict]:
    """Parse smolagents input.value to extract chat messages from ChatMessage repr strings."""
    import ast
    import json as _json
    if not raw or not isinstance(raw, str):
        return []
    try:
        parsed = _json.loads(raw)
    except (_json.JSONDecodeError, TypeError):
        return []
    if not isinstance(parsed, dict) or "messages" not in parsed:
        return []
    messages: list[dict] = []
    for msg_str in parsed["messages"]:
        if not isinstance(msg_str, str):
            continue
        role_match = re.search(r"role=<MessageRole\.\w+:\s*'(\w+)'>", msg_str)
        if not role_match:
            continue
        role = role_match.group(1)
        content_match = re.search(r"content=(\[.+?\])(?:,\s*tool_calls)", msg_str, re.DOTALL)
        text = ""
        if content_match:
            try:
                content_list = ast.literal_eval(content_match.group(1))
                text_parts = [
                    item["text"] for item in content_list
                    if isinstance(item, dict) and "text" in item
                ]
                text = "\n".join(text_parts)
            except Exception:
                text = msg_str
        else:
            text = msg_str
        if text:
            messages.append({"role": role, "content": [{"type": "text", "text": text}]})
    return messages


# ---------------------------------------------------------------------------
# CrewAI helper — parse output.value JSON for AGENT spans
# ---------------------------------------------------------------------------

def _parse_crewai_output_value(attrs: dict) -> list[dict]:
    """Parse CrewAI AGENT span output.value to extract chat messages."""
    import json
    raw_val = attrs.get("output.value", "")
    if not raw_val or not isinstance(raw_val, str):
        return []
    try:
        parsed = json.loads(raw_val)
    except (json.JSONDecodeError, TypeError):
        return []
    if not isinstance(parsed, dict) or "messages" not in parsed:
        return []
    messages: list[dict] = []
    for msg in parsed["messages"]:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        messages.append({"role": role, "content": [{"type": "text", "text": content}]})
    return messages


# ---------------------------------------------------------------------------
# Format a single span as markdown (body section)
# ---------------------------------------------------------------------------

def _span_body(span: ReadableSpan, images_dir: str, img_counter: list[int]) -> str:
    attrs = dict(span.attributes or {})
    name = span.name
    span_id = format(span.context.span_id, "016x")
    dur_ns = span.end_time - span.start_time if span.end_time and span.start_time else 0
    dur_str = _fmt_dur(dur_ns)

    model = _get_model(attrs)
    inp, out = _get_tokens(attrs)
    kind = attrs.get("openinference.span.kind", "")
    is_llm = _is_llm_span(attrs, name)

    img_prefix = f"span_{span_id}"

    lines: list[str] = []

    if is_llm:
        if model:
            cost = _cost(str(model), inp, out)
            lines.append(
                f"### {name} — {model} · "
                f"{inp:,} in / {out:,} out · {_fmt_cost(cost)} ({dur_str})"
            )
        else:
            lines.append(f"### {name} ({dur_str})")
        lines.append("")

        # Try OpenInference flat attributes first
        input_msgs = _reassemble_messages(attrs, "llm.input_messages")
        output_msgs = _reassemble_messages(attrs, "llm.output_messages")

        # Fall back to GenAI event-based messages
        if not input_msgs and not output_msgs:
            input_msgs, output_msgs = _get_genai_messages(span)

        # Fall back to smolagents input.value parsing
        if not input_msgs and not output_msgs:
            input_msgs = _parse_smolagents_input_value(attrs.get("input.value", ""))

        if input_msgs:
            lines.append(_render_messages(input_msgs, images_dir, img_prefix + "_in", img_counter))

        if output_msgs:
            lines.append(_render_messages(output_msgs, images_dir, img_prefix + "_out", img_counter))

        lines.append("---\n")

    elif kind == "AGENT":
        # CrewAI AGENT spans: try to extract messages from output.value JSON
        lines.append(f"### {name} ({dur_str})")
        if model:
            lines.append(f"*Model: {model}*\n")
        else:
            lines.append("")
        crewai_msgs = _parse_crewai_output_value(attrs)
        if crewai_msgs:
            lines.append(_render_messages(crewai_msgs, images_dir, img_prefix + "_agent", img_counter))
        lines.append("---\n")

    elif kind == "CHAIN":
        lines.append(f"## {name} ({dur_str})\n")

    else:
        lines.append(f"### {name} ({dur_str})\n")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Build the summary header from accumulated stats
# ---------------------------------------------------------------------------

def _build_header(
    trace_id: str,
    models: dict[str, dict],
    first_ts: float | None,
    span_count: int,
    total_dur_ns: int,
) -> str:
    lines: list[str] = []
    total_cost = sum(
        _cost(m, s["inp"], s["out"]) for m, s in models.items()
    )
    ts_str = (
        datetime.fromtimestamp(first_ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
        if first_ts
        else "—"
    )
    lines.append(f"# Trace {trace_id[:16]}… — {_fmt_cost(total_cost)}\n")
    lines.append(f"**{ts_str}** · {span_count} spans · {_fmt_dur(total_dur_ns)}\n")

    if models:
        lines.append("| Model | Calls | Input | Output | Cost |")
        lines.append("|---|---|---|---|---|")
        for m, s in models.items():
            c = _cost(m, s["inp"], s["out"])
            lines.append(
                f"| {m} | {s['calls']} | {s['inp']:,} | {s['out']:,} | {_fmt_cost(c)} |"
            )
        lines.append("")

    lines.append("---\n")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Exporter
# ---------------------------------------------------------------------------

class MarkdownTraceExporter(SpanExporter):
    """Incrementally writes human-readable Markdown, one file per trace."""

    def __init__(self, out_dir: str):
        self.out_dir = out_dir
        self._images_dir = os.path.join(out_dir, "images")
        os.makedirs(out_dir, exist_ok=True)
        # Per-trace bookkeeping
        self._files: dict[str, str] = {}                   # trace_id → filepath
        self._models: dict[str, dict[str, dict]] = {}      # trace_id → {model → stats}
        self._span_counts: dict[str, int] = defaultdict(int)
        self._first_ts: dict[str, float] = {}
        self._min_start: dict[str, int] = {}
        self._max_end: dict[str, int] = {}
        self._img_counter: list[int] = [0]                 # mutable counter shared across calls

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        for span in spans:
            trace_id = format(span.context.trace_id, "032x")

            # Ensure file exists
            if trace_id not in self._files:
                timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
                self._files[trace_id] = os.path.join(
                    self.out_dir, f"{timestamp}_{trace_id}.md"
                )
                self._models[trace_id] = {}
                self._first_ts[trace_id] = span.start_time / 1e9 if span.start_time else 0
                self._min_start[trace_id] = span.start_time or 0
                self._max_end[trace_id] = span.end_time or 0

            # Track timing bounds
            if span.start_time and span.start_time < self._min_start[trace_id]:
                self._min_start[trace_id] = span.start_time
            if span.end_time and span.end_time > self._max_end[trace_id]:
                self._max_end[trace_id] = span.end_time

            # Accumulate model stats
            attrs = dict(span.attributes or {})
            model = _get_model(attrs)
            if model:
                model = str(model)
                if model not in self._models[trace_id]:
                    self._models[trace_id][model] = {"inp": 0, "out": 0, "calls": 0}
                inp, out = _get_tokens(attrs)
                self._models[trace_id][model]["inp"] += inp
                self._models[trace_id][model]["out"] += out
                self._models[trace_id][model]["calls"] += 1

            self._span_counts[trace_id] += 1

            # Append body immediately (images saved to disk inline)
            body = _span_body(span, self._images_dir, self._img_counter)
            with open(self._files[trace_id], "a") as f:
                f.write(body)

        return SpanExportResult.SUCCESS

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Rewrite each file with the summary header prepended."""
        for trace_id, filepath in self._files.items():
            if not os.path.exists(filepath):
                continue
            with open(filepath) as f:
                body = f.read()

            dur_ns = self._max_end.get(trace_id, 0) - self._min_start.get(trace_id, 0)
            header = _build_header(
                trace_id,
                self._models.get(trace_id, {}),
                self._first_ts.get(trace_id),
                self._span_counts.get(trace_id, 0),
                dur_ns,
            )
            with open(filepath, "w") as f:
                f.write(header + body)
        return True

    def shutdown(self) -> None:
        self.force_flush()
