main.py

"""CLI entrypoint for Inventory Optimization AI Agent."""

from __future__ import annotations

import argparse
import json
import logging
import sys
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List

from agent.graph import build_graph
from agent.state import AgentState
from cli_helpers import apply_overrides, generate_report, parse_scenario_overrides, print_comparison, print_table
from tools.load_data import load_threshold_config

DISCLAIMER = (
    "⚠️ ADVISORY ONLY: This analysis is generated by an AI decision-support tool "
    "using synthetic mock data. All recommendations require human review and "
    "validation before any action is taken. This tool does not place orders, "
    "modify systems, or reflect real commercial inventory. Always consult your "
    "supply chain team before acting on these suggestions."
)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Inventory Optimization AI Agent")
    parser.add_argument("--data", default="data/inventory_mock.csv", help="Path to CSV/JSON inventory data")
    parser.add_argument("--config", default="config/thresholds.yaml", help="Path to YAML config")
    parser.add_argument("--output", default=None, help="Optional output JSON file path")
    parser.add_argument(
        "--scenario",
        action="append",
        default=[],
        help="Override values (e.g. --scenario lead_time=14 --scenario healthy_dos_min=12)",
    )
    parser.add_argument("--no-report", action="store_true", help="Disable markdown report generation")
    parser.add_argument(
        "--format",
        choices=["table", "json"],
        default="table",
        help="Console output format",
    )
    parser.add_argument(
        "--agent-mode",
        choices=["deterministic", "hybrid", "full"],
        default="deterministic",
        help="Select orchestration mode",
    )
    parser.add_argument(
        "--mode",
        choices=["fast", "thinking"],
        default="fast",
        help="Top-level mode: fast (no graph) or thinking (graph-enabled)",
    )
    parser.add_argument(
        "--fast-template-only",
        action="store_true",
        help="In fast mode, skip LLM and use template explanations only",
    )
    parser.add_argument("--model", default=None, help="Override Ollama model (e.g. gemma3:4b)")
    parser.add_argument("--sku", default=None, help="Analyze only one SKU id (e.g. SKU-001)")
    parser.add_argument("--skus", default=None, help="Analyze a comma-separated SKU list (e.g. SKU-001,SKU-003)")
    return parser.parse_args()


def run_analysis(config: Dict[str, Any]) -> Dict[str, Any]:
    """Execute full LangGraph analysis and return output payload."""
    state: AgentState = {
        "run_id": str(uuid.uuid4()),
        "started_at": datetime.now(timezone.utc).isoformat(),
        "config": config,
        "raw_records": [],
        "sku_records": [],
        "sku_metrics": [],
        "sku_contexts": [],
        "rule_results": {},
        "recommendations": [],
        "llm_prompts": {},
        "llm_responses": {},
        "llm_reasoning": {},
        "llm_reasoning_by_sku": {},
        "llm_retries": {},
        "flow_events": [],
        "tool_call_logs": [],
        "llm_batch_events": [],
        "agent_step_count": 0,
        "agent_max_steps": int(config.get("agent_max_steps", 3)),
        "agent_scratchpad": [],
        "agent_tool_history": [],
        "agent_seen_action_fingerprints": [],
        "agent_done": False,
        "agent_pending_action": None,
        "agent_fallback_reason": "",
        "current_node": "",
        "errors": [],
        "warnings": [],
        "partial_data": False,
        "graph_source": "default",
        "graph_runtime_stats": {},
        "output_valid": False,
        "final_output": None,
    }

    app = build_graph()
    try:
        final_state = app.invoke(state)
    except Exception as exc:
        state["errors"].append(
            {
                "node": state.get("current_node", "unknown"),
                "sku_id": "",
                "message": str(exc),
            }
        )
        state["warnings"].append("run_failed_before_output")
        return {
            "run_id": state["run_id"],
            "generated_at": datetime.now(timezone.utc).isoformat(),
            "summary": {
                "total_skus_analyzed": 0,
                "critical_count": 0,
                "watch_count": 0,
                "healthy_count": 0,
                "overstock_count": 0,
                "skus_skipped": 0,
                "overall_health": "poor",
                "top_priority_skus": [],
            },
            "recommendations": [],
            "metadata": {
                "llm_model": config.get("ollama", {}).get("model", "llama3.2:1b"),
                "graph_source": state.get("graph_source", "default"),
                "config_version": "1.0.0",
                "execution_time_ms": 0,
                "partial_data": True,
                "agent_mode": config.get("agent_mode", "deterministic"),
                "agent_steps_executed": state.get("agent_step_count", 0),
                "agent_tool_history": state.get("agent_tool_history", []),
                "flow_events": state.get("flow_events", []),
                "tool_call_logs": state.get("tool_call_logs", []),
                "llm_batch_events": state.get("llm_batch_events", []),
                "agent_fallback_reason": state.get("agent_fallback_reason", ""),
                "graph_runtime_stats": state.get("graph_runtime_stats", {}),
                "errors": state.get("errors", []),
                "warnings": state.get("warnings", []),
            },
            "disclaimer": DISCLAIMER,
        }
    return final_state.get("final_output") or {
        "run_id": state["run_id"],
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "summary": {
            "total_skus_analyzed": 0,
            "critical_count": 0,
            "watch_count": 0,
            "healthy_count": 0,
            "overstock_count": 0,
            "skus_skipped": 0,
            "overall_health": "poor",
            "top_priority_skus": [],
        },
        "recommendations": [],
        "metadata": {
            "llm_model": config.get("ollama", {}).get("model", "llama3.2:1b"),
            "graph_source": "default",
            "config_version": "1.0.0",
            "execution_time_ms": 0,
            "partial_data": True,
            "agent_mode": config.get("agent_mode", "deterministic"),
            "agent_steps_executed": final_state.get("agent_step_count", 0),
            "agent_tool_history": final_state.get("agent_tool_history", []),
            "flow_events": final_state.get("flow_events", []),
            "tool_call_logs": final_state.get("tool_call_logs", []),
            "llm_batch_events": final_state.get("llm_batch_events", []),
            "agent_fallback_reason": final_state.get("agent_fallback_reason", ""),
            "errors": final_state.get("errors", []),
            "warnings": final_state.get("warnings", []),
        },
        "disclaimer": DISCLAIMER,
    }


def main() -> None:
    """Execute the agent workflow and render output."""
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s - %(message)s")
    args = parse_args()

    base_config = load_threshold_config(Path(args.config))
    base_config["data_path"] = str(Path(args.data))
    base_config["config_path"] = str(Path(args.config))
    base_config["mode"] = args.mode
    base_config["fast_template_only"] = bool(args.fast_template_only)
    if args.model:
        base_config.setdefault("ollama", {})["model"] = str(args.model).strip()

    if isinstance(base_config.get("agent", {}), dict):
        agent_cfg = base_config.get("agent", {})  
    else:
        agent_cfg = {}

    if args.mode == "fast":
        base_config["agent_mode"] = "deterministic"
    else:
        base_config["agent_mode"] = args.agent_mode or str(agent_cfg.get("mode", "hybrid"))
    base_config["agent_max_steps"] = int(agent_cfg.get("max_steps", base_config.get("agent_max_steps", 3)))

    selected_skus: List[str] = []
    if args.sku:
        selected_skus.append(args.sku.strip())
    if args.skus:
        selected_skus.extend([item.strip() for item in args.skus.split(",") if item.strip()])
    if selected_skus:
        base_config["analysis_sku_ids"] = sorted(set(selected_skus))

    overrides = parse_scenario_overrides(args.scenario)
    scenario_config = apply_overrides(base_config, overrides) if overrides else base_config
    if overrides:
        scenario_config["scenario_overrides"] = overrides
    payload = run_analysis(scenario_config)

    if overrides:
        baseline_payload = run_analysis(base_config)
        print_comparison(baseline_payload, payload)

    output_path = Path(args.output) if args.output else Path("results") / f"{payload['run_id']}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False, default=str), encoding="utf-8")

    report_path = None if args.no_report else generate_report(payload, output_path, disclaimer=DISCLAIMER)

    if args.format == "json":
        print(json.dumps(payload, indent=2, ensure_ascii=False, default=str))
    else:
        print_table(payload.get("recommendations", []))
        print()
        print(json.dumps(payload.get("summary", {}), indent=2))
        print()
        print(f"Output written to: {output_path}")
        if report_path is not None:
            print(f"Report written to: {report_path}")
        print(DISCLAIMER)


if __name__ == "__main__":
    main()