app/agent/tools.py

from typing import List, Dict, Any, Optional
import uuid
from datetime import datetime
from app.core.security import UserContext
from app.core.document_indexer import DocumentIndexer
from app.core.data_store import DataStore
from app.config import SNAPSHOT_DATETIME

def tool_document_search(
    query: str,
    user_context: UserContext,
    indexer: DocumentIndexer
) -> Dict[str, Any]:
    """
    Tool 1: Document Search & Retrieval.
    Searches policies, SOPs, agreements, and ops guides with access control & authority ranking.
    """
    documents = indexer.search_documents(query=query, user_context=user_context, top_k=5)
    return {
        "tool_name": "document_search",
        "query": query,
        "results_count": len(documents),
        "documents": documents
    }

def tool_structured_data_lookup(
    entity_type: str,  # "account", "order", "ticket", "summary"
    entity_id: Optional[str],
    user_context: UserContext,
    data_store: DataStore,
    filters: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Tool 2: Structured-data Lookup.
    Queries account, order, or ticket data enforcing data privacy scaping.
    """
    entity = entity_type.lower()
    if entity == "account":
        if entity_id:
            account = data_store.get_account(entity_id, user_context)
            return {
                "tool_name": "structured_data_lookup",
                "entity_type": "account",
                "data": account if account else {"error": f"Account {entity_id} not found or access denied."}
            }
        else:
            accounts = data_store.get_accounts(user_context)
            return {
                "tool_name": "structured_data_lookup",
                "entity_type": "account",
                "data": accounts
            }

    elif entity == "order":
        if entity_id:
            order = data_store.get_order(entity_id, user_context)
            if not order:
                return {"tool_name": "structured_data_lookup", "entity_type": "order", "error": f"Order {entity_id} not found or access denied."}
            
            # Enrich with delay calculation and account info
            delay = data_store.calculate_order_delay_hours(entity_id)
            elapsed_cancel = data_store.calculate_cancellation_elapsed_minutes(entity_id)
            account_info = data_store.get_account(order["account_id"], user_context)
            return {
                "tool_name": "structured_data_lookup",
                "entity_type": "order",
                "data": {
                    **order,
                    "pickup_delay_hours": delay,
                    "cancellation_elapsed_minutes": elapsed_cancel,
                    "account_name": account_info["account_name"] if account_info else "Unknown"
                }
            }
        else:
            target_acc = filters.get("account_id") if filters else None
            orders = data_store.get_orders(user_context, account_id=target_acc)
            return {
                "tool_name": "structured_data_lookup",
                "entity_type": "order",
                "data": orders
            }

    elif entity == "ticket":
        if entity_id:
            ticket = data_store.get_ticket(entity_id, user_context)
            return {
                "tool_name": "structured_data_lookup",
                "entity_type": "ticket",
                "data": ticket if ticket else {"error": f"Ticket {entity_id} not found or access denied."}
            }
        else:
            target_acc = filters.get("account_id") if filters else None
            status = filters.get("status") if filters else None
            tickets = data_store.get_tickets(user_context, account_id=target_acc, status=status)
            return {
                "tool_name": "structured_data_lookup",
                "entity_type": "ticket",
                "data": tickets
            }

    elif entity == "summary":
        return {
            "tool_name": "structured_data_lookup",
            "entity_type": "summary",
            "data": {
                "snapshot_time": data_store.readme_info.get("Dataset snapshot", str(data_store.snapshot_datetime)),
                "total_accounts": len(data_store.get_accounts(user_context)),
                "total_orders": len(data_store.get_orders(user_context)),
                "total_tickets": len(data_store.get_tickets(user_context))
            }
        }

    return {"error": f"Invalid entity_type {entity_type}"}

def tool_calculate_cancellation_fee(
    order_id: str,
    user_context: UserContext,
    data_store: DataStore,
    indexer: DocumentIndexer
) -> Dict[str, Any]:
    """
    Tool 2 (Calculator): Order Cancellation Fee & Eligibility Evaluator.
    Combines order status, timestamp elapsed calculation, SOP v4 rules, and signed Customer Agreement overrides.
    """
    order = data_store.get_order(order_id, user_context)
    if not order:
        return {"error": f"Order {order_id} not found or access denied."}

    account_id = order["account_id"]
    account = data_store.get_account(account_id, user_context)
    elapsed_minutes = data_store.calculate_cancellation_elapsed_minutes(order_id)
    order_status = order["status"].upper()

    # Default SOP v4 Rules
    # DRAFT: fee = 0
    # BOOKED, not picked up: <= 30 mins -> 0; > 30 mins -> INR 250 fee UNLESS contract waives it
    # PICKED_UP: Do not cancel, return-to-origin applies
    # DELIVERED: Cannot be cancelled

    fee_waived = False
    override_source = None
    contract_file = account.get("contract_file") if account else None

    # Check Customer Agreement Overrides (Precedence Level 4)
    if account_id == "ACCT-001":  # Northstar Logistics
        # Northstar Enterprise Agreement Section 2: "Northstar may cancel any BOOKED shipment before pickup with no cancellation fee, regardless of how long ago the shipment was booked."
        if order_status in ["BOOKED", "DRAFT"]:
            fee_waived = True
            override_source = "05_Northstar_Logistics_Enterprise_Agreement.pdf (Section 2)"

    standard_fee = 0
    if order_status == "BOOKED":
        if elapsed_minutes is not None and elapsed_minutes > 30:
            standard_fee = 250
        else:
            standard_fee = 0

    final_fee = 0 if fee_waived else standard_fee

    cancellation_allowed = order_status in ["DRAFT", "BOOKED"]

    return {
        "tool_name": "calculate_cancellation_fee",
        "order_id": order_id,
        "account_id": account_id,
        "account_name": account.get("account_name") if account else "Unknown",
        "order_status": order_status,
        "elapsed_minutes_since_booking": elapsed_minutes,
        "cancellation_allowed": cancellation_allowed,
        "standard_sop_fee_inr": standard_fee,
        "contract_fee_waived": fee_waived,
        "final_cancellation_fee_inr": final_fee,
        "governing_source": override_source if fee_waived else "03_Cancellation_and_Service_Credit_SOP_v4.pdf",
        "precedence_explanation": (
            f"Northstar's signed Enterprise Agreement (Level 4 Authority) waives cancellation fees for BOOKED shipments before pickup, "
            f"overriding standard SOP v4 (Level 2 Authority) which would charge INR 250 after 30 minutes."
            if fee_waived else
            f"Governed by SOP v4: {elapsed_minutes} minutes elapsed since booking. Fee is INR {final_fee}."
        )
    }

def tool_calculate_service_credit(
    order_id: str,
    user_context: UserContext,
    data_store: DataStore,
    indexer: DocumentIndexer
) -> Dict[str, Any]:
    """
    Tool 2 (Calculator): Failed-pickup Service Credit Evaluator.
    Calculates delay threshold, carrier fault check, SOP v4 default vs signed agreement rules.
    """
    order = data_store.get_order(order_id, user_context)
    if not order:
        return {"error": f"Order {order_id} not found or access denied."}

    account_id = order["account_id"]
    account = data_store.get_account(account_id, user_context)
    delay_hours = data_store.calculate_order_delay_hours(order_id)
    carrier_fault = order.get("carrier_fault", False)
    customer_fault = order.get("customer_fault", False)
    shipment_fee = float(order.get("shipment_fee_inr", 0))

    eligible = False
    credit_amount = 0.0
    governing_source = "03_Cancellation_and_Service_Credit_SOP_v4.pdf"
    explanation = ""

    # Check Customer Agreement Overrides
    if account_id == "ACCT-002":  # LumenWorks
        # LumenWorks Agreement Section 3: "If a pickup is more than 4 hours past the end of the scheduled pickup window, the carrier is at fault, and customer is not at fault, LumenWorks receives a fixed INR 300 service credit."
        governing_source = "06_LumenWorks_Service_Agreement.pdf (Section 3)"
        if delay_hours is not None and delay_hours > 4.0 and carrier_fault and not customer_fault:
            eligible = True
            credit_amount = 300.0
            explanation = f"LumenWorks Agreement requires pickup delay > 4 hours (actual delay: {delay_hours} hrs). Fixed credit of INR 300 applies."
        else:
            eligible = False
            credit_amount = 0.0
            explanation = f"Ineligible for credit under LumenWorks Agreement: Delay is {delay_hours} hours (must exceed 4.0 hours), carrier fault={carrier_fault}."

    else:
        # Standard SOP v4 Section 2: Delay > 2 hours, carrier fault, no customer fault. Credit = min(500, 10% of fee)
        if delay_hours is not None and delay_hours > 2.0 and carrier_fault and not customer_fault:
            eligible = True
            credit_amount = min(500.0, 0.10 * shipment_fee)
            explanation = f"Eligible under SOP v4: Delay is {delay_hours} hrs (>2.0 hrs threshold), carrier fault confirmed. Credit is min(500, 10% of INR {shipment_fee}) = INR {credit_amount}."
        else:
            eligible = False
            credit_amount = 0.0
            explanation = f"Ineligible under standard SOP v4: Delay is {delay_hours} hrs (threshold > 2.0 hrs), carrier fault={carrier_fault}."

    requires_manager_approval = credit_amount > 1000.0

    return {
        "tool_name": "calculate_service_credit",
        "order_id": order_id,
        "account_id": account_id,
        "account_name": account.get("account_name") if account else "Unknown",
        "delay_hours": delay_hours,
        "carrier_fault": carrier_fault,
        "customer_fault": customer_fault,
        "shipment_fee_inr": shipment_fee,
        "eligible": eligible,
        "calculated_credit_inr": credit_amount,
        "requires_manager_approval": requires_manager_approval,
        "governing_source": governing_source,
        "explanation": explanation
    }

def tool_prepare_state_action(
    action_name: str,  # "escalate_ticket", "update_ticket", "create_followup_task", "approve_service_credit"
    params: Dict[str, Any],
    user_context: UserContext
) -> Dict[str, Any]:
    """
    Tool 3: State-Changing Action Drafter.
    Generates a PENDING_CONFIRMATION payload requiring explicit user confirmation in the UI.
    """
    action_id = f"ACT-{uuid.uuid4().hex[:8].upper()}"

    descriptions = {
        "escalate_ticket": f"Escalate ticket {params.get('ticket_id')} to P1 / Tier-2 Operations",
        "update_ticket": f"Update ticket {params.get('ticket_id')} status to '{params.get('status')}' and assign to {params.get('assigned_to', 'Unassigned')}",
        "create_followup_task": f"Create follow-up engineering task: '{params.get('task_title')}'",
        "approve_service_credit": f"Apply INR {params.get('amount_inr')} service credit to Order {params.get('order_id')}"
    }

    return {
        "tool_name": "execute_action",
        "status": "PENDING_CONFIRMATION",
        "action_id": action_id,
        "action_name": action_name,
        "action_title": descriptions.get(action_name, f"Execute action {action_name}"),
        "parameters": params,
        "requested_by": user_context.user_id,
        "timestamp": datetime.now().isoformat(),
        "confirmation_required": True,
        "message": f"Action '{descriptions.get(action_name, action_name)}' prepared. Please confirm execution."
    }