"""Trace summarization pipeline using Predict.

Outputs markdown-frontmatter-ready metadata + summary.
When --memory-root is provided, writes the summary markdown file
directly to memory_root/summaries/YYYYMMDD/HHMMSS/{slug}.md using python-frontmatter.

Traces are compacted by adapters (tool outputs stripped), so most traces fit in a
single LLM call. For rare oversized traces, a sequential refine/fold pattern
processes chunks in order and accumulates into a single summary.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any

import dspy
import frontmatter
from pydantic import BaseModel, Field

from lerim.config.logging import logger
from lerim.memory.memory_record import slugify
from lerim.config.settings import get_config, reload_config
from lerim.memory.utils import (
    call_with_fallback,
    configure_dspy_lms,
    estimate_tokens,
    window_transcript,
    window_transcript_jsonl,
)
from lerim.runtime.cost_tracker import capture_dspy_cost
from lerim.sessions import catalog as session_db


class TraceSummaryCandidate(BaseModel):
    """Structured trace summary payload from LLM. Deterministic fields (date, time, coding_agent, etc.) are merged in Python."""

    title: str = Field(description="Short trace title for markdown frontmatter.")
    description: str = Field(
        description="One-line description of what the session achieved."
    )
    user_intent: str = Field(
        description="""\
The user's overall intention and request for this chat session. \
Not the literal query, but the broader goal the user was trying to \
accomplish. At most 150 words.""",
    )
    session_narrative: str = Field(
        description="""\
What actually happened over the course of the chat: actions taken, \
problems encountered, solutions applied, and final outcome. \
At most 200 words.""",
    )
    tags: list[str] = Field(
        default_factory=list,
        description="Group/cluster labels for this summary. No limit.",
    )


class TraceSummarySignature(dspy.Signature):
    """Summarize a raw coding-agent session trace into a structured summary.

    Produce a concise summary with:
    - user_intent: the user's overall goal (at most 150 words)
    - session_narrative: what happened chronologically (at most 200 words)
    - title and description: concise and specific
    - tags: descriptive group/cluster labels

    Ground all claims in transcript evidence. Avoid invented details.
    """

    transcript: str = dspy.InputField(
        desc="Raw session transcript text (JSONL or JSON, schema varies by agent)"
    )
    guidance: str = dspy.InputField(
        desc="Optional lead-agent natural language guidance about focus areas and trace context"
    )
    summary_payload: TraceSummaryCandidate = dspy.OutputField(
        desc="Structured summary payload with title/description/user_intent/session_narrative/tags fields"
    )


class RefineSummarySignature(dspy.Signature):
    """Refine a session summary with the next chunk of the trace.

    You have a running summary from earlier chunks and a new chunk of the trace.
    Update the summary to incorporate new information from this chunk while
    preserving important details from the running summary.
    """

    running_summary: str = dspy.InputField(desc="Summary from previous chunks (JSON)")
    trace_chunk: str = dspy.InputField(desc="Next chunk of the session trace")
    chunk_position: str = dspy.InputField(desc="e.g. 'chunk 2 of 4'")
    summary_payload: TraceSummaryCandidate = dspy.OutputField(
        desc="Updated summary incorporating the new chunk"
    )


def _summarize_trace(
    transcript: str,
    *,
    metadata: dict[str, Any] | None = None,
    guidance: str = "",
) -> dict[str, Any]:
    """Run Predict summarization and return validated summary with metadata merged.

    Single call when trace fits in context. Sequential refine/fold for oversized traces.
    Deterministic fields (date, time, coding_agent, etc.) are merged from metadata
    after the LLM call, not generated by the LLM.
    """
    if not transcript.strip():
        raise RuntimeError("session_trace_empty")
    config = get_config()
    max_window_tokens = config.summarize_role.max_window_tokens
    overlap_tokens = config.summarize_role.window_overlap_tokens
    lms = configure_dspy_lms("summarize")
    meta = metadata or {}
    guid = guidance.strip()

    history_start = len(lms[0].history)
    trace_tokens = estimate_tokens(transcript)

    if trace_tokens <= int(max_window_tokens * 0.85):
        # Fast path: single call — trace fits in context
        logger.info("Summarization: single call ({} est. tokens)", trace_tokens)
        w_start = time.time()
        summarizer = dspy.Predict(TraceSummarySignature)
        _, result = call_with_fallback(
            summarizer, lms, transcript=transcript, guidance=guid
        )
        logger.info("Summarization: done ({:.1f}s)", time.time() - w_start)
    else:
        # Slow path: refine/fold — process chunks sequentially
        if "\n{" in transcript:
            windows = window_transcript_jsonl(
                transcript, max_window_tokens, overlap_tokens
            )
        else:
            windows = window_transcript(transcript, max_window_tokens, overlap_tokens)
        logger.info(
            "Summarization: {} windows (refine/fold), {} est. tokens",
            len(windows),
            trace_tokens,
        )

        # First chunk: full summarization
        w_start = time.time()
        summarizer = dspy.Predict(TraceSummarySignature)
        _, result = call_with_fallback(
            summarizer, lms, transcript=windows[0], guidance=guid
        )
        logger.info("  Chunk 1/{}: done ({:.1f}s)", len(windows), time.time() - w_start)

        # Subsequent chunks: refine with running summary
        refiner = dspy.Predict(RefineSummarySignature)
        for i, window in enumerate(windows[1:], 2):
            running = getattr(result, "summary_payload", None)
            if isinstance(running, TraceSummaryCandidate):
                running_json = running.model_dump_json()
            elif isinstance(running, dict):
                running_json = json.dumps(running)
            else:
                running_json = str(running)
            w_start = time.time()
            _, result = call_with_fallback(
                refiner,
                lms,
                running_summary=running_json,
                trace_chunk=window,
                chunk_position=f"chunk {i} of {len(windows)}",
            )
            logger.info(
                "  Chunk {}/{}: done ({:.1f}s)",
                i,
                len(windows),
                time.time() - w_start,
            )

    capture_dspy_cost(lms[0], history_start)

    payload = getattr(result, "summary_payload", None)
    if isinstance(payload, TraceSummaryCandidate):
        candidate = payload
    elif isinstance(payload, dict):
        candidate = TraceSummaryCandidate.model_validate(payload)
    else:
        raise RuntimeError("dspy summary_payload must be TraceSummaryCandidate or dict")
    llm_payload = candidate.model_dump(mode="json", exclude_none=True)

    # Merge deterministic metadata fields (not generated by the LLM)
    now_iso = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    llm_payload["date"] = meta.get("date", now_iso[:10])
    llm_payload["time"] = meta.get("time", now_iso[11:19])
    llm_payload["coding_agent"] = meta.get("coding_agent") or meta.get(
        "agent_type", "unknown"
    )
    llm_payload["raw_trace_path"] = meta.get("raw_trace_path", "")
    llm_payload["run_id"] = meta.get("run_id", "")
    llm_payload["repo_name"] = meta.get("repo_name", "")
    return llm_payload


def write_summary_markdown(
    payload: dict[str, Any],
    memory_root: Path,
    *,
    run_id: str = "",
) -> Path:
    """Write summary markdown with frontmatter to memory_root/summaries/YYYYMMDD/HHMMSS/{slug}.md."""
    title = str(payload.get("title") or "untitled")
    user_intent = str(payload.get("user_intent") or "")
    session_narrative = str(payload.get("session_narrative") or "")
    summary_body = (
        f"## User Intent\n\n{user_intent}\n\n## What Happened\n\n{session_narrative}"
    )
    now_iso = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    fm_dict: dict[str, Any] = {
        "id": slugify(title),
        "title": title,
        "created": now_iso,
        "source": run_id,
        "description": payload.get("description", ""),
        "date": payload.get("date", now_iso[:10]),
        "time": payload.get("time", now_iso[11:19]),
        "coding_agent": payload.get("coding_agent", "unknown"),
        "raw_trace_path": payload.get("raw_trace_path", ""),
        "run_id": payload.get("run_id") or run_id,
        "repo_name": payload.get("repo_name", ""),
        "tags": payload.get("tags", []),
    }

    slug = slugify(title)
    date_compact = re.sub(r"[^0-9]", "", str(payload.get("date", now_iso[:10])))[:8]
    time_compact = re.sub(r"[^0-9]", "", str(payload.get("time", now_iso[11:19])))[:6]
    if len(date_compact) != 8 or len(time_compact) != 6:
        date_compact = datetime.now(timezone.utc).strftime("%Y%m%d")
        time_compact = datetime.now(timezone.utc).strftime("%H%M%S")
    summaries_dir = memory_root / "summaries" / date_compact / time_compact
    summaries_dir.mkdir(parents=True, exist_ok=True)
    summary_path = summaries_dir / f"{slug}.md"

    post = frontmatter.Post(summary_body, **fm_dict)
    summary_path.write_text(frontmatter.dumps(post) + "\n", encoding="utf-8")
    return summary_path


def summarize_trace_from_session_file(
    session_file_path: Path,
    *,
    metadata: dict[str, Any] | None = None,
    guidance: str = "",
) -> dict[str, Any]:
    """Summarize one session trace file into markdown-ready metadata + <=300 word summary."""
    if not session_file_path.exists() or not session_file_path.is_file():
        raise FileNotFoundError(f"session_file_missing:{session_file_path}")
    transcript = session_file_path.read_text(encoding="utf-8")
    meta = metadata or {}
    meta.setdefault("raw_trace_path", str(session_file_path))
    return _summarize_trace(
        transcript,
        metadata=meta,
        guidance=guidance,
    )


if __name__ == "__main__":
    """Run CLI summary mode by trace path or run a real-path self-test."""
    parser = argparse.ArgumentParser(
        prog="python -m lerim.memory.summarization_pipeline"
    )
    parser.add_argument("--trace-path")
    parser.add_argument("--output")
    parser.add_argument(
        "--memory-root",
        help="When provided, write summary markdown to memory_root/summaries/",
    )
    parser.add_argument("--metadata-json", default="{}")
    parser.add_argument("--guidance", default="")
    args = parser.parse_args()

    if args.trace_path:
        session_file = Path(args.trace_path).expanduser()
        metadata = json.loads(args.metadata_json)
        payload = summarize_trace_from_session_file(
            session_file,
            metadata=metadata if isinstance(metadata, dict) else {},
            guidance=args.guidance,
        )

        # Write summary markdown and output pointer
        if not args.memory_root:
            raise SystemExit("--memory-root is required")
        mr = Path(args.memory_root).expanduser().resolve()
        run_id = (metadata if isinstance(metadata, dict) else {}).get("run_id", "")
        summary_path = write_summary_markdown(payload, mr, run_id=run_id)
        output_data = {"summary_path": str(summary_path)}

        encoded = json.dumps(output_data, ensure_ascii=True, indent=2) + "\n"
        if args.output:
            output_path = Path(args.output).expanduser()
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_text(encoded, encoding="utf-8")
        else:
            sys.stdout.write(encoded)
    else:
        with TemporaryDirectory() as tmp_dir:
            tmp_path = Path(tmp_dir)
            cfg_path = tmp_path / "test_config.toml"
            cfg_path.write_text(
                f'[data]\ndir = "{tmp_path}"\n\n[memory]\nscope = "global_only"\n',
                encoding="utf-8",
            )
            prev_cfg = os.environ.get("LERIM_CONFIG")
            os.environ["LERIM_CONFIG"] = str(cfg_path)
            reload_config()
            session_db.init_sessions_db()

            run_id = "summary-self-test-1"
            session_path = tmp_path / "sessions" / f"{run_id}.jsonl"
            session_path.parent.mkdir(parents=True, exist_ok=True)
            session_path.write_text(
                '{"role":"user","content":"Fix queue heartbeat drift and duplicate claims."}\n'
                '{"role":"assistant","content":"Implemented heartbeat every 15 seconds and bounded retries."}\n',
                encoding="utf-8",
            )
            session_db.index_session_for_fts(
                run_id=run_id,
                agent_type="codex",
                repo_name="lerim",
                content="queue heartbeat fix",
                session_path=str(session_path),
            )
            payload = summarize_trace_from_session_file(
                session_path,
                metadata={
                    "run_id": run_id,
                    "coding_agent": "codex",
                    "repo_name": "lerim",
                },
            )

        if prev_cfg is None:
            os.environ.pop("LERIM_CONFIG", None)
        else:
            os.environ["LERIM_CONFIG"] = prev_cfg
        reload_config()
        assert isinstance(payload, dict)
        assert payload["title"]
        assert payload["description"]
        assert payload["coding_agent"] == "codex"
        assert payload["raw_trace_path"].endswith(f"{run_id}.jsonl")
        assert payload["user_intent"]
        assert len(payload["user_intent"].split()) <= 150
        assert payload["session_narrative"]
        assert len(payload["session_narrative"].split()) <= 200
