environments/trace_env/core/env.py

# """


from __future__ import annotations
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Optional

from .schemas import TraceAction, TraceObservation, EpisodeState
from .world_model import SemanticWorldModel
from ..agents.planner import PlannerAgent
from ..agents.retriever import RetrieverAgent
from ..agents.memory import MemoryAgent
from ..agents.verifier import VerifierAgent
from ..rewards.reward_fn import compute_reward
from ..rewards.anti_hack import AntiHackGuard


class TraceEnv:
    """
    OpenEnv-compatible environment for the Trace project.

    Episode lifecycle:
        reset(task) -> observation
        step(action) -> observation, reward, done, info
        state()      -> EpisodeState (for logging/debugging)

    The environment simulates a user with a federated digital footprint.
    The agent is given a long-horizon instruction (e.g. "Audit all receipts
    from 2022-2024") and must plan sub-tasks, retrieve data from virtual
    sources, and synthesize a verified result.
    """

    DEFAULT_MAX_STEPS = 20          # hard limit per episode
    DEFAULT_TIMEOUT_SECONDS = 300   # wall-clock timeout (per-step)

    def __init__(self, config: dict):
        self.config = config
        self.max_steps = config.get("max_steps", self.DEFAULT_MAX_STEPS)
        self.timeout_seconds = config.get("timeout_seconds", self.DEFAULT_TIMEOUT_SECONDS)

        self.world_model = SemanticWorldModel(config)
        self.planner = PlannerAgent(config)
        self.retriever = RetrieverAgent(config)
        self.memory = MemoryAgent(config)
        self.verifier = VerifierAgent(config)
        self.anti_hack = AntiHackGuard()

        self._episode_id: Optional[str] = None
        self._steps: int = 0
        self._step_start_time: float = 0.0
        self._state: Optional[EpisodeState] = None

    # ------------------------------------------------------------------
    # OpenEnv interface
    # ------------------------------------------------------------------

    def reset(self, task: dict) -> TraceObservation:
        """
        Start a fresh episode with a new task.

        Args:
            task: {
                "instruction": str,          # natural-language goal
                "difficulty": "easy"|"medium"|"hard",
                "available_sources": list[str],  # e.g. ["gmail", "drive"]
                "ground_truth": dict,        # for reward computation
            }

        Returns:
            TraceObservation: initial observation for the agent.
        """
        self._episode_id = str(uuid.uuid4())
        self._steps = 0
        self._step_start_time = time.time()

        self.world_model.initialize(task)
        self.memory.reset()
        self.anti_hack.reset()

        self._state = EpisodeState(
            episode_id=self._episode_id,
            task=task,
            plan=[],
            retrieved_data=[],
            verified=False,
            steps=0,
            done=False,
        )

        obs = TraceObservation(
            episode_id=self._episode_id,
            step=0,
            instruction=task["instruction"],
            available_sources=task["available_sources"],
            context="",
            memory_summary=self.memory.summarize(),
            world_state=self.world_model.snapshot(),
        )
        return obs

    def step(self, action: TraceAction) -> tuple[TraceObservation, float, bool, dict]:
        """
        Execute one agent action and return the next state.

        Action types:
          - PLAN:     decompose instruction into sub-tasks
          - RETRIEVE: fetch data from a virtual source
          - MEMORIZE: store a finding into episodic memory
          - VERIFY:   verify the current plan/answer against world model
          - ANSWER:   submit the final synthesized answer

        Returns:
            (observation, reward, done, info)
        """
        assert self._state is not None, "Call reset() before step()"

        self._steps += 1
        self._state.steps = self._steps

        # ── Timeout / step-limit guards ─────────────────────────────────
        # Per-step timeout: only the current step's processing time matters,
        # not idle time between API calls.
        self._step_start_time = time.time()
        if self._steps > self.max_steps:
            return self._terminate(reason="max_steps")

        # ── Anti-hack validation ─────────────────────────────────────────
        hack_flag = self.anti_hack.check(action)
        if hack_flag:
            reward = compute_reward(
                action, self._state, hack_penalty=True
            )
            info_dict = {"hack": hack_flag}
            obs = self._build_obs(f"[ANTI-HACK] {hack_flag}", metadata=info_dict)
            return obs, reward, False, info_dict

        # ── Dispatch action ──────────────────────────────────────────────
        result_context = ""

        action_type = action.action_type.strip().upper()

        if action_type == "PLAN":
            plan = self.planner.decompose(
                action.content, self._state.task
            )
            self._state.plan = plan
            result_context = f"Plan created: {plan}"

        elif action_type == "RETRIEVE":
            data = self.retriever.fetch(
                source=action.source,
                query=action.content,
                world_model=self.world_model,
                metadata=action.metadata,
            )
            if not isinstance(data, list):
                data = [data]

            # Inject real data into world model so visible_items updates
            self.world_model.inject_real_data(action.source, data)
            self._state.retrieved_data.extend(data)

            result_context = f"Retrieved {len(data)} items from {action.source}"
            info = {}

            # ── Gmail processing: merge and summarize all retrieved transactions ──
            if action.source == "gmail" and data:
                try:
                    from ..tools.transaction_parser import parse_transactions_bulk
                    from ..tools.dashboard_renderer import render_dashboard
                    
                    # We parse everything we've retrieved so far to ensure deduplication
                    # and a cumulative summary.
                    parsed_all = parse_transactions_bulk(self._state.retrieved_data)
                    summary = parsed_all.get("summary", {})
                    transactions = parsed_all.get("transactions", [])
                    
                    total_spend = summary.get("total_spend", 0.0)
                    tx_count = summary.get("count", 0)
                    by_category = summary.get("by_category", {})
                    top_category = next(iter(by_category.keys()), "unknown")
                    top_category_spend = by_category.get(top_category, 0.0)
                    
                    dashboard_html = render_dashboard(parsed_all)

                    result_context = (
                        f"Step summary: Retrieved {len(data)} new items. "
                        f"Cumulative Audit: {tx_count} total transactions | "
                        f"Total Spend: ₹{total_spend:,.2f} | "
                        f"Top Category: {top_category} (₹{top_category_spend:,.2f})"
                    )
                    info = {
                        "gmail_query": action.content,
                        "transactions_summary": summary,
                        "transactions": transactions,
                        "dashboard_html": dashboard_html,
                        "dashboard_generated": True,
                        "cumulative_count": tx_count,
                        "cumulative_spend": total_spend,
                    }
                except Exception as e:
                    info = {
                        "dashboard_generated": False,
                        "dashboard_error": str(e),
                    }
            
            elif action.source == "sheets":
                try:
                    from ..tools.sheets_tool import fetch_and_summarize
                    from ..tools.transaction_parser import parse_transactions_bulk
                    
                    summary = fetch_and_summarize()
                    sheet_txs = summary.get("transactions", [])
                    
                    # Deduplicate before extending: only add Sheets rows
                    # whose IDs are not already present from Gmail retrieval
                    existing_ids = {
                        item.get("id") for item in self._state.retrieved_data
                        if item.get("id")
                    }
                    # Keep all sheet transactions; parse_transactions_bulk will handle merging
                    self._state.retrieved_data.extend(sheet_txs)
                    
                    # Build cumulative summary from ALL retrieved data (Gmail + Sheets)
                    parsed_all = parse_transactions_bulk(self._state.retrieved_data)
                    summary = parsed_all.get("summary", {})
                    all_txs = parsed_all.get("transactions", [])
                    
                    # Calculate overlapping items for logging
                    gmail_ids = {
                        item.get("id") for item in self._state.retrieved_data
                        if item.get("_source") != "sheets" and item.get("id")
                    }
                    overlapping = sum(1 for tx in sheet_txs if tx.get("id") in gmail_ids)
                    
                    result_context = (
                        f"Retrieved {len(sheet_txs)} items from Google Sheets "
                        f"({len(sheet_txs) - overlapping} new, {overlapping} already in Gmail). "
                        f"Cumulative Audit: {summary.get('count', 0)} total transactions | "
                        f"Total Spend: ₹{summary.get('total_spend', 0.0):,.2f}"
                    )
                    info = {
                        "source": "sheets",
                        "sheets_count": len(sheet_txs),
                        "new_from_sheets": len(sheet_txs) - overlapping,
                        "transactions": all_txs,  # merged Gmail + Sheets
                        "transactions_summary": summary,
                    }
                except Exception as e:
                    result_context = f"Error retrieving from Sheets: {e}"
                    info = {"error": str(e)}
            else:
                info = {}

        elif action_type == "MEMORIZE":
            self.memory.store(action.content, action.metadata)
            result_context = "Stored to episodic memory."

        elif action_type == "VERIFY":
            verification = self.verifier.verify(
                claim=action.content,
                world_model=self.world_model,
                memory=self.memory,
            )
            self._state.verified = verification["passed"]
            result_context = f"Verification: {verification}"

        elif action_type == "SYNC":
            # Sync retrieved transactions to Google Sheets
            try:
                from ..tools.sheets_tool import append_transactions, fetch_and_summarize
                
                # Get transactions from the current state (parsed if available in info)
                # In a real scenario, we'd pull from world model or previous info.
                # For the demo, we'll use the last retrieved transactions if they exist.
                # However, retriever already puts data into world_model.
                
                # We need to get the parsed transactions.
                # Let's assume we want to sync whatever we found in the last Gmail pass.
                # But a cleaner way is to sync all transactions in the world model.
                
                all_tx = []
                # In a real implementation, we'd query the world model for all transactions.
                # For now, let's use the retrieved_data directly if it looks like Gmail data.
                # Or better, the env maintains a list of parsed transactions.
                
                from ..tools.transaction_parser import parse_transactions_bulk
                parsed = parse_transactions_bulk(self._state.retrieved_data)
                transactions = parsed.get("transactions", [])
                
                sheet_url = append_transactions(transactions)
                if sheet_url:
                    # After sync, fetch the full summary to verify
                    ledger_summary = fetch_and_summarize()
                    total_ledger = ledger_summary.get("total_spend", 0.0)
                    
                    result_context = (
                        f"Synced {len(transactions)} transactions to Google Sheets: {sheet_url}. "
                        f"Current Ledger Total: ₹{total_ledger:,.2f}"
                    )
                    info = {
                        "sheet_url": sheet_url,
                        "ledger_summary": ledger_summary,
                        "sync_count": len(transactions)
                    }
                else:
                    result_context = "Failed to sync to Google Sheets. Check credentials."
                    info = {"error": "Sync failed"}
            except Exception as e:
                result_context = f"Error during SYNC: {e}"
                info = {"error": str(e)}

        elif action_type == "EXPORT":
            # Export retrieved transactions to a DOCX report
            try:
                from ..tools.transaction_parser import parse_transactions_bulk
                from ..tools.report_tool import export_transactions_to_docx
                
                # Use all currently retrieved data (Gmail + Sheets)
                parsed = parse_transactions_bulk(self._state.retrieved_data)
                transactions = parsed.get("transactions", [])
                
                report_path = export_transactions_to_docx(transactions)
                
                result_context = (
                    f"Exported {len(transactions)} transactions to DOCX report at: {report_path}."
                )
                info = {
                    "report_path": report_path,
                    "export_count": len(transactions)
                }
            except Exception as e:
                result_context = f"Error during EXPORT: {e}"
                info = {"error": str(e)}

        elif action_type == "ANSWER":
            self._state.final_answer = action.content
            reward = compute_reward(action, self._state)
            info_dict = {"final_answer": action.content}
            obs = self._build_obs("Episode complete.", metadata=info_dict)
            self._state.done = True
            return obs, reward, True, info_dict

        else:
            result_context = f"Unknown action type: {action.action_type}"

        # ── Intermediate reward & next observation ───────────────────────
        info_dict = info if "info" in locals() else {}
        reward = compute_reward(action, self._state)
        obs = self._build_obs(result_context, metadata=info_dict)
        return obs, reward, False, info_dict

    def state(self) -> EpisodeState:
        """Return full episode state (for logging/debugging)."""
        return self._state

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

    def _build_obs(self, context: str, metadata: Optional[dict] = None) -> TraceObservation:
        return TraceObservation(
            episode_id=self._episode_id,
            step=self._steps,
            instruction=self._state.task["instruction"],
            available_sources=self._state.task["available_sources"],
            context=context,
            memory_summary=self.memory.summarize(),
            world_state=self.world_model.snapshot(),
            metadata=metadata or {},
        )

    def _terminate(self, reason: str) -> tuple[TraceObservation, float, bool, dict]:
        info_dict = {"termination_reason": reason}
        obs = self._build_obs(f"Episode terminated: {reason}", metadata=info_dict)
        reward = compute_reward(None, self._state, terminal_penalty=True)
        self._state.done = True
        return obs, reward, True, info_dict