"""
Memory Store Module

Session memory for tracking task progress:
- Task state management
- Action history
- LLM interaction logging
- Persistence support
"""

import json
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Optional


class TaskStatus(Enum):
    """Status of a task."""

    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class ActionRecord:
    """Record of a single agent action."""

    timestamp: str
    action_type: str  # "tool_call", "llm_response", "user_input", "error"
    tool_name: Optional[str] = None
    tool_args: Optional[dict] = None
    result: Optional[str] = None
    is_error: bool = False

    def to_dict(self) -> dict:
        return {
            "timestamp": self.timestamp,
            "action_type": self.action_type,
            "tool_name": self.tool_name,
            "tool_args": self.tool_args,
            "result": self.result[:500] if self.result else None,
            "is_error": self.is_error,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "ActionRecord":
        return cls(**data)


@dataclass
class TaskState:
    """State of the current task."""

    task_id: str
    description: str
    status: TaskStatus = TaskStatus.PENDING
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
    iterations: int = 0
    plan: list[str] = field(default_factory=list)
    current_step: int = 0
    summary: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "description": self.description,
            "status": self.status.value,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "iterations": self.iterations,
            "plan": self.plan,
            "current_step": self.current_step,
            "summary": self.summary,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "TaskState":
        data["status"] = TaskStatus(data["status"])
        return cls(**data)


class MemoryStore:
    """
    Session memory for the AI agent.
    
    Features:
    - Task state tracking
    - Action history recording
    - LLM conversation logging
    - Optional file persistence
    """

    def __init__(self, persist_path: Optional[Path] = None):
        self.persist_path = persist_path
        self.current_task: Optional[TaskState] = None
        self.action_history: list[ActionRecord] = []
        self.conversation: list[dict] = []
        self.metadata: dict[str, Any] = {}

        if persist_path and persist_path.exists():
            self._load()

    # Task Management

    def start_task(self, task_id: str, description: str) -> TaskState:
        """Start a new task."""
        self.current_task = TaskState(
            task_id=task_id,
            description=description,
            status=TaskStatus.IN_PROGRESS,
        )
        self.action_history = []
        self.conversation = []
        self._save()
        return self.current_task

    def update_task_status(self, status: TaskStatus, summary: Optional[str] = None) -> None:
        """Update the current task status."""
        if self.current_task:
            self.current_task.status = status
            self.current_task.updated_at = datetime.now().isoformat()
            if summary:
                self.current_task.summary = summary
            self._save()

    def set_task_plan(self, plan: list[str]) -> None:
        """Set the execution plan for the current task."""
        if self.current_task:
            self.current_task.plan = plan
            self._save()

    def advance_step(self) -> None:
        """Move to the next step in the plan."""
        if self.current_task:
            self.current_task.current_step += 1
            self.current_task.iterations += 1
            self.current_task.updated_at = datetime.now().isoformat()
            self._save()

    def increment_iteration(self) -> int:
        """Increment and return the iteration count."""
        if self.current_task:
            self.current_task.iterations += 1
            self._save()
            return self.current_task.iterations
        return 0

    def complete_task(self, summary: str) -> None:
        """Mark the current task as completed."""
        self.update_task_status(TaskStatus.COMPLETED, summary)

    def fail_task(self, error: str) -> None:
        """Mark the current task as failed."""
        self.update_task_status(TaskStatus.FAILED, error)

    # Action Recording

    def record_action(
        self,
        action_type: str,
        tool_name: Optional[str] = None,
        tool_args: Optional[dict] = None,
        result: Optional[str] = None,
        is_error: bool = False,
    ) -> ActionRecord:
        """Record an agent action."""
        record = ActionRecord(
            timestamp=datetime.now().isoformat(),
            action_type=action_type,
            tool_name=tool_name,
            tool_args=tool_args,
            result=result,
            is_error=is_error,
        )
        self.action_history.append(record)
        self._save()
        return record

    def record_tool_call(
        self,
        tool_name: str,
        tool_args: dict,
        result: str,
        is_error: bool = False,
    ) -> ActionRecord:
        """Record a tool call action."""
        return self.record_action(
            action_type="tool_call",
            tool_name=tool_name,
            tool_args=tool_args,
            result=result,
            is_error=is_error,
        )

    def record_llm_response(self, content: str) -> ActionRecord:
        """Record an LLM response."""
        return self.record_action(
            action_type="llm_response",
            result=content,
        )

    def record_error(self, error: str) -> ActionRecord:
        """Record an error."""
        return self.record_action(
            action_type="error",
            result=error,
            is_error=True,
        )

    # Conversation Management

    def add_message(self, role: str, content: str) -> None:
        """Add a message to the conversation history."""
        self.conversation.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat(),
        })
        self._save()

    def add_tool_result(
        self,
        tool_call_id: str,
        tool_name: str,
        result: str,
    ) -> None:
        """Add a tool result to the conversation."""
        self.conversation.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "name": tool_name,
            "content": result,
            "timestamp": datetime.now().isoformat(),
        })
        self._save()

    def get_conversation(self, max_messages: Optional[int] = None) -> list[dict]:
        """Get conversation history."""
        if max_messages:
            return self.conversation[-max_messages:]
        return self.conversation.copy()

    # Queries

    def get_recent_actions(self, count: int = 10) -> list[ActionRecord]:
        """Get the most recent actions."""
        return self.action_history[-count:]

    def get_error_count(self) -> int:
        """Get the count of errors in current session."""
        return sum(1 for action in self.action_history if action.is_error)

    def get_tool_usage(self) -> dict[str, int]:
        """Get tool usage statistics."""
        usage: dict[str, int] = {}
        for action in self.action_history:
            if action.tool_name:
                usage[action.tool_name] = usage.get(action.tool_name, 0) + 1
        return usage

    # Persistence

    def _save(self) -> None:
        """Save state to file if persist_path is set."""
        if not self.persist_path:
            return

        data = {
            "task": self.current_task.to_dict() if self.current_task else None,
            "actions": [a.to_dict() for a in self.action_history[-100:]],  # Keep last 100
            "conversation": self.conversation[-50:],  # Keep last 50 messages
            "metadata": self.metadata,
        }

        self.persist_path.parent.mkdir(parents=True, exist_ok=True)
        self.persist_path.write_text(json.dumps(data, indent=2))

    def _load(self) -> None:
        """Load state from file."""
        if not self.persist_path or not self.persist_path.exists():
            return

        try:
            data = json.loads(self.persist_path.read_text())

            if data.get("task"):
                self.current_task = TaskState.from_dict(data["task"])

            self.action_history = [
                ActionRecord.from_dict(a) for a in data.get("actions", [])
            ]
            self.conversation = data.get("conversation", [])
            self.metadata = data.get("metadata", {})

        except Exception:
            # Start fresh if load fails
            pass

    def clear(self) -> None:
        """Clear all stored data."""
        self.current_task = None
        self.action_history = []
        self.conversation = []
        self.metadata = {}

        if self.persist_path and self.persist_path.exists():
            self.persist_path.unlink()

    def get_session_summary(self) -> str:
        """Get a summary of the current session."""
        if not self.current_task:
            return "No active task"

        task = self.current_task
        tool_usage = self.get_tool_usage()
        tool_summary = ", ".join(f"{k}: {v}" for k, v in tool_usage.items())

        return f"""Task: {task.description}
Status: {task.status.value}
Iterations: {task.iterations}
Actions: {len(self.action_history)}
Errors: {self.get_error_count()}
Tool Usage: {tool_summary or 'none'}
Plan Progress: Step {task.current_step + 1}/{len(task.plan)} ({len(task.plan)} total steps)"""
