"""Timeline Agent — correlates evidence across time."""

from __future__ import annotations

import json

from base_agent import BaseAgent
from evidence_graph import EvidenceGraph
from llm_client import LLMClient
from tool_registry import TOOL_CATALOG


class TimelineAgent(BaseAgent):
    name = "timeline"
    role = (
        "Timeline forensic analyst. You build chronological timelines from filesystem "
        "MAC timestamps and correlate events across all phenomena categories in the "
        "evidence graph to reconstruct the sequence of activities on the system."
    )

    def __init__(self, llm: LLMClient, graph: EvidenceGraph) -> None:
        super().__init__(llm, graph)
        self._register_tools()

    def _register_tools(self) -> None:
        # Filesystem timeline tool from catalog
        td = TOOL_CATALOG.get("build_filesystem_timeline")
        if td:
            self.register_tool(td.name, td.description, td.input_schema, td.executor)

        # Custom tool to get all phenomena with timestamps for correlation
        self.register_tool(
            name="get_timestamped_phenomena",
            description="Get all phenomena that have timestamps, sorted chronologically. Use for timeline correlation.",
            input_schema={"type": "object", "properties": {}},
            executor=self._get_timestamped_phenomena,
        )

        # Tool to add temporal edges between phenomena
        self.register_tool(
            name="add_temporal_edge",
            description="Add a temporal relationship between two phenomena (before, after, or concurrent).",
            input_schema={
                "type": "object",
                "properties": {
                    "source_id": {"type": "string", "description": "ID of the earlier/source phenomenon."},
                    "target_id": {"type": "string", "description": "ID of the later/target phenomenon."},
                    "relation": {
                        "type": "string",
                        "enum": ["before", "after", "concurrent"],
                        "description": "Temporal relationship.",
                    },
                },
                "required": ["source_id", "target_id", "relation"],
            },
            executor=self._add_temporal_edge,
        )

    async def _get_timestamped_phenomena(self) -> str:
        items = [
            ph for ph in self.graph.phenomena.values()
            if ph.timestamp
        ]
        items.sort(key=lambda ph: ph.timestamp or "")

        if not items:
            return "No phenomena with timestamps found."

        lines = []
        for ph in items:
            lines.append(f"{ph.timestamp} | [{ph.category}] {ph.title} ({ph.id})")
            lines.append(f"  {ph.description[:150]}")
        return "\n".join(lines)

    async def _add_temporal_edge(
        self, source_id: str, target_id: str, relation: str,
    ) -> str:
        try:
            await self.graph.add_edge(
                source_id=source_id,
                target_id=target_id,
                edge_type="temporal",
                metadata={"relation": relation},
                created_by=self.name,
            )
            return f"Temporal edge added: {source_id} —[{relation}]→ {target_id}"
        except ValueError as e:
            return f"Error: {e}"