from __future__ import annotations

from dataclasses import dataclass, field
import json
from typing import Any, Literal

from backend.task_graph import TaskGraph, TaskNode
from backend.tool_catalog import ToolCapability


ToolPlanFailureMode = Literal["abort", "continue", "ask_user"]


@dataclass(slots=True)
class ToolPlanStep:
    step_id: str
    tool: str
    purpose: str = ""
    args: dict[str, Any] = field(default_factory=dict)
    args_from: dict[str, Any] = field(default_factory=dict)
    expects: list[str] = field(default_factory=list)
    bind: dict[str, str] = field(default_factory=dict)
    on_failure: ToolPlanFailureMode = "abort"

    def to_dict(self) -> dict[str, Any]:
        return {
            "step_id": self.step_id,
            "tool": self.tool,
            "purpose": self.purpose,
            "args": dict(self.args),
            "args_from": dict(self.args_from),
            "expects": list(self.expects),
            "bind": dict(self.bind),
            "on_failure": self.on_failure,
        }


@dataclass(slots=True)
class ToolPlan:
    goal: str
    domain: str
    strategy: str = ""
    confidence: float = 0.0
    requires_clarification: bool = False
    clarification_question: str = ""
    steps: list[ToolPlanStep] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "goal": self.goal,
            "domain": self.domain,
            "strategy": self.strategy,
            "confidence": self.confidence,
            "requires_clarification": self.requires_clarification,
            "clarification_question": self.clarification_question,
            "steps": [step.to_dict() for step in self.steps],
            "metadata": dict(self.metadata),
        }


@dataclass(slots=True)
class ToolPlanValidationResult:
    valid: bool
    errors: list[str] = field(default_factory=list)
    blocked_tools: list[str] = field(default_factory=list)
    disabled_tools: list[str] = field(default_factory=list)
    approval_required_tools: list[str] = field(default_factory=list)
    unknown_tools: list[str] = field(default_factory=list)
    unavailable_tools: list[str] = field(default_factory=list)
    unresolved_bindings: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "valid": self.valid,
            "errors": list(self.errors),
            "blocked_tools": list(self.blocked_tools),
            "disabled_tools": list(self.disabled_tools),
            "approval_required_tools": list(self.approval_required_tools),
            "unknown_tools": list(self.unknown_tools),
            "unavailable_tools": list(self.unavailable_tools),
            "unresolved_bindings": list(self.unresolved_bindings),
        }


def tool_plan_step_from_payload(payload: dict | None) -> ToolPlanStep | None:
    raw = dict(payload or {})
    step_id = str(raw.get("step_id") or "").strip()
    tool = str(raw.get("tool") or "").strip()
    if not step_id or not tool:
        return None
    failure_mode = str(raw.get("on_failure") or "abort").strip().lower()
    if failure_mode not in {"abort", "continue", "ask_user"}:
        failure_mode = "abort"
    expects = [str(item).strip() for item in (raw.get("expects") or []) if str(item).strip()]
    bind = {
        str(key).strip(): str(value).strip()
        for key, value in dict(raw.get("bind") or {}).items()
        if str(key).strip() and str(value).strip()
    }
    return ToolPlanStep(
        step_id=step_id,
        tool=tool,
        purpose=str(raw.get("purpose") or "").strip(),
        args=dict(raw.get("args") or {}),
        args_from=dict(raw.get("args_from") or {}),
        expects=expects,
        bind=bind,
        on_failure=failure_mode,
    )


def tool_plan_from_payload(payload: dict | None) -> ToolPlan | None:
    raw = dict(payload or {})
    goal = str(raw.get("goal") or "").strip()
    domain = str(raw.get("domain") or "").strip().lower()
    if not goal or not domain:
        return None
    steps: list[ToolPlanStep] = []
    for item in (raw.get("steps") or []):
        step = tool_plan_step_from_payload(item if isinstance(item, dict) else None)
        if step is not None:
            steps.append(step)
    return ToolPlan(
        goal=goal,
        domain=domain,
        strategy=str(raw.get("strategy") or "").strip(),
        confidence=float(raw.get("confidence") or 0.0),
        requires_clarification=bool(raw.get("requires_clarification", False)),
        clarification_question=str(raw.get("clarification_question") or "").strip(),
        steps=steps,
        metadata=dict(raw.get("metadata") or {}),
    )


def build_tool_plan_prompt(base_prompt: str, capability_summary: list[dict[str, Any]]) -> str:
    prompt = str(base_prompt or "").strip()
    capabilities = list(capability_summary or [])
    return (
        f"{prompt}\n\n"
        "Before calling any tool, return ONLY a valid ToolPlan JSON object. Do not add prose.\n"
        "ToolPlan schema:\n"
        "{\n"
        '  "goal": "",\n'
        '  "domain": "email|calendar|search|general|code",\n'
        '  "strategy": "",\n'
        '  "confidence": 0.0,\n'
        '  "requires_clarification": false,\n'
        '  "clarification_question": "",\n'
        '  "steps": [\n'
        "    {\n"
        '      "step_id": "step_1",\n'
        '      "tool": "",\n'
        '      "purpose": "",\n'
        '      "args": {},\n'
        '      "args_from": {},\n'
        '      "expects": [],\n'
        '      "bind": {},\n'
        '      "on_failure": "abort|continue|ask_user"\n'
        "    }\n"
        "  ]\n"
        "}\n"
        "Rules:\n"
        "- Use only tools listed in Available tool capabilities.\n"
        "- If the task is ambiguous, set requires_clarification=true and leave steps empty.\n"
        "- Use args_from only for bindings from earlier plan steps, like step_1.primary_message_id.\n"
        "- For update/delete actions that need a message_id or event_id, plan the lookup step first.\n"
        "- Return JSON only.\n\n"
        f"Available tool capabilities:\n{json.dumps(capabilities, ensure_ascii=False)}"
    )


def tool_plan_to_task_graph(plan: ToolPlan) -> TaskGraph:
    tasks: list[TaskNode] = []
    prior_step_id = ""
    for step in plan.steps:
        args = dict(step.args or step.args_from)
        depends_on = [prior_step_id] if prior_step_id else []
        tasks.append(TaskNode(
            id=step.step_id,
            kind="tool_call",
            tool_name=step.tool,
            args=args,
            depends_on=depends_on,
            side_effect=False,
        ))
        prior_step_id = step.step_id
    return TaskGraph(
        goal=plan.goal,
        protocol="tool_plan",
        domain=plan.domain,
        display_goal=plan.goal,
        tasks=tasks,
        metadata={
            "strategy": plan.strategy,
            "confidence": plan.confidence,
            "requires_clarification": plan.requires_clarification,
            **dict(plan.metadata or {}),
        },
    )


def validate_tool_plan(
    plan: ToolPlan,
    *,
    catalog: dict[str, ToolCapability],
    blocked_tools: set[str] | None = None,
    disabled_tools: set[str] | None = None,
    approval_tools: set[str] | None = None,
    allow_side_effects: bool = True,
    max_steps: int = 8,
) -> ToolPlanValidationResult:
    blocked = blocked_tools or set()
    disabled = disabled_tools or set()
    approval = approval_tools or set()

    errors: list[str] = []
    blocked_hits: list[str] = []
    disabled_hits: list[str] = []
    approval_hits: list[str] = []
    unknown_tools: list[str] = []
    unavailable_tools: list[str] = []
    unresolved_bindings: list[str] = []

    if plan.requires_clarification:
        if plan.steps:
            errors.append("Tool plan cannot contain steps when clarification is required")
        if not plan.clarification_question:
            errors.append("Tool plan requires a clarification_question when clarification is required")

    if len(plan.steps) > max_steps:
        errors.append(f"Tool plan exceeds max_steps ({len(plan.steps)} > {max_steps})")

    step_ids: set[str] = set()
    available_bindings: dict[str, set[str]] = {}
    produced_resources: set[str] = set()

    for index, step in enumerate(plan.steps, start=1):
        if step.step_id in step_ids:
            errors.append(f"Duplicate tool plan step_id '{step.step_id}'")
        else:
            step_ids.add(step.step_id)

        if not step.args and not step.args_from:
            errors.append(f"Tool plan step '{step.step_id}' must define args or args_from")
        if step.args and step.args_from:
            errors.append(f"Tool plan step '{step.step_id}' cannot define both args and args_from")

        capability = catalog.get(step.tool)
        if capability is None:
            unknown_tools.append(step.tool)
            continue
        if not capability.available:
            unavailable_tools.append(step.tool)
        if step.tool in blocked:
            blocked_hits.append(step.tool)
        if step.tool in disabled or not capability.enabled:
            disabled_hits.append(step.tool)
        if step.tool in approval or capability.requires_approval:
            approval_hits.append(step.tool)
        if capability.side_effect and not allow_side_effects:
            errors.append(f"Tool plan step '{step.step_id}' uses side effects in a read-only validation mode")

        _validate_required_args(step, capability, errors)
        _validate_binding_references(step, available_bindings, unresolved_bindings, errors)
        _validate_prerequisites(step, capability, produced_resources, errors)

        if capability.planner_hints.get("produces"):
            produced_resources.update(str(item).strip() for item in capability.planner_hints.get("produces", []) if str(item).strip())
        if step.expects:
            produced_resources.update(step.expects)
        available_bindings[step.step_id] = set(step.bind.keys())

        if step.on_failure not in {"abort", "continue", "ask_user"}:
            errors.append(f"Tool plan step '{step.step_id}' has invalid on_failure '{step.on_failure}'")
        if not step.purpose:
            errors.append(f"Tool plan step '{step.step_id}' must include a purpose")

    valid = not any((errors, blocked_hits, disabled_hits, approval_hits, unknown_tools, unavailable_tools, unresolved_bindings))
    return ToolPlanValidationResult(
        valid=valid,
        errors=errors,
        blocked_tools=sorted(set(blocked_hits)),
        disabled_tools=sorted(set(disabled_hits)),
        approval_required_tools=sorted(set(approval_hits)),
        unknown_tools=sorted(set(unknown_tools)),
        unavailable_tools=sorted(set(unavailable_tools)),
        unresolved_bindings=sorted(set(unresolved_bindings)),
    )


def _validate_required_args(step: ToolPlanStep, capability: ToolCapability, errors: list[str]) -> None:
    schema = dict(capability.input_schema or {})
    required_fields = [str(item).strip() for item in (schema.get("required") or []) if str(item).strip()]
    provided_fields = set(str(key).strip() for key in {**step.args, **step.args_from}.keys() if str(key).strip())
    missing = [field for field in required_fields if field not in provided_fields]
    if missing:
        errors.append(
            f"Tool plan step '{step.step_id}' missing required args for '{step.tool}': {', '.join(missing)}"
        )


def _validate_binding_references(
    step: ToolPlanStep,
    available_bindings: dict[str, set[str]],
    unresolved_bindings: list[str],
    errors: list[str],
) -> None:
    for arg_name, arg_value in step.args_from.items():
        if not isinstance(arg_value, str):
            continue
        reference = str(arg_value).strip()
        if "." not in reference:
            continue
        source_step, binding_name = reference.split(".", 1)
        source_step = source_step.strip()
        binding_name = binding_name.strip()
        if not source_step or not binding_name:
            unresolved_bindings.append(f"{step.step_id}:{arg_name}")
            errors.append(f"Tool plan step '{step.step_id}' has invalid binding reference '{reference}'")
            continue
        known_bindings = available_bindings.get(source_step)
        if not known_bindings or binding_name not in known_bindings:
            unresolved_bindings.append(f"{step.step_id}:{arg_name}")
            errors.append(
                f"Tool plan step '{step.step_id}' references unknown binding '{reference}'"
            )


def _validate_prerequisites(
    step: ToolPlanStep,
    capability: ToolCapability,
    produced_resources: set[str],
    errors: list[str],
) -> None:
    required = [str(item).strip() for item in (capability.planner_hints.get("requires_previous_tool") or []) if str(item).strip()]
    for requirement in required:
        if requirement == "message_lookup" and "message_ids" not in produced_resources:
            errors.append(f"Tool plan step '{step.step_id}' requires a prior message lookup")
        elif requirement == "calendar_lookup" and "event_ids" not in produced_resources:
            errors.append(f"Tool plan step '{step.step_id}' requires a prior calendar lookup")
        elif requirement == "file_lookup" and "file_ids" not in produced_resources and "artifacts" not in produced_resources:
            errors.append(f"Tool plan step '{step.step_id}' requires a prior file lookup")