"""AgenLang Runtime - executes contracts with safety, audit, and settlement."""

from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Dict, Any, Tuple
import secrets
from cryptography.hazmat.primitives import hmac
from cryptography.hazmat.primitives.hashes import SHA256
from cryptography.hazmat.backends import default_backend
from .contract import Contract
from .memory import Memory
from .tools import TOOLS_REGISTRY
import asyncio

class Runtime:
    """Executes an AgenLang contract safely and produces the SER."""

    def __init__(self, contract: Contract, depth=0):
        if depth > 5:
            raise ValueError("Max subcontract depth exceeded")
        self.depth = depth
        self.contract = contract
        self.start_time = datetime.now(timezone.utc)
        self.execution_id = contract.contract_id
        self.steps_executed = 0
        self.replay_data = []
        self.last_output = None
        self.memory = Memory(self.execution_id, contract.memory_contract.get("data_subject", "anonymous"))
        self.replay_key = secrets.token_bytes(32)

    async def execute(self) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """Execute the contract and return (result, ser)."""
        if self.contract.constraints.get("joule_budget", 0) <= 0:
            raise ValueError("Joule budget exhausted or invalid")

        workflow_type = self.contract.workflow.get("type", "sequence")

        if workflow_type == "parallel":
            await self._execute_parallel()
        else:
            self._execute_sequence()

        result = {
            "status": "success",
            "output": f"Executed goal: {self.contract.goal}",
            "steps_completed": self.steps_executed
        }

        end_time = datetime.now(timezone.utc)
        joules_used = 42.0 * self.steps_executed
        ser = {
            "execution_id": self.execution_id,
            "timestamps": {
                "start": self.start_time.isoformat() + "Z",
                "end": end_time.isoformat() + "Z"
            },
            "decision_points": [],
            "resource_usage": {
                "joules_used": joules_used,
                "usd_cost": 0.01 * self.steps_executed,
                "efficiency_score": 0.92
            },
            "safety_checks": {
                "capability_violations": 0,
                "intent_anchor_verified": True
            },
            "replay_ref": f"{self.execution_id}.replay",
            "replay_key": self.replay_key.hex(),
            "settlement_receipt": {
                "joule_recipient": self.contract.settlement["joule_recipient"],
                "rate": self.contract.settlement["rate"],
                "total_joules_owed": joules_used * self.contract.settlement["rate"]
            }
        }

        self._save_replay()

        handoff_keys = self.contract.memory_contract.get("handoff_keys", [])
        if handoff_keys:
            step_data = {entry["step"]: entry["output"] for entry in self.replay_data}
            self.memory.handoff(handoff_keys, step_data)
            print(f"Memory handoff: saved keys {handoff_keys}")

        if self.contract.memory_contract.get("purge_on_complete", True):
            self.memory.purge()

        return result, ser

    def _execute_sequence(self):
        for step in self.contract.workflow.get("steps", []):
            self._execute_step_with_handler(step)

    async def _execute_parallel(self):
        steps = self.contract.workflow.get("steps", [])
        tasks = [asyncio.create_task(self._execute_step_with_handler_async(step)) for step in steps]
        await asyncio.gather(*tasks)

    def _execute_step_with_handler(self, step: Dict):
        max_retries = step.get("on_error", {}).get("retry", 0)
        fallback = step.get("on_error", {}).get("fallback")
        escalate_to = step.get("on_error", {}).get("escalate_to")
        notify_intent = step.get("on_error", {}).get("notify_intent_anchor", True)

        for attempt in range(max_retries + 1):
            try:
                self._execute_step(step)
                return
            except Exception as e:
                if attempt < max_retries:
                    print(f"Step failed (attempt {attempt+1}), retrying...")
                    continue
                else:
                    if fallback:
                        print(f"Falling back to: {fallback}")
                        return
                    if escalate_to and notify_intent:
                        print(f"Escalating to {escalate_to} and notifying intent_anchor")
                        raise ValueError(f"Escalated: {str(e)}")
                    raise

    async def _execute_step_with_handler_async(self, step: Dict):
        self._execute_step_with_handler(step)

    def _resolve_args(self, args: Dict) -> Dict:
        """Replace "$prev" values with the last step's output."""
        resolved = {}
        for k, v in args.items():
            if v == "$prev" and self.last_output is not None:
                resolved[k] = self.last_output
            else:
                resolved[k] = v
        return resolved

    def _execute_step(self, step: Dict):
        action = step.get("action")
        if action == "tool":
            tool_name = step.get("target")
            tool = TOOLS_REGISTRY.get(tool_name)
            if tool:
                required_caps = tool["capabilities"]
                if all(cap in [c["capability"] for c in self.contract.capability_attestations] for cap in required_caps):
                    module_name, func_name = tool["entry_point"].rsplit(".", 1)
                    module = __import__(module_name, fromlist=[func_name])
                    func = getattr(module, func_name)
                    args = self._resolve_args(step.get("args", {}))
                    output = func(args)
                    self.last_output = output
                    self.steps_executed += 1
                    self.replay_data.append({"step": tool_name, "output": output})
                    print(f"Executed tool: {tool_name} - Output: {output}")
                else:
                    raise ValueError(f"Capability violation for tool {tool_name}")
            else:
                raise ValueError(f"Unknown tool {tool_name}")
        elif action == "subcontract":
            sub_contract_data = step.get("args", {})
            sub_contract = Contract.from_dict(sub_contract_data)
            sub_runtime = Runtime(sub_contract, depth=self.depth + 1)
            sub_result, sub_ser = asyncio.run(sub_runtime.execute())
            self.replay_data.append({"step": "subcontract", "output": sub_result})
            print(f"Executed subcontract - Result: {sub_result['output']}")
        else:
            print(f"Executing step: {action} {step.get('target')}")

    def _save_replay(self):
        replay_content = json.dumps(self.replay_data).encode('utf-8')
        h = hmac.HMAC(self.replay_key, SHA256(), backend=default_backend())
        h.update(replay_content)
        hmac_value = h.finalize()
        replay_path = Path(f"{self.execution_id}.replay")
        replay_path.write_bytes(replay_content + hmac_value)
        print(f"Replay saved to {replay_path} with HMAC integrity")

    def to_ser_json(self, ser: Dict) -> str:
        return json.dumps(ser, indent=2)
