agent/pipeline.py

import os
from typing import Tuple, List, Dict
import json
import os
from datetime import datetime

from . import mcp_client, topology, parsing
from .oob_tools import (
    oob_get_last_backup,
    oob_perform_backup,
    oob_detect_drift,
    oob_get_alerts,
)
from .oob_tools.utils import load_json
from .oob_tools import root_cause
from .recovery import execute_recovery_plan
from .oob_tools.utils import save_json


def format_plan_markdown(steps):
    """
    steps: list of dicts with keys:
      - "description" (str)
      - "json" (dict)
    """
    lines = ["### Planned Steps\n"]
    for idx, step in enumerate(steps, start=1):
        lines.append(f"**Step {idx}:** {step['description']}")
        lines.append("```json")
        lines.append(parsing.to_pretty_json(step["json"]))
        lines.append("```")
        lines.append("")
    return "\n".join(lines)


def format_risk_markdown(overall_risk, step_risks, topo_summary):
    lines = ["### Overall Risk\n"]
    lines.append(f"- Level: **{overall_risk.get('level', 'unknown')}**")
    score = overall_risk.get("score")
    if score is not None:
        lines.append(f"- Score: `{score}`")

    if overall_risk.get("notes"):
        lines.append(f"- Notes: {overall_risk['notes']}")

    lines.append("\n### Per-Step Risk\n")
    for idx, r in enumerate(step_risks, start=1):
        lines.append(f"**Step {idx}:** {r.get('level', 'unknown')}")
        if "score" in r:
            lines.append(f"- Score: `{r['score']}`")
        if "summary" in r:
            lines.append(f"- Summary: {r['summary']}")
        lines.append("")

    lines.append("### Topology Impact\n")
    lines.append(topo_summary or "_No topology data yet._")

    return "\n".join(lines)


def format_diffs_markdown(diff_summaries):
    lines = ["### Config Diff Summary\n"]
    if not diff_summaries:
        lines.append("_No diff summaries available._")
        return "\n".join(lines)

    for idx, d in enumerate(diff_summaries, start=1):
        lines.append(f"**Step {idx}:** {d.get('title', 'Change')}")
        if "body" in d:
            lines.append("```")
            lines.append(d["body"])
            lines.append("```")
        lines.append("")
    return "\n".join(lines)


def format_rollback_markdown(rollback_steps):
    lines = ["### Rollback Plan\n"]
    if not rollback_steps:
        lines.append("_No rollback steps generated._")
        return "\n".join(lines)

    for idx, step in enumerate(rollback_steps, start=1):
        lines.append(f"**Step {idx}:** {step['description']}")
        if "commands" in step and step["commands"]:
            lines.append("```")
            for cmd in step["commands"]:
                lines.append(cmd)
            lines.append("```")
        lines.append("")
    return "\n".join(lines)


def format_tool_log_markdown(tool_calls):
    lines = ["### MCP Tool Call Log\n"]
    if not tool_calls:
        lines.append("_No tool calls executed._")
        return "\n".join(lines)

    for call in tool_calls:
        lines.append(f"- Tool: `{call.get('tool', 'unknown')}`")
        lines.append("  - Arguments:")
        lines.append("    ```json")
        lines.append(parsing.to_pretty_json(call.get("arguments", {})))
        lines.append("    ```")
        if "response" in call:
            lines.append("  - Response:")
            lines.append("    ```json")
            lines.append(parsing.to_pretty_json(call["response"]))
            lines.append("    ```")
        lines.append("")
    return "\n".join(lines)


def format_troubleshoot_log(tool_calls: List[Dict]) -> str:
    return format_tool_log_markdown(tool_calls)


def format_actions_markdown(actions: List[Dict]) -> str:
    lines = ["### Recommended Actions\n"]
    if not actions:
        lines.append("_No recommended actions available._")
        return "\n".join(lines)
    for act in actions:
        lines.append(f"**{act.get('title', 'Action')}**")
        if act.get("reason"):
            lines.append(f"- Reason: {act['reason']}")
        if act.get("actions"):
            lines.append("- Steps:")
            for step in act["actions"]:
                lines.append(f"  - `{json.dumps(step)}`")
        lines.append("")
    return "\n".join(lines)


def format_confidence_markdown(confidence: List[Dict]) -> str:
    lines = ["### Root Cause Confidence\n"]
    if not confidence:
        lines.append("_No confidence data available._")
        return "\n".join(lines)
    for item in confidence:
        lines.append(f"- {item.get('label')}: **{item.get('score_pct')}%**")
    return "\n".join(lines)


def format_device_state_markdown(device: str, device_state: Dict, bootstrap: Dict, backup: Dict, drift: Dict, alerts: List[Dict]) -> str:
    lines = ["### Device State Snapshot\n"]
    lines.append(f"- Device: **{device}**")
    lines.append(f"- Status: {device_state.get('status', 'unknown')}")
    lines.append(f"- Role/Pod: {device_state.get('role', 'n/a')} / {device_state.get('pod', 'n/a')}")
    lines.append(f"- Last config hash: {device_state.get('last_config_hash', 'n/a')}")
    lines.append(f"- Last update: {device_state.get('last_update', 'n/a')}")
    lines.append(f"- Bootstrapped: {bootstrap.get('bootstrapped', False)} at {bootstrap.get('bootstrap_time', 'n/a') if bootstrap else 'n/a'}")
    lines.append(f"- Backup: {backup.get('last_backup_at') or 'none'} (hash: {backup.get('config_hash')})")
    lines.append(f"- Drift: {drift.get('drift_detected')} (current: {drift.get('current_hash')}, last: {drift.get('last_known_hash')})")
    if alerts:
        lines.append(f"- Alerts (last {min(3, len(alerts))}):")
        for alert in alerts[-3:]:
            lines.append(f"  - [{alert.get('timestamp')}] {alert.get('message')}")
    else:
        lines.append("- Alerts: none")
    return "\n".join(lines)


def compute_health_score(backup: Dict, drift: Dict, alerts: List[Dict], last_mcp: Dict) -> Tuple[int, str]:
    drift_factor = 1.0 if drift.get("drift_detected") else 0.0
    alert_factor = min(1.0, len(alerts) / 3.0) if alerts else 0.0
    backup_age_factor = 1.0 if not backup.get("last_backup_at") else 0.2
    risk_level = (last_mcp.get("risk") or "").lower() if last_mcp else ""
    risk_map = {"critical": 1.0, "high": 0.7, "medium": 0.4, "low": 0.1}
    last_mcp_risk_factor = risk_map.get(risk_level, 0.0)

    health = 1 - 0.4 * drift_factor - 0.3 * alert_factor - 0.2 * backup_age_factor - 0.1 * last_mcp_risk_factor
    health_pct = max(0, min(int(round(health * 100)), 100))

    if health_pct >= 90:
        status = "Healthy"
    elif health_pct >= 70:
        status = "Unstable"
    elif health_pct >= 50:
        status = "Degraded"
    else:
        status = "Critical"

    return health_pct, status


def analyze_change(change_text: str) -> Tuple[str, str, str, str, str]:
    """
    Main entrypoint called by app.py.

    Returns:
      plan_markdown,
      risk_markdown,
      diffs_markdown,
      rollback_markdown,
      tool_log_markdown
    """
    if not change_text.strip():
        msg = "_Please describe your network change to begin analysis._"
        return msg, msg, msg, msg, msg

    # 1. Parse input into atomic steps (stubbed)
    steps = parsing.parse_change_request(change_text)

    # 2. Simulate via MCP (stubbed) and collect risk + tool logs
    step_risks, tool_calls = mcp_client.simulate_steps_with_mcp(steps)

    # 3. Compute topology summary (stubbed)
    topo_summary = topology.summarize_topology_impact(steps)

    # 4. Aggregate overall risk (very naive for now)
    overall_risk = {
        "level": "medium" if step_risks else "unknown",
        "score": sum(r.get("score", 0) for r in step_risks) / max(len(step_risks), 1)
        if step_risks
        else None,
        "notes": "Naive average of per-step scores (placeholder).",
    }

    # 5. Generate simple diff summaries (stubbed)
    diff_summaries = parsing.build_diff_summaries(steps)

    # 6. Generate rollback steps (stubbed)
    rollback_steps = parsing.build_rollback_plan(steps)

    # 7. Format for UI
    plan_md = format_plan_markdown(steps)
    risk_md = format_risk_markdown(overall_risk, step_risks, topo_summary)
    diffs_md = format_diffs_markdown(diff_summaries)
    rollback_md = format_rollback_markdown(rollback_steps)
    tool_log_md = format_tool_log_markdown(tool_calls)

    return plan_md, risk_md, diffs_md, rollback_md, tool_log_md


def oob_troubleshoot(device: str, issue_text: str) -> Tuple[str, str, str, str, str]:
    """
    Runs the Overgrowth OOB diagnosis pipeline.
    Returns markdown summary, tool log markdown, actions markdown, confidence markdown, device state markdown.
    """
    if not device:
        msg = "_Select a device to run troubleshooting._"
        return msg, msg, msg, msg, msg

    tool_calls: List[Dict] = []

    # Backup status; auto-backup if none exists to keep state fresh during this action.
    backup_info = oob_get_last_backup(device)
    tool_calls.append(
        {
            "tool": "oob_get_last_backup",
            "arguments": {"device": device},
            "response": backup_info,
        }
    )
    if backup_info.get("last_backup_at") is None:
        auto_backup = oob_perform_backup(device)
        tool_calls.append(
            {
                "tool": "oob_perform_backup",
                "arguments": {"device": device},
                "response": auto_backup,
            }
        )
        backup_info = auto_backup

    # Drift detection
    drift_info = oob_detect_drift(device)
    tool_calls.append(
        {
            "tool": "oob_detect_drift",
            "arguments": {"device": device},
            "response": drift_info,
        }
    )

    # Alerts
    alerts = oob_get_alerts(device)
    tool_calls.append(
        {
            "tool": "oob_get_alerts",
            "arguments": {"device": device},
            "response": {"alerts": alerts},
        }
    )

    # Bootstrap / device state introspection
    infra_root = os.path.join(os.path.dirname(__file__), "..", "infra")
    bootstrap_state = load_json(os.path.join(infra_root, "bootstrap.json"))
    device_state = load_json(os.path.join(infra_root, "device_state.json"))
    bootstrap_info = bootstrap_state.get(device)
    dev_state = device_state.get(device, {})

    # Topology impact for context
    topo_summary = topology.summarize_topology_impact([{"json": {"device": device}}])
    # richer topo info for root-cause
    topo_raw = load_json(os.path.join(infra_root, "topology.json"))
    topo_device_info = {}
    for d in topo_raw.get("devices", []):
        if d.get("id") == device:
            topo_device_info = d
            break
    topo_device_info.update(bootstrap_info or {})
    topo_device_info.update(dev_state or {})

    # Get last MCP history for this device
    mcp_history = load_json(os.path.join(infra_root, "mcp_history.json"))
    last_mcp = None
    for entry in reversed(mcp_history.get("recent", [])):
        if entry.get("device") == device:
            last_mcp = entry
            break

    # Root-cause inference
    ranked_causes, recommended_actions, narrative, confidence, stability_proxy = root_cause.infer_root_cause(
        device=device,
        issue_text=issue_text,
        backup_info=backup_info,
        drift_info=drift_info,
        topology_info=topo_device_info,
        alerts=alerts,
        last_mcp_result=last_mcp,
    )

    health_pct, health_status = compute_health_score(backup_info, drift_info, alerts, last_mcp)

    lines = ["### OOB Troubleshooting Summary\n", "---"]
    lines.append(f"**Device:** **{device}**")
    lines.append(f"**Issue:** {issue_text or '_No issue text provided_'}")
    lines.append(f"**Device Stability:** {health_pct} ({health_status})")
    lines.append("---")
    lines.append("**Backup status**")
    lines.append(
        f"- Last backup: {backup_info.get('last_backup_at') or 'none'} "
        f"(hash: {backup_info.get('config_hash') or 'n/a'})"
    )
    lines.append(f"- Snapshots: {len(backup_info.get('snapshots', []))}")
    lines.append("\n**Drift status**")
    lines.append(
        f"- Drift detected: {drift_info.get('drift_detected')} "
        f"(current: {drift_info.get('current_hash')}, "
        f"last known: {drift_info.get('last_known_hash')})"
    )
    lines.append(f"- Last checked: {drift_info.get('last_checked')}")
    lines.append("\n**Bootstrap / device state**")
    if bootstrap_info:
        lines.append(
            f"- Bootstrapped: {bootstrap_info.get('bootstrapped', False)} at {bootstrap_info.get('bootstrap_time')}"
        )
    else:
        lines.append("- Bootstrapped: unknown / not recorded.")
    lines.append(f"- Device status: {dev_state.get('status', 'unknown')}")
    lines.append(f"- Last config hash: {dev_state.get('last_config_hash', 'n/a')}")
    lines.append("\n**Alerts**")
    if alerts:
        for alert in alerts[-3:]:
            lines.append(f"- [{alert.get('timestamp')}] {alert.get('message')}")
    else:
        lines.append("- No alerts on record.")

    lines.append("\n**Topology impact**")
    lines.append(topo_summary)
    lines.append("---")

    lines.append("\n**Potential root causes**")
    for rc in ranked_causes:
        lines.append(f"- {rc.get('title')} (score: {rc.get('score')})")
        if rc.get("details"):
            lines.append(f"  - {rc['details']}")

    lines.append("\n**Expert narrative**")
    lines.append(narrative)

    summary_md = "\n".join(lines)
    tool_log_md = format_troubleshoot_log(tool_calls)
    actions_md = format_actions_markdown(recommended_actions)
    confidence_md = format_confidence_markdown(confidence)
    device_state_md = format_device_state_markdown(
        device, dev_state, bootstrap_info or {}, backup_info, drift_info, alerts
    )

    # Add root-cause tool call log for transparency
    tool_calls.append(
        {
            "tool": "oob_root_cause_infer",
            "arguments": {
                "device": device,
                "issue_text": issue_text,
            },
            "response": {
                "ranked_causes": ranked_causes,
                "recommended_actions": recommended_actions,
                "narrative": narrative,
                "confidence": confidence,
                "stability_proxy": stability_proxy,
            },
        }
    )
    tool_log_md = format_troubleshoot_log(tool_calls)

    return summary_md, tool_log_md, actions_md, confidence_md, device_state_md


def perform_recovery(device: str, issue_text: str) -> Tuple[str, str]:
    """
    Re-runs troubleshooting to derive actions, executes them, and returns execution summary and log.
    """
    if not device:
        msg = "_Select a device to execute recovery._"
        return msg, msg

    # Run the same data gathering as troubleshooting to derive actions.
    infra_root = os.path.join(os.path.dirname(__file__), "..", "infra")
    backup_info = oob_get_last_backup(device)
    drift_info = oob_detect_drift(device)
    alerts = oob_get_alerts(device)
    bootstrap_state = load_json(os.path.join(infra_root, "bootstrap.json"))
    device_state = load_json(os.path.join(infra_root, "device_state.json"))
    bootstrap_info = bootstrap_state.get(device)
    dev_state = device_state.get(device, {})
    topo_raw = load_json(os.path.join(infra_root, "topology.json"))
    topo_device_info = {}
    for d in topo_raw.get("devices", []):
        if d.get("id") == device:
            topo_device_info = d
            break
    topo_device_info.update(bootstrap_info or {})
    topo_device_info.update(dev_state or {})
    mcp_history = load_json(os.path.join(infra_root, "mcp_history.json"))
    last_mcp = None
    for entry in reversed(mcp_history.get("recent", [])):
        if entry.get("device") == device:
            last_mcp = entry
            break

    ranked_causes, recommended_actions, narrative, confidence, stability_proxy = root_cause.infer_root_cause(
        device=device,
        issue_text=issue_text,
        backup_info=backup_info,
        drift_info=drift_info,
        topology_info=topo_device_info,
        alerts=alerts,
        last_mcp_result=last_mcp,
    )

    health_before, health_status_before = compute_health_score(backup_info, drift_info, alerts, last_mcp)

    # Build flat action list from first recommendation, or empty.
    flat_actions = []
    if recommended_actions:
        # Prefer the first recommended action list; append clear_alerts to finish.
        flat_actions.extend(recommended_actions[0].get("actions", []))
    flat_actions.append({"type": "clear_alerts", "device": device})

    state_context = {
        "device": device,
        "root_cause": ranked_causes[0]["title"] if ranked_causes else "unknown",
        "outcome": "needs further work",
    }

    summary_md, exec_tool_calls, snapshot = execute_recovery_plan(flat_actions, state_context)

    # Compute health after using updated snapshot
    backup_after = snapshot.get("backups", {})
    drift_after = snapshot.get("drift", {})
    alerts_after = snapshot.get("alerts", [])
    # Last MCP remains the same unless actions included MCP call (history already updated in simulate call)
    mcp_history_after = load_json(os.path.join(infra_root, "mcp_history.json"))
    last_mcp_after = None
    for entry in reversed(mcp_history_after.get("recent", [])):
        if entry.get("device") == device:
            last_mcp_after = entry
            break

    health_after, health_status_after = compute_health_score(backup_after, drift_after, alerts_after, last_mcp_after)

    narrative_after = (
        f"Overgrowth executed recovery on {device}. Stability {health_before} ({health_status_before}) "
        f"-> {health_after} ({health_status_after}). "
        f"Actions run: {', '.join([a.get('type') for a in flat_actions])}. "
        "Reasoning: aligned device with stable snapshot, reconciled drift, cleared alerts, and refreshed backups."
    )

    # Add post-exec narrative to summary
    summary_md = summary_md + "\n\n" + narrative_after

    # Append synapse/audit trail with outcome
    outcome = "stabilized" if health_after >= 90 else "needs further work" if health_after >= 70 else "rollback suggested"
    synapse = load_json(os.path.join(infra_root, "synapse_log.json")) or {"events": []}
    synapse.setdefault("events", []).append(
        {
            "timestamp": datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
            "device": device,
            "root_cause_summary": ranked_causes[0]["title"] if ranked_causes else "unknown",
            "actions_executed": flat_actions,
            "outcome": outcome,
            "health_before": health_before,
            "health_after": health_after,
        }
    )
    synapse["events"] = synapse["events"][-200:]
    save_json(os.path.join(infra_root, "synapse_log.json"), synapse)

    # Build log markdown
    exec_log_md = format_troubleshoot_log(exec_tool_calls)

    return summary_md, exec_log_md


def load_synapse_log() -> str:
    """
    Returns a markdown-formatted view of the synapse log.
    """
    infra_root = os.path.join(os.path.dirname(__file__), "..", "infra")
    log = load_json(os.path.join(infra_root, "synapse_log.json")) or {"events": []}
    events = log.get("events", [])
    if not events:
        return "_No synapse events logged yet._"

    lines = ["### Overgrowth Synapse Log", "| Timestamp | Device | Root Cause | Outcome | Health Δ |", "| --- | --- | --- | --- | --- |"]
    for ev in reversed(events[-50:]):
        delta = f"{ev.get('health_before', '?')}→{ev.get('health_after', '?')}"
        lines.append(
            f"| {ev.get('timestamp', '?')} | {ev.get('device', '?')} | {ev.get('root_cause_summary', '?')} | "
            f"{ev.get('outcome', '?')} | {delta} |"
        )
    return "\n".join(lines)


def reset_state() -> str:
    """
    Resets Overgrowth persistent state to defaults. Dangerous; intended for demos.
    """
    infra_root = os.path.join(os.path.dirname(__file__), "..", "infra")
    targets = [
        ("device_state.json", {}),
        ("backups.json", {}),
        ("drift.json", {}),
        ("bootstrap.json", {}),
        ("alerts.json", {}),
        ("mcp_history.json", {"recent": []}),
        ("synapse_log.json", {"events": []}),
    ]
    for fname, data in targets:
        save_json(os.path.join(infra_root, fname), data)
    return "_Overgrowth state reset. All caches and histories cleared._"