"""
GitMem Coding - Coding Context Store

Local JSON-based storage backend for coding context data.
Caches file contents read by AI agents for cross-session retrieval,
reducing token usage by avoiding redundant file reads.

Storage Structure:
    .gitmem_coding/
    ├── agents/
    │   └── {agent_id}/
    │       ├── file_contexts.json    # Cached file contents + metadata
    │       ├── coding_sessions.json  # Session read/retrieve logs
    │       └── settings.json         # Agent-specific config
    ├── index/
    │   └── file_index.json           # Path-based lookup index
    └── config.json                   # Global config
"""

import os
import json
import uuid
import hashlib
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
from pathlib import Path
import threading


from .models import FileContext, CodingSession, ContextStatus, TOKENS_PER_CHAR_RATIO, CodeChunk
from .chunking_engine import ChunkingEngine
from .ast_skeleton import CodeFlowGenerator, detect_language, retrieve_path, BSTIndex

# Staleness threshold — files not accessed in this many days are considered stale
STALE_DAYS_THRESHOLD = 30


class CodingContextStore:
    """
    Local JSON-based storage for coding context data.
    
    Thread-safe storage with automatic persistence.
    Provides hash-based freshness detection and semantic chunking.
    
    Features:
        - File content CRUD operations
        - Semantic Chunking & Deduplication
        - Hash-based staleness detection
        - Keyword search across cached files
        - Token savings tracking
        - Session-level analytics
    """
    
    def __init__(self, root_path: str = "./.gitmem_coding"):
        """
        Initialize the coding context store.
        
        Args:
            root_path: Root directory for storage (defaults to .gitmem_coding)
        """
        self.root_path = Path(root_path).absolute()
        self.agents_path = self.root_path / "agents"
        self.index_path = self.root_path / "index"
        self.chunks_path = self.root_path / "chunks.json" # Global chunk registry
        self.config_path = self.root_path / "config.json"
        
        # Thread safety
        self._lock = threading.RLock()
        
        # Initialize directories
        self._ensure_dirs()
        self.global_index_path = self.index_path / "global_index.json"
        
        self._load_config()
        self._global_index = self._load_global_index()
    
    # =========================================================================
    # Internal Setup
    # =========================================================================
    
    def _ensure_dirs(self):
        """Create necessary directory structure."""
        self.root_path.mkdir(parents=True, exist_ok=True)
        self.agents_path.mkdir(exist_ok=True)
        self.index_path.mkdir(exist_ok=True)
        
        # Initialize global chunk registry if missing
        if not self.chunks_path.exists():
            with open(self.chunks_path, 'w', encoding='utf-8') as f:
                json.dump({}, f)

    def _load_config(self):
        """Load global configuration."""
        if self.config_path.exists():
            with open(self.config_path, 'r', encoding='utf-8') as f:
                self._config = json.load(f)
        else:
            self._config = {
                "version": "1.0.0",
                "created_at": datetime.now().isoformat(),
                "stale_days_threshold": STALE_DAYS_THRESHOLD,
                "module": "gitmem_coding"
            }
            self._save_config()
    
    def _save_config(self):
        """Save global configuration."""
        with open(self.config_path, 'w', encoding='utf-8') as f:
            json.dump(self._config, f, indent=2)
            
    def _load_chunks(self) -> Dict[str, Any]:
        """Load global chunk registry."""
        with self._lock:
            if self.chunks_path.exists():
                try:
                    with open(self.chunks_path, 'r', encoding='utf-8') as f:
                        return json.load(f)
                except json.JSONDecodeError:
                    return {}
            return {}

    def _save_chunks(self, chunks: Dict[str, Any]):
        """Save global chunk registry."""
        with self._lock:
            with open(self.chunks_path, 'w', encoding='utf-8') as f:
                json.dump(chunks, f, indent=2)

    def _load_global_index(self) -> Dict[str, List[str]]:
        """Load global inverted index (Symbol -> [FilePaths])."""
        with self._lock:
            if self.global_index_path.exists():
                try:
                    with open(self.global_index_path, 'r', encoding='utf-8') as f:
                        return json.load(f)
                except json.JSONDecodeError:
                    return {}
            return {}

    def _save_global_index(self):
        """Save global inverted index."""
        with self._lock:
            with open(self.global_index_path, 'w', encoding='utf-8') as f:
                json.dump(self._global_index, f, indent=2)

    def _update_global_index(self, file_path: str, symbols: List[str]):
        """
        Update global index for a file.
        1. Remove file from all entries (cleanup old).
        2. Add file to new symbol entries.
        """
        normalized_path = os.path.normpath(file_path)
        
        # 1. Cleanup old
        self._remove_from_global_index(normalized_path, save=False)
        
        # 2. Add new
        for symbol in symbols:
            if symbol not in self._global_index:
                self._global_index[symbol] = []
            if normalized_path not in self._global_index[symbol]:
                self._global_index[symbol].append(normalized_path)
        
        self._save_global_index()

    def _remove_from_global_index(self, file_path: str, save: bool = True):
        """Remove a file from the global index."""
        normalized_path = os.path.normpath(file_path)
        empty_keys = []
        
        for symbol, files in self._global_index.items():
            if normalized_path in files:
                files.remove(normalized_path)
                if not files:
                    empty_keys.append(symbol)
        
        for k in empty_keys:
            del self._global_index[k]
            
        if save:
            self._save_global_index()

    def _get_agent_path(self, agent_id: str) -> Path:
        """Get the storage path for an agent."""
        safe_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in agent_id)
        return self.agents_path / safe_id
    
    def _ensure_agent(self, agent_id: str):
        """Ensure agent directory and files exist."""
        agent_path = self._get_agent_path(agent_id)
        agent_path.mkdir(exist_ok=True)
        
        defaults = {
            "file_contexts.json": [],
            "coding_sessions.json": [],
            "settings.json": {"created_at": datetime.now().isoformat()}
        }
        
        for filename, default_content in defaults.items():
            filepath = agent_path / filename
            if not filepath.exists():
                with open(filepath, 'w', encoding='utf-8') as f:
                    json.dump(default_content, f, indent=2)
    
    def _load_agent_data(self, agent_id: str, data_type: str) -> list:
        """Load agent data from JSON file."""
        self._ensure_agent(agent_id)
        filepath = self._get_agent_path(agent_id) / f"{data_type}.json"
        
        with self._lock:
            if filepath.exists():
                with open(filepath, 'r', encoding='utf-8') as f:
                    return json.load(f)
        return []
    
    def _save_agent_data(self, agent_id: str, data_type: str, data: list):
        """Save agent data to JSON file."""
        self._ensure_agent(agent_id)
        filepath = self._get_agent_path(agent_id) / f"{data_type}.json"
        
        with self._lock:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, default=str)
    
    # =========================================================================
    # File Context Operations
    # =========================================================================
    
    def store_file_context(
        self,
        agent_id: str,
        file_path: str,
        content: str,
        language: str = "other",
        session_id: str = "",
        keywords: List[str] = None,
        content_summary: str = "",
        storage_mode: str = "skeleton"
    ) -> Dict[str, Any]:
        """
        Store or update file content in the coding context cache.
        
        Uses CodeFlowGenerator to create a structured Flow Tree and BST Index.
        Stores the tree and index as JSON.
        """
        content_hash = FileContext.compute_hash(content)
        file_name = os.path.basename(file_path)
        normalized_path = os.path.normpath(file_path)
        original_token_estimate = FileContext.estimate_tokens(content)
        
        # ── Generate Code Flow Tree & Index ──
        flow_data = {}
        skeleton_token_estimate = 0
        compression_ratio = 0.0
        
        if storage_mode == "skeleton" or storage_mode == "flow":
            # Use new CodeFlowGenerator
            generator = CodeFlowGenerator()
            flow_data = generator.generate(content, file_path)
            
            # Serialize for storage
            # We store the tree and index in 'compact_skeleton' field for compatibility 
            # (or add new fields, but let's reuse 'compact_skeleton' as a JSON string for now to minimize schema changes)
            compact_skeleton = json.dumps(flow_data)
            
            # --- Global Index Update ---
            if "index" in flow_data:
                # BST Index is { "key": ..., "left": ..., "right": ... }
                # We need to traverse it to get all keys
                def _get_all_keys(node):
                    keys = []
                    if not node: return keys
                    keys.append(node["key"])
                    keys.extend(_get_all_keys(node["left"]))
                    keys.extend(_get_all_keys(node["right"]))
                    return keys
                
                all_symbols = _get_all_keys(flow_data["index"])
                self._update_global_index(normalized_path, all_symbols)
            # ---------------------------
            
            skeleton_token_estimate = len(compact_skeleton) // 4 # Rough token estimate
            if original_token_estimate > 0:
                compression_ratio = round(
                    1.0 - (skeleton_token_estimate / original_token_estimate), 4
                )
        else:
            compact_skeleton = ""

        # Get file modification time if accessible
        file_mod_time = ""
        try:
            if os.path.exists(normalized_path):
                file_mod_time = datetime.fromtimestamp(
                    os.path.getmtime(normalized_path)
                ).isoformat()
        except (OSError, ValueError):
            pass
        
        # --- Chunking & Deduplication (Keep existing logic) ---
        chunker = ChunkingEngine.get_chunker(language)
        chunks = chunker.chunk_file(content, file_path)
        
        chunk_registry = self._load_chunks()
        new_chunks_count = 0
        file_chunk_hashes = []
        
        for chunk in chunks:
            file_chunk_hashes.append(chunk.hash_id)
            if chunk.hash_id not in chunk_registry:
                chunk_registry[chunk.hash_id] = chunk.to_dict()
                new_chunks_count += 1
        
        if new_chunks_count > 0:
            self._save_chunks(chunk_registry)
        
        # ── Determine what to persist ──
        stored_content = content if storage_mode == "full" else ""
        effective_token_estimate = (
            original_token_estimate if storage_mode == "full" 
            else skeleton_token_estimate
        )
        
        # --- File Context Update ---
        contexts = self._load_agent_data(agent_id, "file_contexts")
        
        existing_idx = None
        for i, ctx in enumerate(contexts):
            if os.path.normpath(ctx.get("file_path", "")) == normalized_path:
                existing_idx = i
                break
        
        now = datetime.now().isoformat()
        
        # Reuse or Create Context Logic
        if existing_idx is not None:
            existing = contexts[existing_idx]
            was_same_content = existing.get("content_hash") == content_hash
            
            existing.update({
                "content": stored_content,
                "content_hash": content_hash,
                "compact_skeleton": compact_skeleton, # Stores the JSON tree string
                "storage_mode": storage_mode,
                "size_bytes": len(content.encode('utf-8')),
                "line_count": content.count('\n') + 1,
                "token_estimate": effective_token_estimate,
                "original_token_estimate": original_token_estimate,
                "skeleton_token_estimate": skeleton_token_estimate,
                "compression_ratio": compression_ratio,
                "last_accessed_at": now,
                "access_count": existing.get("access_count", 0) + 1,
                "file_modified_at": file_mod_time,
                "language": language,
                "session_id": session_id,
                "chunk_hashes": file_chunk_hashes
            })
            
            if content_summary: existing["content_summary"] = content_summary
            if keywords:
                old_kw = set(existing.get("keywords", []))
                old_kw.update(keywords)
                existing["keywords"] = list(old_kw)
            
            contexts[existing_idx] = existing
            self._save_agent_data(agent_id, "file_contexts", contexts)
            
            return {
                "status": "updated",
                "context_id": existing["id"],
                "file_path": normalized_path,
                "content_changed": not was_same_content,
                "storage_mode": storage_mode,
                "token_estimate": effective_token_estimate,
                "compression_ratio": compression_ratio,
                "message": f"File Code Flow updated."
            }
        else:
            context_id = str(uuid.uuid4())
            new_context = {
                "id": context_id,
                "file_path": normalized_path,
                "file_name": file_name,
                "content": stored_content,
                "content_hash": content_hash,
                "content_summary": content_summary,
                "compact_skeleton": compact_skeleton,
                "storage_mode": storage_mode,
                "language": language,
                "line_count": content.count('\n') + 1,
                "size_bytes": len(content.encode('utf-8')),
                "agent_id": agent_id,
                "session_id": session_id,
                "access_count": 1,
                "token_estimate": effective_token_estimate,
                "original_token_estimate": original_token_estimate,
                "skeleton_token_estimate": skeleton_token_estimate,
                "compression_ratio": compression_ratio,
                "created_at": now,
                "last_accessed_at": now,
                "file_modified_at": file_mod_time,
                "keywords": keywords or [],
                "chunk_hashes": file_chunk_hashes
            }
            contexts.append(new_context)
            self._save_agent_data(agent_id, "file_contexts", contexts)
            
            return {
                "status": "created",
                "context_id": context_id,
                "file_path": normalized_path,
                "storage_mode": storage_mode,
                "token_estimate": effective_token_estimate,
                "compression_ratio": compression_ratio,
                "message": f"File Code Flow created."
            }

    def store_file_context_from_path(
        self,
        agent_id: str,
        file_path: str,
        language: str = "auto",
        session_id: str = "",
        keywords: List[str] = None,
        content_summary: str = "",
        storage_mode: str = "skeleton"
    ) -> Dict[str, Any]:
        """
        Store a file by reading it from disk.
        """
        normalized_path = os.path.normpath(file_path)
        
        if not os.path.exists(normalized_path):
            return {
                "status": "error",
                "error": f"File not found: {normalized_path}",
                "file_path": normalized_path
            }
        
        try:
            with open(normalized_path, 'r', encoding='utf-8', errors='replace') as f:
                content = f.read()
        except (OSError, IOError) as e:
            return {
                "status": "error",
                "error": f"Failed to read file: {e}",
                "file_path": normalized_path
            }
        
        if language == "auto":
            language = detect_language(file_path)
        
        return self.store_file_context(
            agent_id=agent_id,
            file_path=file_path,
            content=content,
            language=language,
            session_id=session_id,
            keywords=keywords,
            content_summary=content_summary,
            storage_mode=storage_mode
        )

    # Helper method found in analysis of ast_skeleton import replacement
    def retrieve_file_context(
        self,
        agent_id: str,
        file_path: str
    ) -> Dict[str, Any]:
        """
        Retrieve cached file content (Code Flow Tree) with freshness check.
        """
        normalized_path = os.path.normpath(file_path)
        contexts = self._load_agent_data(agent_id, "file_contexts")
        
        found = next((ctx for ctx in contexts if os.path.normpath(ctx.get("file_path", "")) == normalized_path), None)
        
        if found is None:
            return {
                "status": "cache_miss",
                "file_path": normalized_path,
                "message": "File not found in coding context cache."
            }
        
        # Check freshness (simplified for this task)
        freshness_status = ContextStatus.UNKNOWN.value
        try:
             if os.path.exists(normalized_path):
                with open(normalized_path, 'r', encoding='utf-8', errors='replace') as f:
                    current_hash = FileContext.compute_hash(f.read())
                    freshness_status = ContextStatus.FRESH.value if current_hash == found.get("content_hash") else ContextStatus.STALE.value
             else:
                freshness_status = ContextStatus.MISSING.value
        except:
            pass

        return {
            "status": "cache_hit",
            "freshness": freshness_status,
            "file_path": normalized_path,
            "code_flow": json.loads(found.get("compact_skeleton", "{}")), # Return parsed JSON tree
            "storage_mode": found.get("storage_mode"),
            "message": f"Code Flow retrieved. Status: {freshness_status}"
        }
    
    def search_code_flow(
        self,
        agent_id: str,
        query: str
    ) -> List[Dict[str, Any]]:
        """
        Search for code paths using the BST index.
        Returns a list of paths from Root to the matching Node.
        """
        contexts = self._load_agent_data(agent_id, "file_contexts")
        results = []
        
        # Optimize with Global Index
        relevant_files = self._global_index.get(query, [])
        
        # Helper to load specific file context
        def _get_context_by_path(path):
            # Optimally we should have a map, but iterating list is OK for local scope
            # Or better: construct a temporary map once?
            # Actually, `contexts` is a list.
            for c in contexts:
                if os.path.normpath(c.get("file_path", "")) == path:
                    return c
            return None

        from .ast_skeleton import retrieve_path
        
        # If we have matches in global index, only search those files
        # But we must also support partial matches? 
        # The requirement says "string related queries". 
        # Global index is exact symbol match.
        # If query is not in global index, should we fallback to full scan? 
        # For now, let's assume exact symbol match for speed. 
        # OR: check if user wants fuzzy search. 
        # The prompt says "optimized to perform well with string related queries".
        # Let's try Global Index first. If result empty, maybe fallback?
        # NO, strict O(1) means trust index. 
        
        processed_files = set()
        
        if relevant_files:
            # Search only relevant files
            for file_path in relevant_files:
                ctx = _get_context_by_path(file_path)
                if not ctx: continue
                
                try:
                    flow_json = ctx.get("compact_skeleton", "{}")
                    if not flow_json: continue
                    flow_data = json.loads(flow_json)
                    
                    paths = retrieve_path(flow_data["tree"], flow_data["index"], query)
                    for path in paths:
                        results.append({
                             "file_path": ctx.get("file_path"),
                             "matched_symbol": query,
                             "flow_path": path
                        })
                    processed_files.add(file_path)
                except Exception:
                    continue
        else:
             # Fallback? Maybe strict index is better.
             # If I search "create_flow", I expect results.
             pass
             
        return results

    def list_code_flows(self, agent_id: str, limit: int = 50, offset: int = 0) -> Dict[str, Any]:
        """List all stored code flow structures."""
        contexts = self._load_agent_data(agent_id, "file_contexts")
        total = len(contexts)
        sliced = contexts[offset : offset + limit]
        
        summaries = []
        for ctx in sliced:
            summaries.append({
                "context_id": ctx.get("id"),
                "file_path": ctx.get("file_path"),
                "language": ctx.get("language"),
                "last_accessed": ctx.get("last_accessed_at"),
                "size_bytes": ctx.get("size_bytes", 0)
            })
            
        return {
            "total": total,
            "items": summaries
        }

    def delete_code_flow(self, agent_id: str, file_path: str) -> bool:
        """Delete a code flow entry."""
        contexts = self._load_agent_data(agent_id, "file_contexts")
        normalized_path = os.path.normpath(file_path)
        
        initial_len = len(contexts)
        contexts = [c for c in contexts if os.path.normpath(c.get("file_path", "")) != normalized_path]
        
        if len(contexts) < initial_len:
            self._save_agent_data(agent_id, "file_contexts", contexts)
            self._remove_from_global_index(normalized_path)
            return True
        return False


    
    # =========================================================================
    # Session Tracking
    # =========================================================================
    
    def record_session_activity(
        self,
        agent_id: str,
        session_id: str,
        action: str,
        file_path: str,
        tokens: int = 0
    ):
        """
        Record a file store/retrieve action within a session.
        
        Args:
            agent_id: The agent ID
            session_id: Current session ID
            action: "store" or "retrieve"
            file_path: File path involved
            tokens: Token count for the operation
        """
        sessions = self._load_agent_data(agent_id, "coding_sessions")
        
        # Find or create session
        session = None
        session_idx = None
        for i, s in enumerate(sessions):
            if s.get("session_id") == session_id:
                session = s
                session_idx = i
                break
        
        if session is None:
            session = {
                "session_id": session_id,
                "agent_id": agent_id,
                "files_stored": [],
                "files_retrieved": [],
                "tokens_stored": 0,
                "tokens_retrieved": 0,
                "cache_hits": 0,
                "cache_misses": 0,
                "started_at": datetime.now().isoformat(),
                "ended_at": ""
            }
            sessions.append(session)
            session_idx = len(sessions) - 1
        
        normalized_path = os.path.normpath(file_path)
        
        if action == "store":
            if normalized_path not in session["files_stored"]:
                session["files_stored"].append(normalized_path)
            session["tokens_stored"] += tokens
        elif action == "retrieve":
            if normalized_path not in session["files_retrieved"]:
                session["files_retrieved"].append(normalized_path)
            session["tokens_retrieved"] += tokens
            session["cache_hits"] += 1
        elif action == "miss":
            session["cache_misses"] += 1
        
        sessions[session_idx] = session
        
        # Keep only last 100 sessions
        if len(sessions) > 100:
            sessions = sessions[-100:]
        
        self._save_agent_data(agent_id, "coding_sessions", sessions)
    
    def get_token_savings_report(self, agent_id: str) -> Dict[str, Any]:
        """
        Get a comprehensive token savings report.
        
        Args:
            agent_id: The agent ID
        
        Returns:
            Dict with token savings breakdown
        """
        sessions = self._load_agent_data(agent_id, "coding_sessions")
        contexts = self._load_agent_data(agent_id, "file_contexts")
        
        total_tokens_stored = sum(s.get("tokens_stored", 0) for s in sessions)
        total_tokens_retrieved = sum(s.get("tokens_retrieved", 0) for s in sessions)
        total_cache_hits = sum(s.get("cache_hits", 0) for s in sessions)
        total_cache_misses = sum(s.get("cache_misses", 0) for s in sessions)
        
        # Current cache value
        cache_token_value = sum(c.get("token_estimate", 0) for c in contexts)
        
        return {
            "total_tokens_stored": total_tokens_stored,
            "total_tokens_retrieved_from_cache": total_tokens_retrieved,
            "estimated_token_savings": total_tokens_retrieved,
            "total_cache_hits": total_cache_hits,
            "total_cache_misses": total_cache_misses,
            "hit_rate": round(
                (total_cache_hits / (total_cache_hits + total_cache_misses) * 100)
                if (total_cache_hits + total_cache_misses) > 0 else 0, 2
            ),
            "current_cache_token_value": cache_token_value,
            "files_in_cache": len(contexts),
            "sessions_tracked": len(sessions)
        }
    
    # =========================================================================
    # Cleanup Operations
    # =========================================================================
    
    def cleanup_stale_contexts(
        self,
        agent_id: str,
        days_threshold: int = None
    ) -> Dict[str, Any]:
        """
        Remove cached contexts not accessed within the threshold period.
        
        Args:
            agent_id: The agent ID
            days_threshold: Days since last access to consider stale
        
        Returns:
            Dict with cleanup results
        """
        threshold = days_threshold or self._config.get("stale_days_threshold", STALE_DAYS_THRESHOLD)
        cutoff = (datetime.now() - timedelta(days=threshold)).isoformat()
        
        contexts = self._load_agent_data(agent_id, "file_contexts")
        original_count = len(contexts)
        
        # Keep contexts accessed after the cutoff
        fresh = [c for c in contexts if c.get("last_accessed_at", "") >= cutoff]
        removed_count = original_count - len(fresh)
        
        if removed_count > 0:
            self._save_agent_data(agent_id, "file_contexts", fresh)
        
        return {
            "status": "ok",
            "removed_count": removed_count,
            "remaining_count": len(fresh),
            "threshold_days": threshold
        }
    
    # =========================================================================
    # Agent Operations
    # =========================================================================
    
    def list_agents(self) -> List[str]:
        """List all agents with coding context data."""
        agents = []
        if self.agents_path.exists():
            for path in self.agents_path.iterdir():
                if path.is_dir():
                    agents.append(path.name)
        return agents
    
    def delete_agent_data(self, agent_id: str) -> bool:
        """Delete all coding context data for an agent."""
        import shutil
        agent_path = self._get_agent_path(agent_id)
        if agent_path.exists():
            shutil.rmtree(agent_path)
            return True
        return False

