"""
Unified memory plugin with automatic environment detection.

Provides production-ready memory for DAITA agents.
Project-scoped by default, global as opt-in.
"""

import os
from typing import List, Optional
from ..base import BasePlugin
from ...core.tools import AgentTool, tool
from .metadata import MemoryMetadata


class MemoryPlugin(BasePlugin):
    """
    Production-ready memory plugin for DAITA agents.

    Features:
    - Automatic context injection before each run (on_before_run)
    - Hybrid semantic + keyword recall with importance weighting and temporal decay
    - Access tracking: recalled memories increment access_count for smarter pruning
    - Curator-maintained MEMORY.md: clean, categorized, never append-bloated
    - update_memory: agents can supersede stale or resolved facts

    Storage:
    - Project-scoped (default): .daita/memory/workspaces/{workspace}/
    - Global (opt-in): ~/.daita/memory/workspaces/{workspace}/

    Usage:
        # Isolated, project-scoped (default)
        agent.add_plugin(MemoryPlugin())

        # Shared workspace across agents
        memory = MemoryPlugin(workspace="research_team")
        agent1.add_plugin(memory)
        agent2.add_plugin(memory)

        # Global scope
        agent.add_plugin(MemoryPlugin(scope="global"))
    """

    def __init__(
        self,
        workspace: Optional[str] = None,
        scope: str = "project",
        auto_curate: str = "on_stop",
        curation_provider: Optional[str] = None,
        curation_model: Optional[str] = None,
        curation_api_key: Optional[str] = None,
        embedding_provider: str = "openai",
        embedding_model: str = "text-embedding-3-small"
    ):
        """
        Args:
            workspace: Workspace name (default: derived from agent_id for isolation)
            scope: "project" (default) or "global"
            auto_curate: When to curate: "on_stop" (default) or "manual"
            curation_provider: LLM for curation (default: "openai")
            curation_model: Model for curation (default: "gpt-4o-mini")
            curation_api_key: API key for curation (default: global settings)
            embedding_provider: Provider for embeddings (default: "openai")
            embedding_model: Embedding model (default: "text-embedding-3-small")
        """
        self.workspace = workspace
        self.scope = scope
        self.auto_curate = auto_curate
        self.curation_provider = curation_provider
        self.curation_model = curation_model
        self.curation_api_key = curation_api_key
        self.embedding_provider = embedding_provider
        self.embedding_model = embedding_model

        self._agent_id = None
        self.backend = None
        self.curator = None
        self.environment = None

    def initialize(self, agent_id: str):
        """Called by Agent.add_plugin() to inject agent context."""
        if self.backend is not None:
            self._agent_id = agent_id
            return

        self._agent_id = agent_id

        # Derive stable workspace name from agent_id (strip UUID suffix)
        if self.workspace:
            workspace = self.workspace
        elif '_' in agent_id and len(agent_id) > 9:
            workspace = '_'.join(agent_id.split('_')[:-1])
        else:
            workspace = agent_id

        runtime = os.getenv('DAITA_RUNTIME', 'local')

        if runtime == 'lambda':
            from .supabase_backend import SupabaseMemoryBackend

            org_id = os.getenv('DAITA_ORG_ID') or os.getenv('DAITA_ORGANIZATION_ID')
            project_name = os.getenv('DAITA_PROJECT_NAME')

            if not org_id or not project_name:
                raise RuntimeError(
                    "Missing DAITA_ORG_ID and DAITA_PROJECT_NAME for cloud memory."
                )

            self.backend = SupabaseMemoryBackend(
                org_id=org_id,
                project_name=project_name,
                workspace=workspace,
                agent_id=agent_id,
                scope=self.scope,
                embedding_provider=self.embedding_provider,
                embedding_model=self.embedding_model
            )
            self.environment = "cloud_supabase"
            print(f"Memory: cloud Supabase pgvector, {self.scope}-scoped, workspace='{workspace}'")
        else:
            from .local_backend import LocalMemoryBackend

            self.backend = LocalMemoryBackend(
                workspace=workspace,
                agent_id=agent_id,
                scope=self.scope,
                embedding_provider=self.embedding_provider,
                embedding_model=self.embedding_model
            )
            self.environment = "local"

        from .cloud_curator import CloudMemoryCurator

        self.curator = CloudMemoryCurator(
            backend=self.backend,
            agent_id=agent_id,
            llm_provider=self.curation_provider,
            llm_model=self.curation_model,
            api_key=self.curation_api_key
        )

        if self.environment == "local":
            scope_label = f"{self.scope}-scoped"
            workspace_label = "shared" if self.workspace else "isolated"
            print(f"Memory: {scope_label}, {workspace_label} workspace='{workspace}'")
            print(f"  Location: {self.backend.workspace_dir}")

        if self.auto_curate != "manual":
            print(f"  Auto-curation: {self.auto_curate} (provider: {self.curator.llm_provider}, model: {self.curator.llm_model})")

    def get_tools(self) -> List[AgentTool]:
        """Get memory tools for the LLM to use."""
        if self.backend is None:
            raise RuntimeError(
                "MemoryPlugin not initialized. Add via agent.add_plugin()."
            )

        backend = self.backend

        # Instantiate once, reused across all remember() calls in this session.
        from .contradiction import ContradictionChecker
        _checker = ContradictionChecker(
            llm=self.curator.llm,
            recall_fn=backend.recall,
            importance_threshold=0.7,
        )

        @tool
        async def remember(
            content: str,
            importance: float = 0.5,
            category: Optional[str] = None
        ):
            """
            Store information in memory for future recall.

            Use for facts, decisions, action items, and context worth preserving.

            Args:
                content: The information to store (be specific and self-contained)
                importance: How critical this memory is (0.0-1.0):
                    0.9+ = critical (security issues, hard deadlines, major decisions)
                    0.7-0.8 = important (key facts, events, action items)
                    0.5-0.6 = useful (general context, preferences)
                    0.3-0.4 = low priority (minor notes)
                category: Optional tag (e.g. "security", "project", "contact", "event")

            Returns:
                Confirmation with chunk_id
            """
            importance = max(0.0, min(1.0, importance))

            # Semantic dedup: skip if a near-identical memory already exists.
            # Compare against raw_semantic_score (not the hybrid score) to prevent
            # importance-boost inflation from falsely flagging distinct facts as
            # duplicates. E.g. "FK: api_keys.org_id → org" and
            # "FK: agents.org_id → org" share template tokens (raw ~0.80) but
            # score above 0.92 after a 0.9-importance boost — they are NOT duplicates.
            existing = await backend.recall(content, limit=1, score_threshold=0.6)
            if existing:
                raw_sim = (
                    existing[0].get("raw_semantic_score")
                    or existing[0].get("score_breakdown", {}).get("semantic")
                    or existing[0].get("score", 0)
                )
                if raw_sim >= 0.92:
                    return {
                        "status": "duplicate_skipped",
                        "message": "A near-identical memory already exists",
                        "chunk_id": existing[0].get("chunk_id"),
                        "existing_score": round(float(raw_sim), 3),
                    }

            # Contradiction/evolution check for important facts (importance >= 0.7).
            # - CONTRADICTION: new fact is a reasoning error conflicting with an
            #   established memory — block storage and explain what conflicts.
            # - EVOLUTION: a real-world change (e.g. FK dropped, constraint removed)
            #   that makes an existing memory outdated — auto-replace it so memory
            #   stays accurate without forcing the agent to call update_memory().
            if importance >= 0.7:
                conflict = await _checker.check(content, importance)
                if conflict.status == "contradiction":
                    return {
                        "status": "conflict_detected",
                        "message": (
                            "New fact contradicts an existing high-importance memory. "
                            "If this is a real change, use update_memory() to explicitly "
                            "replace the old fact."
                        ),
                        "conflicting_chunk_id": conflict.conflicting_chunk_id,
                        "conflicting_content": conflict.conflicting_content,
                        "conflict_reason": conflict.conflict_reason,
                    }
                elif conflict.status == "evolution":
                    # Real-world change: replace the stale fact automatically.
                    await backend.update_memory(
                        conflict.conflicting_content,
                        content,
                        importance,
                    )
                    return {
                        "status": "evolution_replaced",
                        "message": "Existing memory was outdated and has been replaced.",
                        "replaced_chunk_id": conflict.conflicting_chunk_id,
                        "conflict_reason": conflict.conflict_reason,
                    }

            metadata = MemoryMetadata(
                content=content,
                importance=importance,
                source='agent_inferred',
                category=category
            )
            return await backend.remember(content, category=category, metadata=metadata)

        @tool
        async def recall(
            query: str,
            limit: int = 5,
            score_threshold: float = 0.6,
            min_importance: Optional[float] = None,
            max_importance: Optional[float] = None,
            category: Optional[str] = None
        ):
            """
            Search memories by meaning using hybrid semantic + keyword search.

            Results are ranked by relevance, weighted by importance, and
            adjusted for age (recent memories rank slightly higher).

            Args:
                query: What you're looking for (natural language)
                limit: Maximum results (default: 5)
                score_threshold: Minimum relevance score 0-1 (default: 0.6)
                min_importance: Filter to memories with importance >= this value
                max_importance: Filter to memories with importance <= this value
                category: Restrict search to a specific category (e.g. "fk_constraint")

            Returns:
                Ranked list of relevant memories with scores and metadata
            """
            results = await backend.recall(
                query=query,
                limit=limit,
                score_threshold=score_threshold,
                min_importance=min_importance,
                max_importance=max_importance,
                category=category
            )

            # Track access for usage-based pruning signals
            chunk_ids = [r['chunk_id'] for r in results if r.get('chunk_id')]
            if chunk_ids:
                backend.search.track_access(chunk_ids)

            import json as _json
            return _json.loads(_json.dumps(results, default=str))

        @tool
        async def list_by_category(
            category: str,
            min_importance: float = 0.0,
            limit: int = 100
        ):
            """
            Enumerate ALL stored memories in a category without semantic ranking.

            PREFER this over recall() whenever the question asks to list or count
            everything of a known type:
              - "Which tables have jsonb columns?"  -> list_by_category("column_type")
              - "List all FK constraints"           -> list_by_category("fk_constraint")
              - "What are the row counts?"          -> list_by_category("table_stats")
              - "What schema observations exist?"   -> list_by_category("schema_observation")

            Use recall() only when you need semantic similarity ranking (e.g.
            "what do I know about the users table?"). For exhaustive enumeration,
            always use this tool — recall() truncates at its limit and may miss facts.

            Args:
                category: The category tag to enumerate (e.g. "fk_constraint", "column_type")
                min_importance: Only return memories with importance >= this value (default: 0.0)
                limit: Maximum results (default: 100)

            Returns:
                All matching memories ordered by importance descending
            """
            results = await backend.list_by_category(
                category=category,
                min_importance=min_importance,
                limit=limit
            )
            import json as _json
            return _json.loads(_json.dumps(results, default=str))

        @tool
        async def update_memory(
            query: str,
            new_content: str,
            importance: float = 0.5
        ):
            """
            Replace an existing memory with updated information.

            Use when a fact has changed, been resolved, or was incorrect.
            Finds the closest matching memory, removes it, and stores the update.

            Args:
                query: Description of the memory to find and replace
                new_content: The corrected or updated information
                importance: Importance score for the updated memory (0.0-1.0)

            Returns:
                Result with count of memories replaced
            """
            return await backend.update_memory(query, new_content, importance)

        @tool
        async def read_memory(file: str = "MEMORY.md"):
            """
            Read a memory file. Accepts shorthand values or absolute paths.

            Args:
                file: "MEMORY.md" (default) for long-term memory summary,
                      "today" for today's interaction log,
                      or an absolute file path

            Returns:
                Full file contents
            """
            if file == "MEMORY.md":
                return await backend.read_memory_md()
            elif file == "today":
                return await backend.read_today_log()
            else:
                return await backend.read_memory(file)

        @tool
        async def list_memories():
            """
            List available memory files (MEMORY.md and today's log).

            Returns:
                Summary of available memory files with paths and sizes
            """
            from datetime import date as _date
            today = _date.today().isoformat()
            files = []
            try:
                content = await backend.read_memory_md()
                if content and not content.startswith("# Long-Term Memory\n\n(No"):
                    files.append({"file": "MEMORY.md", "size_bytes": len(content.encode())})
            except Exception:
                pass
            try:
                log_content = await backend.read_today_log()
                if log_content and not log_content.startswith(f"# Daily Log - {today}\n\n(No"):
                    files.append({"file": f"logs/{today}.md", "size_bytes": len(log_content.encode())})
            except Exception:
                pass
            return files

        return [remember, recall, list_by_category, update_memory, read_memory, list_memories]

    async def on_before_run(self, prompt: str) -> Optional[str]:
        """
        Auto-inject relevant memories into agent context before each run.

        Called by the agent framework before the first LLM turn. Returns a
        formatted memory block that gets prepended to the system prompt,
        so the agent arrives with relevant context already loaded — no
        explicit recall call needed.
        """
        if self.backend is None:
            return None

        try:
            results = await self.backend.recall(prompt, limit=10, score_threshold=0.55)

            # Fallback: no semantic matches — inject top-3 by importance
            if not results:
                all_results = await self.backend.recall(prompt, limit=10, score_threshold=0.0)
                results = sorted(
                    all_results,
                    key=lambda r: r.get('metadata', {}).get('importance', 0),
                    reverse=True
                )[:3]
            else:
                # Always surface top-3 high-importance facts (importance >= 0.8),
                # deduplicated against the semantic results already in the list.
                seen = {r['chunk_id'] for r in results if r.get('chunk_id')}
                top_important = await self.backend.recall(
                    prompt, limit=10, score_threshold=0.0, min_importance=0.8
                )
                for r in sorted(
                    top_important,
                    key=lambda r: r.get('metadata', {}).get('importance', 0),
                    reverse=True
                )[:3]:
                    if r.get('chunk_id') not in seen:
                        results.append(r)
                        seen.add(r['chunk_id'])

            if not results:
                return None

            # Track access for these auto-injected memories
            chunk_ids = [r['chunk_id'] for r in results if r.get('chunk_id')]
            if chunk_ids:
                self.backend.search.track_access(chunk_ids)

            lines = ["## Relevant Memory"]
            for r in results:
                importance = r.get('metadata', {}).get('importance', 0.5)
                category = r.get('metadata', {}).get('category', '')
                tag = f"[{category}] " if category else ""
                lines.append(f"- {tag}{r['content'].strip()} (importance: {importance:.1f})")

            return '\n'.join(lines)
        except Exception as e:
            import logging as _logging
            _logging.getLogger(__name__).warning(
                "Memory recall failed in on_before_run: %s. "
                "Agent will proceed without memory context.", e
            )
            return None

    async def curate(self):
        """Manually trigger curation. Returns CurationResult."""
        if not self.curator:
            raise RuntimeError("Curator not initialized")
        return await self.curator.curate()

    async def on_agent_stop(self):
        """Flush pending changes and auto-curate if enabled."""
        if hasattr(self.backend, 'flush'):
            await self.backend.flush()

        if self.auto_curate == "on_stop":
            result = await self.curate()
            if result.success:
                msg = f"Memory curated: +{result.facts_added} new"
                if result.memories_updated > 0:
                    msg += f", ~{result.memories_updated} updated"
                if result.memories_pruned > 0:
                    msg += f", -{result.memories_pruned} pruned"
                if getattr(result, 'existing_memories', 0) > 0:
                    msg += f" ({result.existing_memories} stored by agent)"
                msg += f", ${result.cost_usd:.4f}"
                print(msg)

    def get_pending_metrics(self) -> dict:
        if self.backend is None:
            return {"memory_count_delta": 0, "memory_retrieval_count": 0}
        if hasattr(self.backend, 'get_pending_metrics'):
            return self.backend.get_pending_metrics()
        return {
            "memory_count_delta": 0,
            "memory_retrieval_count": getattr(self.backend, '_retrieval_count', 0)
        }

    async def mark_important(self, query: str, importance: float, source: str = 'user_explicit') -> dict:
        """Mark memories matching query with a specific importance score."""
        if not 0.0 <= importance <= 1.0:
            raise ValueError(f"Importance must be 0.0-1.0, got {importance}")

        matches = await self.backend.recall(query, limit=100, score_threshold=0.7)
        if not matches:
            return {"status": "success", "updated": 0, "message": "No matching memories found"}

        updated = 0
        for match in matches:
            chunk_id = match.get('chunk_id')
            if chunk_id:
                await self.backend.update_chunk_metadata(chunk_id, {
                    'importance': importance,
                    'source': source,
                })
                updated += 1

        return {"status": "success", "updated": updated, "message": f"Updated importance to {importance} for {updated} memories"}

    async def pin(self, query: str) -> dict:
        """Pin memories matching query (importance=1.0, never pruned)."""
        return await self.mark_important(query, 1.0, source='user_explicit')

    async def forget(self, query: str) -> dict:
        """Delete memories matching query."""
        matches = await self.backend.recall(query, limit=100, score_threshold=0.7)
        if not matches:
            return {"status": "success", "deleted": 0, "message": "No matching memories found"}

        chunk_ids = [m['chunk_id'] for m in matches if m.get('chunk_id')]
        await self.backend.delete_chunks(chunk_ids)

        return {"status": "success", "deleted": len(chunk_ids), "message": f"Deleted {len(chunk_ids)} memories"}

    def configure(self, **kwargs) -> dict:
        """Update plugin configuration at runtime. Supports: auto_curate."""
        updated = {}
        if 'auto_curate' in kwargs:
            value = kwargs['auto_curate']
            if isinstance(value, str) and value in ['on_stop', 'manual']:
                self.auto_curate = value
                updated['auto_curate'] = self.auto_curate
            else:
                raise ValueError("auto_curate must be 'on_stop' or 'manual'")
        return {"status": "success", "updated": updated}
