"""YAML config parsing and restricted expression parser for workflows."""
from __future__ import annotations

import ast
from typing import Any, Callable

import yaml

from logos_adk.errors import ConfigError
from logos_adk.state_store import build_state_store
from logos_adk.workflow.context import PipelineContext


# --- Restricted Expression Parser ---

_ALLOWED_OPS = {
    ast.Gt: lambda a, b: a > b,
    ast.Lt: lambda a, b: a < b,
    ast.GtE: lambda a, b: a >= b,
    ast.LtE: lambda a, b: a <= b,
    ast.Eq: lambda a, b: a == b,
    ast.NotEq: lambda a, b: a != b,
    ast.In: lambda a, b: a in b,
    ast.NotIn: lambda a, b: a not in b,
}


def parse_expression(expr: str) -> Callable[[PipelineContext], bool]:
    """Parse a restricted expression string into a callable.

    Only allows: Name, Attribute (dotted access), Constant, Compare.
    No function calls, imports, assignments, etc.
    """
    try:
        tree = ast.parse(expr, mode="eval")
    except SyntaxError as exc:
        raise ConfigError(f"Invalid expression: {expr!r}: {exc}") from exc

    _validate_ast(tree.body, expr)

    def evaluate(ctx: PipelineContext) -> bool:
        return _eval_node(tree.body, ctx)

    return evaluate


def _validate_ast(node: ast.AST, expr: str) -> None:
    """Validate that AST only contains allowed node types."""
    allowed = (ast.Compare, ast.Name, ast.Attribute, ast.Constant, ast.Load,
                ast.Gt, ast.Lt, ast.GtE, ast.LtE, ast.Eq, ast.NotEq, ast.In, ast.NotIn)
    if not isinstance(node, allowed):
        raise ConfigError(
            f"Disallowed expression construct in {expr!r}: {type(node).__name__}"
        )
    for child in ast.iter_child_nodes(node):
        _validate_ast(child, expr)


def _eval_node(node: ast.AST, ctx: PipelineContext) -> Any:
    """Evaluate an AST node against a PipelineContext."""
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.Name):
        return ctx[node.id]
    if isinstance(node, ast.Attribute):
        obj = _eval_node(node.value, ctx)
        if isinstance(obj, dict):
            return obj[node.attr]
        return getattr(obj, node.attr)
    if isinstance(node, ast.Compare):
        left = _eval_node(node.left, ctx)
        for op, comparator in zip(node.ops, node.comparators):
            right = _eval_node(comparator, ctx)
            op_fn = _ALLOWED_OPS.get(type(op))
            if op_fn is None:
                raise ConfigError(f"Unsupported operator: {type(op).__name__}")
            if not op_fn(left, right):
                return False
            left = right
        return True
    raise ConfigError(f"Cannot evaluate node: {type(node).__name__}")


# --- YAML Config Parser ---

def load_pipeline_from_dict(data: dict[str, Any]) -> Any:
    """Parse a pipeline config dict into a Pipeline instance."""
    from logos_adk.reliability import RetryPolicy
    from logos_adk.workflow.pipeline import Pipeline

    if not isinstance(data, dict):
        raise ConfigError("pipeline config must be a mapping")

    name = data.get("name")
    if not name:
        raise ConfigError("pipeline config requires 'name'")

    steps_data = data.get("steps")
    if steps_data is None:
        raise ConfigError("pipeline config requires 'steps'")

    retry = None
    retry_data = data.get("retry")
    if retry_data and isinstance(retry_data, dict):
        retry = RetryPolicy(
            max_attempts=retry_data.get("max_attempts", 3),
            base_delay=retry_data.get("base_delay", 0.25),
            backoff=retry_data.get("backoff", "exponential"),
        )

    timeout = data.get("timeout")

    state_store = _create_pipeline_state_store(data)

    pipeline = Pipeline(name, retry=retry, timeout=timeout, state_store=state_store)

    for step_data in steps_data:
        _parse_step(pipeline, step_data)

    return pipeline


def _create_pipeline_state_store(data: dict[str, Any]) -> Any:
    store_type = data.get("state_store")
    if store_type is None:
        return None

    from logos_adk.config import StateConfig

    return build_state_store(
        StateConfig(
            backend=store_type,
            dsn=data.get("state_store_dsn"),
            path=data.get("state_store_path", ".logos_adk/pipeline_state.db"),
            table_name=data.get("state_store_table_name", "agent_state_versions"),
        ),
        default_persistent=False,
    )


def load_graph_from_dict(
    data: dict[str, Any],
    *,
    agents: dict[str, Any] | None = None,
    functions: dict[str, Callable[[PipelineContext], Any]] | None = None,
) -> Any:
    """Parse a graph config dict into a Graph instance."""
    from logos_adk.workflow.graph import Graph

    if not isinstance(data, dict):
        raise ConfigError("graph config must be a mapping")

    name = data.get("name")
    if not isinstance(name, str) or not name:
        raise ConfigError("graph config requires 'name'")

    nodes_data = data.get("nodes")
    if not isinstance(nodes_data, list) or not nodes_data:
        raise ConfigError("graph config requires non-empty 'nodes'")

    graph = Graph(name, max_visits=data.get("max_visits", 64))

    for node_data in nodes_data:
        _parse_graph_node(graph, node_data, agents=agents, functions=functions)

    edges_data = data.get("edges", [])
    if not isinstance(edges_data, list):
        raise ConfigError("'edges' must be a list when provided")
    for edge_data in edges_data:
        _parse_graph_edge(graph, edge_data)

    entrypoint = data.get("entrypoint")
    if entrypoint is not None:
        if not isinstance(entrypoint, str) or not entrypoint:
            raise ConfigError("'entrypoint' must be a non-empty string")
        graph.entrypoint(entrypoint)

    terminal_nodes = data.get("terminal_nodes", [])
    if terminal_nodes:
        if not isinstance(terminal_nodes, list) or not all(isinstance(item, str) and item for item in terminal_nodes):
            raise ConfigError("'terminal_nodes' must be a list of non-empty strings")
        graph.terminal(*terminal_nodes)

    return graph


def load_graph_config(
    path: str,
    *,
    agents: dict[str, Any] | None = None,
    functions: dict[str, Callable[[PipelineContext], Any]] | None = None,
) -> Any:
    """Load a Graph from a YAML config file."""
    with open(path) as f:
        data = yaml.safe_load(f)
    graph = load_graph_from_dict(data, agents=agents, functions=functions)
    graph._source_path = path
    return graph


def save_graph_config(graph: Any, path: str | None = None) -> None:
    """Save a Graph instance back to a YAML config file."""
    save_path = path or getattr(graph, "_source_path", None)
    if not save_path:
        raise ConfigError(f"No source path found for graph '{graph.name}', cannot save.")

    if not hasattr(graph, "to_dict"):
        raise ConfigError(f"Graph '{graph.name}' does not support to_dict() serialization.")

    data = graph.to_dict()
    with open(save_path, "w") as f:
        yaml.safe_dump(data, f, sort_keys=False)


def _parse_step(pipeline: Any, data: dict[str, Any]) -> None:
    """Parse a single step config and add to pipeline."""
    from logos_adk.workflow.steps import ApprovalStep, ConditionalStep, LoopStep, ParallelStep
    from logos_adk.workflow.pipeline import _StepConfig

    if "approval" in data:
        info = data["approval"] if isinstance(data["approval"], dict) else {}
        s = ApprovalStep(info.get("name", "approval"), info.get("message"))
        pipeline._steps_config.append(_StepConfig(s))
        return

    if "conditional" in data:
        info = data["conditional"]
        condition = parse_expression(info["if"])
        then_step = _parse_single_step(info["then"])
        else_step = _parse_single_step(info["else"]) if "else" in info else None
        s = ConditionalStep(info.get("name", "conditional"), condition, then_step, else_step)
        pipeline._steps_config.append(_StepConfig(s))
        return

    if "loop" in data:
        info = data["loop"]
        body = _parse_single_step(info["step"])
        until = parse_expression(info["until"])
        s = LoopStep(
            info.get("name", "loop"),
            body,
            until,
            info.get("max_iterations", 10),
        )
        pipeline._steps_config.append(_StepConfig(s))
        return

    if "parallel" in data:
        info = data["parallel"]
        sub_steps = [_parse_single_step(sd) for sd in info["steps"]]
        s = ParallelStep(info.get("name", "parallel"), sub_steps)
        pipeline._steps_config.append(_StepConfig(s))
        return

    if "agent" in data:
        raise ConfigError(
            f"Agent step '{data.get('name', data['agent'])}' requires agents to be passed programmatically. "
            "Use load_pipeline_from_dict() with an agents dict."
        )

    raise ConfigError(f"Unknown step type in config: {list(data.keys())}")


def _parse_single_step(data: dict[str, Any]) -> Any:
    """Parse a single step dict into a Step instance."""
    from logos_adk.workflow.steps import ApprovalStep

    if "approval" in data:
        info = data["approval"] if isinstance(data["approval"], dict) else {}
        return ApprovalStep(info.get("name", "approval"), info.get("message"))

    raise ConfigError(f"Cannot parse inline step: {list(data.keys())}")


def _parse_graph_node(
    graph: Any,
    data: dict[str, Any],
    *,
    agents: dict[str, Any] | None,
    functions: dict[str, Callable[[PipelineContext], Any]] | None,
) -> None:
    from logos_adk.workflow.steps import ApprovalStep

    if not isinstance(data, dict):
        raise ConfigError("graph node config must be a mapping")
    name = data.get("name")
    if not isinstance(name, str) or not name:
        raise ConfigError("graph node requires non-empty string 'name'")

    metadata = data.get("metadata")
    if metadata is not None and not isinstance(metadata, dict):
        raise ConfigError(f"graph node '{name}' metadata must be a mapping")

    declared_types = [field for field in ("approval", "agent", "function", "graph", "parallel") if field in data]
    if len(declared_types) != 1:
        raise ConfigError(
            f"graph node '{name}' must declare exactly one of: approval, agent, function, graph, parallel"
        )

    if "approval" in data:
        info = data["approval"] if isinstance(data["approval"], dict) else {}
        target = ApprovalStep(name, info.get("message"), options=info.get("options"))
    elif "agent" in data:
        if agents is None:
            raise ConfigError(
                f"Graph agent node '{name}' requires agents to be passed programmatically"
            )
        agent_ref = data["agent"]
        if isinstance(agent_ref, dict):
            agent_ref = agent_ref.get("ref") or agent_ref.get("name")
        if not isinstance(agent_ref, str) or not agent_ref:
            raise ConfigError(f"graph node '{name}' agent ref must be a non-empty string")
        target = agents.get(agent_ref)
        if target is None:
            raise ConfigError(f"graph node '{name}' references unknown agent '{agent_ref}'")
    elif "function" in data:
        if functions is None:
            raise ConfigError(
                f"Graph function node '{name}' requires functions to be passed programmatically"
            )
        function_ref = data["function"]
        if isinstance(function_ref, dict):
            function_ref = function_ref.get("ref") or function_ref.get("name")
        if not isinstance(function_ref, str) or not function_ref:
            raise ConfigError(f"graph node '{name}' function ref must be a non-empty string")
        target = functions.get(function_ref)
        if target is None:
            raise ConfigError(f"graph node '{name}' references unknown function '{function_ref}'")
    elif "parallel" in data:
        from logos_adk.workflow.steps import ParallelStep, AgentStep, FunctionStep, ApprovalStep
        info = data["parallel"]
        sub_steps = []
        for step_data in info.get("steps", []):
            step_name = step_data.get("name", "sub_step")
            
            if "approval" in step_data:
                app_info = step_data["approval"] if isinstance(step_data["approval"], dict) else {}
                sub_steps.append(ApprovalStep(step_name, app_info.get("message"), options=app_info.get("options")))
            elif "agent" in step_data:
                agent_ref = step_data["agent"]
                if agents is None:
                    raise ConfigError(
                        f"graph node '{name}' parallel sub-step '{step_name}' requires agents to be passed programmatically"
                    )
                target_agent = agents.get(agent_ref)
                if target_agent is None:
                    raise ConfigError(
                        f"graph node '{name}' parallel sub-step '{step_name}' references unknown agent '{agent_ref}'"
                    )
                sub_steps.append(AgentStep(step_name, target_agent))
            elif "function" in step_data:
                func_ref = step_data["function"]
                if functions is None:
                    raise ConfigError(
                        f"graph node '{name}' parallel sub-step '{step_name}' requires functions to be passed programmatically"
                    )
                target_func = functions.get(func_ref)
                if target_func is None:
                    raise ConfigError(
                        f"graph node '{name}' parallel sub-step '{step_name}' references unknown function '{func_ref}'"
                    )
                sub_steps.append(FunctionStep(step_name, target_func))
            else:
                raise ConfigError(
                    f"graph node '{name}' parallel sub-step '{step_name}' has unsupported type"
                )
        target = ParallelStep(name, sub_steps)
    else:
        nested = data["graph"]
        if not isinstance(nested, dict):
            raise ConfigError(f"graph node '{name}' nested graph config must be a mapping")
        nested = dict(nested)
        nested.setdefault("name", name)
        target = load_graph_from_dict(nested, agents=agents, functions=functions)

    graph.node(target, name=name, metadata=metadata)


def _parse_graph_edge(graph: Any, data: dict[str, Any]) -> None:
    if not isinstance(data, dict):
        raise ConfigError("graph edge config must be a mapping")
    source = data.get("source", data.get("from"))
    target = data.get("target", data.get("to"))
    if not isinstance(source, str) or not source:
        raise ConfigError("graph edge requires non-empty 'source' or 'from'")
    if not isinstance(target, str) or not target:
        raise ConfigError("graph edge requires non-empty 'target' or 'to'")
    when = data.get("when")
    if when is not None and not isinstance(when, str):
        raise ConfigError(f"graph edge '{source}->{target}' 'when' must be a string")
    label = data.get("label")
    if label is not None and not isinstance(label, str):
        raise ConfigError(f"graph edge '{source}->{target}' 'label' must be a string")
    graph.edge(
        source,
        target,
        when=parse_expression(when) if when else None,
        label=label,
    )
