"""
helix_cdc/agents/capsule_agent.py
=================================

Agent-in-Capsule: A seed running on semantic substrate with goals.

This is the "thin glue layer" ChatGPT mentioned - just binding what already exists.

Usage:
    from helix_cdc.agents import CapsuleAgent

    agent = CapsuleAgent(
        vault_path="vaults/mistral_language.vault",
        goals=["explain clearly", "be helpful"]
    )

    result = agent.think("What is the meaning of life?")
    print(result['output'])

    # Multi-step reasoning toward goal
    for step_result in agent.step_toward_goal("explain quantum physics", steps=5):
        print(step_result['word'])
"""

from __future__ import annotations

import json
import time
import hashlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Any, List, Optional, Iterator

from helix_cdc.substrate.core import (
    Substrate,
    Seed,
    GlyphState,
    load_vault,
    step_glyph,
    distance,
    gradient_toward,
    merge,
)


# State root for agent logs
AGENT_STATE_ROOT = Path("/tmp/bloomos_state/agents")
AGENT_STATE_ROOT.mkdir(parents=True, exist_ok=True)


@dataclass
class CapsuleAgent:
    """
    Agent-in-Capsule: A seed running on semantic substrate with goals.

    The agent:
    - Lives on a substrate (vault manifold)
    - Has active goals (encoded as seeds)
    - Maintains episodic log
    - Can think (generate) and step toward goals
    """
    vault_path: str
    goals: List[str] = field(default_factory=list)
    agent_id: Optional[str] = None
    temperature: float = 0.8
    goal_step_size: float = 0.05

    # Internal state
    _substrate: Optional[Substrate] = field(default=None, repr=False)
    _goal_seeds: List[Seed] = field(default_factory=list, repr=False)
    _current_seed: Optional[Seed] = field(default=None, repr=False)
    _log_path: Optional[Path] = field(default=None, repr=False)
    _step_count: int = field(default=0, repr=False)

    def __post_init__(self):
        # Generate agent ID if not provided
        if self.agent_id is None:
            self.agent_id = hashlib.sha256(f"{self.vault_path}:{time.time()}".encode()).hexdigest()[:12]

        # Load substrate
        self._substrate = load_vault(self.vault_path)

        # Encode goals as seeds
        self._goal_seeds = [self._substrate.seed_from_text(g) for g in self.goals]

        # Merge goals into single goal seed (if multiple)
        if len(self._goal_seeds) > 1:
            self._merged_goal = self._goal_seeds[0]
            for g in self._goal_seeds[1:]:
                self._merged_goal = merge(self._merged_goal, g)
        elif len(self._goal_seeds) == 1:
            self._merged_goal = self._goal_seeds[0]
        else:
            self._merged_goal = None

        # Setup log
        self._log_path = AGENT_STATE_ROOT / f"{self.agent_id}.jsonl"

    def _log(self, entry: Dict[str, Any]) -> None:
        """Append to agent log."""
        entry["ts"] = time.time()
        entry["agent_id"] = self.agent_id
        entry["step"] = self._step_count
        with self._log_path.open("a") as f:
            f.write(json.dumps(entry, default=str) + "\n")

    def think(self, prompt: str, max_tokens: int = 12) -> Dict[str, Any]:
        """
        Generate a response to prompt, guided by goals.

        Args:
            prompt: Input text
            max_tokens: Max tokens to generate

        Returns:
            Dict with output, tokens, seed trace
        """
        start = time.time()

        # Create seed from prompt
        seed = self._substrate.seed_from_text(prompt)
        self._current_seed = seed

        words = []
        seed_trace = []

        for pos in range(max_tokens):
            state = step_glyph(
                self._substrate,
                seed,
                position=pos,
                temperature=self.temperature,
                goal_seed=self._merged_goal,
                goal_step=self.goal_step_size
            )

            words.append(state.word)
            seed_trace.append({
                "pos": pos,
                "word": state.word,
                "basis_idx": state.basis_idx,
                "seed_hash": state.seed.hash
            })

            seed = state.seed
            self._step_count += 1

        self._current_seed = seed
        output = " ".join(words)
        duration = time.time() - start

        result = {
            "output": output,
            "tokens": words,
            "seed_trace": seed_trace,
            "duration_s": duration,
            "tensors_regenerated": 0,
            "agent_id": self.agent_id,
            "goals": self.goals,
        }

        self._log({"action": "think", "prompt": prompt, "output": output})

        return result

    def step_toward_goal(self, prompt: str, steps: int = 5) -> Iterator[Dict[str, Any]]:
        """
        Multi-step reasoning toward goal, yielding each step.

        Args:
            prompt: Starting prompt
            steps: Number of steps to take

        Yields:
            Dict with step info, word, distance to goal
        """
        seed = self._substrate.seed_from_text(prompt)

        for i in range(steps):
            state = step_glyph(
                self._substrate,
                seed,
                position=i,
                temperature=self.temperature,
                goal_seed=self._merged_goal,
                goal_step=self.goal_step_size * 2  # Larger steps for goal-directed
            )

            dist_to_goal = distance(state.seed, self._merged_goal) if self._merged_goal else None

            step_info = {
                "step": i,
                "word": state.word,
                "basis_idx": state.basis_idx,
                "seed_hash": state.seed.hash,
                "distance_to_goal": dist_to_goal,
            }

            self._log({"action": "step_toward_goal", **step_info})
            self._step_count += 1

            yield step_info

            seed = state.seed

        self._current_seed = seed

    def set_goals(self, goals: List[str]) -> None:
        """Update agent goals."""
        self.goals = goals
        self._goal_seeds = [self._substrate.seed_from_text(g) for g in goals]
        if len(self._goal_seeds) > 1:
            self._merged_goal = self._goal_seeds[0]
            for g in self._goal_seeds[1:]:
                self._merged_goal = merge(self._merged_goal, g)
        elif len(self._goal_seeds) == 1:
            self._merged_goal = self._goal_seeds[0]
        else:
            self._merged_goal = None

        self._log({"action": "set_goals", "goals": goals})

    def distance_to_goal(self) -> Optional[float]:
        """Current distance from agent's seed to goal."""
        if self._current_seed is None or self._merged_goal is None:
            return None
        return distance(self._current_seed, self._merged_goal)

    def get_log(self, limit: int = 20) -> List[Dict[str, Any]]:
        """Read agent log."""
        if not self._log_path.exists():
            return []
        entries = []
        with self._log_path.open() as f:
            for line in f:
                if line.strip():
                    entries.append(json.loads(line))
        return entries[-limit:]

    def info(self) -> Dict[str, Any]:
        """Agent info."""
        return {
            "agent_id": self.agent_id,
            "vault_path": self.vault_path,
            "vault_K": self._substrate.K,
            "vault_D": self._substrate.D,
            "goals": self.goals,
            "step_count": self._step_count,
            "log_path": str(self._log_path),
        }
