"""plan_manager tool — create, track and query execution plans.

Plans are stored as JSON with two levels: phases contain steps.
All mutations propagate effective statuses bottom-up so the stored document
is always consistent.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import Any

from ims_mcp.constants import VALID_PLAN_STATUSES
from ims_mcp.services.plan_store import PlanStore

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

_SENTINEL = object()


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _err(code: str, message: str) -> str:
    return json.dumps({"error": code, "message": message})


def _ok(data: dict) -> str:
    return json.dumps(data)


# ---------------------------------------------------------------------------
# RFC 7396 JSON Merge Patch
# ---------------------------------------------------------------------------

def _merge_patch(target: Any, patch: Any) -> Any:
    """Apply RFC 7396 merge patch.  null values remove keys from dicts."""
    if not isinstance(patch, dict):
        return patch
    if not isinstance(target, dict):
        target = {}
    result = dict(target)
    for key, value in patch.items():
        if value is None:
            result.pop(key, None)
        else:
            result[key] = _merge_patch(result.get(key, _SENTINEL), value)
    return result


def _merge_by_id(existing: list[dict], incoming: list[dict]) -> list[dict] | str:
    """Merge two lists of dicts by their 'id' field (RFC 7396 per item).

    Returns the merged list or an error string if any item lacks an 'id'.
    Existing items absent from *incoming* are preserved.
    New items from *incoming* are appended.
    """
    for item in incoming:
        if "id" not in item or item["id"] is None:
            return "missing_id"
    existing_map: dict[str, dict] = {item["id"]: item for item in existing if "id" in item}
    result: list[dict] = list(existing)
    for patch_item in incoming:
        item_id = patch_item["id"]
        if item_id in existing_map:
            # Find and update in-place in result list
            for i, r in enumerate(result):
                if r.get("id") == item_id:
                    result[i] = _merge_patch(r, patch_item)
                    break
        else:
            result.append(dict(patch_item))
    return result


# ---------------------------------------------------------------------------
# Status helpers
# ---------------------------------------------------------------------------

def _compute_status_from_children(statuses: list[str]) -> str:
    # Priority: complete > failed > blocked > in_progress > open
    if not statuses:
        return "open"
    if all(s == "complete" for s in statuses):
        return "complete"
    if any(s == "failed" for s in statuses):
        return "failed"
    if any(s == "blocked" for s in statuses):
        return "blocked"
    if any(s in ("in_progress", "complete") for s in statuses):
        return "in_progress"
    return "open"


def _propagate_statuses(plan: dict) -> None:
    """Recompute and store effective statuses bottom-up across the plan."""
    phase_statuses = []
    for phase in plan.get("phases", []):
        step_statuses = [step.get("status", "open") for step in phase.get("steps", [])]
        if step_statuses:
            phase["status"] = _compute_status_from_children(step_statuses)
        phase_statuses.append(phase.get("status", "open"))
    if phase_statuses:
        plan["status"] = _compute_status_from_children(phase_statuses)


# ---------------------------------------------------------------------------
# Lookup helpers
# ---------------------------------------------------------------------------

def _find_phase(plan: dict, phase_id: str) -> dict | None:
    for phase in plan.get("phases", []):
        if phase.get("id") == phase_id:
            return phase
    return None


def _find_step(plan: dict, step_id: str) -> tuple[dict, dict] | None:
    """Return (phase, step) for the given step_id, or None."""
    for phase in plan.get("phases", []):
        for step in phase.get("steps", []):
            if step.get("id") == step_id:
                return phase, step
    return None


def _build_step_status_map(plan: dict) -> dict[str, str]:
    return {
        step["id"]: step.get("status", "open")
        for phase in plan.get("phases", [])
        for step in phase.get("steps", [])
        if "id" in step
    }


def _build_phase_status_map(plan: dict) -> dict[str, str]:
    return {
        phase["id"]: phase.get("status", "open")
        for phase in plan.get("phases", [])
        if "id" in phase
    }


def _deps_satisfied(item: dict, status_map: dict[str, str]) -> bool:
    return all(status_map.get(dep) == "complete" for dep in item.get("depends_on", []))


# ---------------------------------------------------------------------------
# Command: upsert
# ---------------------------------------------------------------------------

def cmd_upsert(
    plan: dict | None,
    target_id: str,
    data: dict,
    now: str,
) -> tuple[dict, str] | str:
    """Return (updated_plan, key_suffix) or error string."""

    if target_id == "entire_plan":
        if plan is None:
            # Create new plan
            base: dict = {
                "name": "Unnamed Plan",
                "description": "",
                "status": "open",
                "created_at": now,
                "updated_at": now,
                "phases": [],
            }
            merged = _merge_patch(base, data)
            # Merge phases by id if provided
            if "phases" in data and isinstance(data["phases"], list):
                result = _merge_by_id([], data["phases"])
                if result == "missing_id":
                    return "missing_id"
                # Also validate step ids within each phase
                for phase in result:
                    if isinstance(phase.get("steps"), list):
                        step_check = _merge_by_id([], phase["steps"])
                        if step_check == "missing_id":
                            return "missing_id"
                merged["phases"] = result
            # Ensure defaults on all phases/steps
            for phase in merged.get("phases", []):
                phase.setdefault("status", "open")
                phase.setdefault("steps", [])
                for step in phase.get("steps", []):
                    step.setdefault("status", "open")
            merged["created_at"] = now
            merged["updated_at"] = now
        else:
            # Patch existing plan
            merged = dict(plan)
            merged["updated_at"] = now
            if "phases" in data and isinstance(data["phases"], list):
                existing_phases = plan.get("phases", [])
                result = _merge_by_id(existing_phases, data["phases"])
                if result == "missing_id":
                    return "missing_id"
                # For each patched/new phase, also merge its steps
                merged_phases = []
                for phase in result:
                    orig_phase = _find_phase(plan, phase.get("id", "")) or {}
                    if "steps" in phase and isinstance(phase.get("steps"), list):
                        step_result = _merge_by_id(
                            orig_phase.get("steps", []), phase["steps"]
                        )
                        if step_result == "missing_id":
                            return "missing_id"
                        phase = dict(phase)
                        phase["steps"] = step_result
                    merged_phases.append(phase)
                # Patch remaining top-level fields
                patch_without_phases = {k: v for k, v in data.items() if k != "phases"}
                merged = _merge_patch(merged, patch_without_phases)
                merged["phases"] = merged_phases
                merged["updated_at"] = now
            else:
                merged = _merge_patch(merged, data)
                merged["updated_at"] = now
        _propagate_statuses(merged)
        return merged, "entire_plan"

    if plan is None:
        return "plan_not_found"

    # Phase-level upsert
    phase = _find_phase(plan, target_id)
    if phase is not None:
        updated_plan = dict(plan)
        updated_phases = []
        for p in updated_plan.get("phases", []):
            if p.get("id") == target_id:
                if "steps" in data and isinstance(data["steps"], list):
                    step_result = _merge_by_id(p.get("steps", []), data["steps"])
                    if step_result == "missing_id":
                        return "missing_id"
                    patch_without_steps = {k: v for k, v in data.items() if k != "steps"}
                    patched = _merge_patch(p, patch_without_steps)
                    patched["steps"] = step_result
                else:
                    patched = _merge_patch(p, data)
                updated_phases.append(patched)
            else:
                updated_phases.append(p)
        updated_plan["phases"] = updated_phases
        updated_plan["updated_at"] = now
        _propagate_statuses(updated_plan)
        return updated_plan, target_id

    # Step-level upsert
    found = _find_step(plan, target_id)
    if found is not None:
        parent_phase, _ = found
        updated_plan = dict(plan)
        updated_phases = []
        for p in updated_plan.get("phases", []):
            if p.get("id") == parent_phase.get("id"):
                new_steps = []
                for s in p.get("steps", []):
                    if s.get("id") == target_id:
                        new_steps.append(_merge_patch(s, data))
                    else:
                        new_steps.append(s)
                p = dict(p)
                p["steps"] = new_steps
            updated_phases.append(p)
        updated_plan["phases"] = updated_phases
        updated_plan["updated_at"] = now
        _propagate_statuses(updated_plan)
        return updated_plan, target_id

    # target_id not found as phase or step — treat as new phase if data has
    # no "phase_id" field, else treat as new step.
    if "phase_id" in data:
        # New step
        phase_id = data.get("phase_id")
        parent = _find_phase(plan, phase_id)
        if parent is None:
            return "target_not_found"
        new_step = {k: v for k, v in data.items() if k != "phase_id"}
        new_step["id"] = target_id
        new_step.setdefault("status", "open")
        new_step.setdefault("depends_on", [])
        updated_plan = dict(plan)
        updated_phases = []
        for p in updated_plan.get("phases", []):
            if p.get("id") == phase_id:
                p = dict(p)
                p["steps"] = list(p.get("steps", [])) + [new_step]
            updated_phases.append(p)
        updated_plan["phases"] = updated_phases
        updated_plan["updated_at"] = now
        _propagate_statuses(updated_plan)
        return updated_plan, target_id
    else:
        # New phase — require explicit phase_id via target_id
        return _upsert_new_phase(plan, target_id, data, now)



def _upsert_new_phase(plan: dict, target_id: str, data: dict, now: str) -> tuple[dict, str]:
    """Append a new phase to the plan."""
    new_phase = dict(data)
    new_phase["id"] = target_id
    new_phase.setdefault("name", target_id)
    new_phase.setdefault("description", "")
    new_phase.setdefault("status", "open")
    new_phase.setdefault("steps", [])
    new_phase.setdefault("depends_on", [])
    # Validate steps if provided
    if isinstance(new_phase.get("steps"), list):
        for step in new_phase["steps"]:
            step.setdefault("status", "open")
            step.setdefault("depends_on", [])
    updated_plan = dict(plan)
    updated_plan["phases"] = list(plan.get("phases", [])) + [new_phase]
    updated_plan["updated_at"] = now
    _propagate_statuses(updated_plan)
    return updated_plan, target_id


# ---------------------------------------------------------------------------
# Command: query
# ---------------------------------------------------------------------------

def cmd_query(plan: dict, target_id: str) -> str:
    if target_id == "entire_plan":
        return json.dumps(plan)
    phase = _find_phase(plan, target_id)
    if phase is not None:
        return json.dumps(phase)
    found = _find_step(plan, target_id)
    if found is not None:
        _, step = found
        return json.dumps(step)
    return _err("target_not_found", f"No phase or step with id '{target_id}'")


# ---------------------------------------------------------------------------
# Command: show_status
# ---------------------------------------------------------------------------

def _step_summary(step: dict) -> dict:
    return {"id": step.get("id"), "name": step.get("name", step.get("id")), "status": step.get("status", "open")}


def _phase_summary(phase: dict) -> dict:
    steps = phase.get("steps", [])
    return {
        "id": phase.get("id"),
        "name": phase.get("name", phase.get("id")),
        "status": phase.get("status", "open"),
        "steps": [_step_summary(s) for s in steps],
    }


def _totals(statuses: list[str]) -> dict:
    t: dict[str, int] = {s: 0 for s in VALID_PLAN_STATUSES}
    for s in statuses:
        t[s] = t.get(s, 0) + 1
    total = len(statuses)
    complete = t.get("complete", 0)
    t["total"] = total
    t["progress_pct"] = round(complete / total * 100, 1) if total else 0.0
    return t


def cmd_show_status(plan: dict, target_id: str) -> str:
    if target_id == "entire_plan":
        all_step_statuses = [
            step.get("status", "open")
            for phase in plan.get("phases", [])
            for step in phase.get("steps", [])
        ]
        totals = _totals(all_step_statuses)
        return json.dumps({
            "name": plan.get("name"),
            "status": plan.get("status", "open"),
            "progress_pct": totals["progress_pct"],
            "totals": totals,
            "phases": [_phase_summary(p) for p in plan.get("phases", [])],
        })

    phase = _find_phase(plan, target_id)
    if phase is not None:
        step_statuses = [s.get("status", "open") for s in phase.get("steps", [])]
        return json.dumps({
            "id": phase.get("id"),
            "name": phase.get("name", phase.get("id")),
            "status": phase.get("status", "open"),
            "progress_pct": _totals(step_statuses)["progress_pct"],
            "totals": _totals(step_statuses),
            "steps": [_step_summary(s) for s in phase.get("steps", [])],
        })

    found = _find_step(plan, target_id)
    if found is not None:
        _, step = found
        return json.dumps({
            "id": step.get("id"),
            "name": step.get("name", step.get("id")),
            "status": step.get("status", "open"),
            "prompt": step.get("prompt", ""),
            "depends_on": step.get("depends_on", []),
        })

    return _err("target_not_found", f"No phase or step with id '{target_id}'")


# ---------------------------------------------------------------------------
# Command: update_status
# ---------------------------------------------------------------------------

def cmd_update_status(plan: dict, target_id: str, new_status: str, now: str) -> tuple[dict, str] | str:
    """Update status and propagate UPWARD only.

    - step update  → recompute parent phase from its steps → recompute plan from phases
    - phase update → keep explicit phase status, recompute plan from phases only
    - entire_plan  → set plan status only
    """
    if new_status not in VALID_PLAN_STATUSES:
        return "invalid_status"

    if target_id == "entire_plan":
        updated = dict(plan)
        updated["status"] = new_status
        updated["updated_at"] = now
        return updated, new_status

    phase = _find_phase(plan, target_id)
    if phase is not None:
        # Set phase status explicitly; propagate plan status from phases (upward only).
        updated = dict(plan)
        updated["phases"] = [
            dict(p, status=new_status) if p.get("id") == target_id else p
            for p in plan.get("phases", [])
        ]
        updated["updated_at"] = now
        # Propagate upward: recompute plan status from (now-updated) phase statuses
        phase_statuses = [p.get("status", "open") for p in updated["phases"]]
        updated["status"] = _compute_status_from_children(phase_statuses)
        return updated, new_status

    found = _find_step(plan, target_id)
    if found is not None:
        parent_phase, _ = found
        updated = dict(plan)
        updated_phases = []
        for p in plan.get("phases", []):
            if p.get("id") == parent_phase.get("id"):
                p = dict(p)
                p["steps"] = [
                    dict(s, status=new_status) if s.get("id") == target_id else s
                    for s in p.get("steps", [])
                ]
                # Recompute phase status from its (now-updated) steps
                step_statuses = [s.get("status", "open") for s in p["steps"]]
                p["status"] = _compute_status_from_children(step_statuses)
            updated_phases.append(p)
        updated["phases"] = updated_phases
        updated["updated_at"] = now
        # Propagate upward: recompute plan status from phases
        phase_statuses = [p.get("status", "open") for p in updated["phases"]]
        updated["status"] = _compute_status_from_children(phase_statuses)
        return updated, new_status

    return "target_not_found"


# ---------------------------------------------------------------------------
# Command: next
# ---------------------------------------------------------------------------

def cmd_next(plan: dict, target_id: str, limit: int) -> str:
    step_status_map = _build_step_status_map(plan)
    phase_status_map = _build_phase_status_map(plan)
    ready: list[dict] = []

    phases = plan.get("phases", [])

    # When target_id is a specific step: check just that step
    if target_id not in ("entire_plan",) and _find_phase(plan, target_id) is None:
        found = _find_step(plan, target_id)
        if found is None:
            return _err("target_not_found", f"No phase or step with id '{target_id}'")
        parent_phase, step = found
        # Check phase deps
        if not _deps_satisfied(parent_phase, phase_status_map):
            return json.dumps([])
        if step.get("status", "open") == "open" and _deps_satisfied(step, step_status_map):
            ready.append(_step_entry(parent_phase, step))
        return json.dumps(ready)

    for phase in phases:
        # Filter to target phase if requested
        if target_id != "entire_plan" and phase.get("id") != target_id:
            continue
        # Skip if phase deps not satisfied
        if not _deps_satisfied(phase, phase_status_map):
            continue
        # Skip completed/failed phases entirely
        if phase.get("status") in ("complete", "failed"):
            continue
        for step in phase.get("steps", []):
            if step.get("status", "open") != "open":
                continue
            if _deps_satisfied(step, step_status_map):
                ready.append(_step_entry(phase, step))
                if limit and len(ready) >= limit:
                    return json.dumps(ready)

    return json.dumps(ready)


def _step_entry(phase: dict, step: dict) -> dict:
    entry: dict = {
        "phase_id": phase.get("id"),
        "phase_name": phase.get("name", phase.get("id")),
        "step_id": step.get("id"),
        "step_name": step.get("name", step.get("id")),
        "prompt": step.get("prompt", ""),
    }
    if step.get("model"):
        entry["model"] = step["model"]
    return entry


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

async def plan_manager_tool(
    command: str,
    plan_name: str,
    target_id: str,
    data: dict | None,
    new_status: str | None,
    limit: int,
    store: PlanStore,
) -> str:
    key = f"plan:{plan_name}"

    # ---- query / show_status / next: read-only ----
    if command == "query":
        plan = await store.get(key)
        if plan is None:
            return _err("plan_not_found", f"Plan '{plan_name}' not found")
        return cmd_query(plan, target_id)

    if command == "show_status":
        plan = await store.get(key)
        if plan is None:
            return _err("plan_not_found", f"Plan '{plan_name}' not found")
        return cmd_show_status(plan, target_id)

    if command == "next":
        plan = await store.get(key)
        if plan is None:
            return _err("plan_not_found", f"Plan '{plan_name}' not found")
        return cmd_next(plan, target_id, limit)

    # ---- upsert ----
    if command == "upsert":
        if data is None:
            return _err("missing_data", "upsert requires 'data'")
        plan = await store.get(key)
        now = _now_iso()
        result = cmd_upsert(plan, target_id, data, now)
        if isinstance(result, str):
            # error code
            messages = {
                "missing_id": "Array item is missing required 'id' field",
                "missing_phase_id": "New step upsert requires 'phase_id' in data",
                "plan_not_found": f"Plan '{plan_name}' not found; use target_id='entire_plan' to create it",
                "target_not_found": f"Phase or step '{target_id}' not found in plan",
            }
            return _err(result, messages.get(result, result))
        updated_plan, resolved_target = result
        await store.set(key, updated_plan)
        return _ok({"ok": True, "key": key, "target_id": resolved_target})

    # ---- update_status ----

    if command == "update_status":
        if not new_status:
            return _err("missing_new_status", "update_status requires 'new_status'")
        plan = await store.get(key)
        if plan is None:
            return _err("plan_not_found", f"Plan '{plan_name}' not found")
        now = _now_iso()
        result = cmd_update_status(plan, target_id, new_status, now)
        if isinstance(result, str):
            messages = {
                "invalid_status": f"'{new_status}' is not a valid status. Valid: {sorted(VALID_PLAN_STATUSES)}",
                "target_not_found": f"No phase or step with id '{target_id}'",
            }
            return _err(result, messages.get(result, result))
        updated_plan, resolved_status = result
        await store.set(key, updated_plan)
        return _ok({"ok": True, "target_id": target_id, "new_status": resolved_status})

    return _err("invalid_command", f"Unknown command '{command}'. Valid: upsert, query, show_status, update_status, next")
