"""
Core ToxoLayer implementation.

This module contains the main ToxoLayer class which implements the trainable
intelligence layer that sits on top of Gemini API, providing the illusion
of a fine-tuned model through advanced prompt engineering and memory systems.
"""

import asyncio
import json
import logging
import os
import tempfile
import zipfile
from pathlib import Path
from typing import Dict, Any, List, Optional, Union, Tuple
from dataclasses import dataclass, asdict
from datetime import datetime
from typing import TYPE_CHECKING

try:
    import numpy as np  # type: ignore
    HAS_NUMPY = True
except Exception:  # pragma: no cover
    np = None  # type: ignore
    HAS_NUMPY = False

from .config import ToxoConfig
from .comprehensive_pipeline import ComprehensiveResponseMixin
from .package_writer import PackageWriterMixin
from .phase2_systems import Phase2SystemsMixin
from .rich_prompt_engine import RichPromptEngine
from .soft_prompt_guidance import convert_soft_prompt_to_guidance
from .multimodal import MultimodalMixin
from ..integrations.gemini_client import GeminiClient
from ..integrations.llm_client import MultiProviderLLMClient, create_llm_client
from ..training.soft_prompt import SoftPromptEngine
try:
    from ..training.reranker import ResponseReranker
    RERANKER_AVAILABLE = True
except ImportError:
    ResponseReranker = None
    RERANKER_AVAILABLE = False
from ..memory.domain_memory import DomainMemory
from ..utils.logger import get_logger, PerformanceTimer
from ..utils.exceptions import ToxoError, ModelError, InferenceError, ValidationError

try:
    from .quality_scorer import classify_query
except ImportError:
    classify_query = None  # Optional

# Response depth instructions for verbosity calibration
RESPONSE_DEPTH_INSTRUCTIONS = {
    "concise": (
        "[RESPONSE STYLE]\nBe concise. Answer in 1-3 sentences unless the question clearly requires more. "
        "Prioritize key facts and actionable points. Avoid over-explanation, preamble, or restatement."
    ),
    "balanced": (
        "[RESPONSE STYLE]\nMatch response length to question complexity. Be direct; prefer bullets over paragraphs. "
        "Include only essential detail—cap at 4-6 key points unless explicitly asked for more. Skip preamble."
    ),
    "detailed": (
        "[RESPONSE STYLE]\nProvide thorough explanations when useful. Include context, rationale, and examples when they add value."
    ),
    "procedural_short": (
        "[RESPONSE STYLE]\nProvide 6-8 numbered steps. One clear line per step. "
        "Include key parameters. Aim for 250-450 chars."
    ),
    "comparison": (
        "[RESPONSE STYLE]\nStructured comparison. Provide 4-6 bullets with one line of content each. "
        "Include: body, clarity, best for (adapt to topic). No long paragraphs."
    ),
}
DEPTH_MAX_TOKENS = {"concise": 512, "balanced": 768, "detailed": 1536, "procedural_short": 512, "comparison": 600}

# Query-type-specific instructions
QUERY_TYPE_INSTRUCTIONS = {
    "comparison": (
        "[COMPARISON FORMAT]\n"
        "Method A vs Method B. Bullets: Body, Clarity, Best for (adapt labels). "
        "One line per bullet with concrete contrast. No paragraphs. Prefer structured contrast over explanation."
    ),
    "technical": (
        "[TECHNICAL FORMAT]\nPrefer numeric targets over long explanation. Short. Dense. Expert."
    ),
    "technical_numeric": (
        "[TECHNICAL - NUMERIC ONLY]\nList only numeric values where relevant. No prose."
    ),
    "vague": (
        "[VAGUE QUERY]\nProvide domain assumption briefly, then offer 1 clarifying follow-up if helpful."
    ),
}

DEPTH_ESCALATION = {"concise": "balanced", "balanced": "detailed", "detailed": "detailed"}
MULTIPART_PATTERNS = ("compare", "vs", " versus ", " and ", "both", "each", "difference between", "similarities", "contrast")
TECHNICAL_INDICATORS = ("ratio", "percentage", "%", "target", "metric", "benchmark", "analysis", "breakdown", "step by step")


def _should_escalate_depth(query: str, requested_depth: str, domain: str = "general") -> str:
    """Detect multi-part/technical questions, escalate or use procedural_short. Returns effective depth."""
    if requested_depth == "detailed":
        return requested_depth
    q = query.lower().strip()
    words = len(q.split())
    procedural_markers = ("how to", "how do", "dial in", "process")
    asks_for_guide = any(m in q for m in ("step by step", "step-by-step", "guide", "detailed", "complete", "in-depth", "full"))
    if any(p in q for p in procedural_markers):
        if asks_for_guide or words > 5:
            pass
        else:
            return "procedural_short"
    numeric_request = any(n in q for n in ("ratio", "%", "target", "yield", "give"))
    technical = any(t in q for t in TECHNICAL_INDICATORS)
    if technical and numeric_request:
        return requested_depth
    comparison_markers = ("compare", " vs ", " versus ", "difference between", "similarities", "contrast")
    if any(c in q for c in comparison_markers):
        return "comparison"
    multipart = any(p in q for p in MULTIPART_PATTERNS)
    multi_question = q.count("?") >= 2 or (len(q) > 80 and "?" in q)
    if multipart or technical or multi_question:
        return DEPTH_ESCALATION.get(requested_depth, requested_depth)
    return requested_depth


@dataclass
class ToxoLayerMetadata:
    """Metadata for a Toxo layer."""
    name: str
    version: str
    created_at: str
    updated_at: str
    domain: str
    task_type: str = "general"
    description: str = ""
    tags: List[str] = None
    performance_metrics: Dict[str, float] = None
    training_examples: int = 0
    # User ownership fields (backward compatible with defaults)
    user_id: str = "anonymous"
    user_email: str = ""
    is_public: bool = True  # Default to public for backward compatibility
    
    def __post_init__(self):
        if self.tags is None:
            self.tags = []
        if self.performance_metrics is None:
            self.performance_metrics = {}


class ToxoLayer(PackageWriterMixin, MultimodalMixin, ComprehensiveResponseMixin, Phase2SystemsMixin):
    """
    Main Toxo Layer class implementing trainable intelligence layer.
    
    This class provides the core functionality for creating domain-specific
    AI layers on top of Gemini API through:
    - Soft prompt training and optimization
    - Hierarchical memory systems
    - Response reranking and filtering
    - Document-aware context injection
    - Real-time adaptation from feedback
    """
    
    def __init__(
        self,
        domain: str = "general",
        task_type: str = "general",
        config: Optional[ToxoConfig] = None,
        name: str = None,
        description: str = "",
        existing_soft_prompt: Optional["np.ndarray"] = None,
        user_id: str = "anonymous",
        user_email: str = "",
        is_public: bool = True,
        llm_provider: str = "gemini",
        llm_model: str = None,
        llm_api_key: str = None
    ):
        """
        Initialize a Toxo layer.
        
        Args:
            domain: Domain of expertise (e.g., 'finance', 'legal', 'medical')
            task_type: Type of task ('general', 'classification', 'generation', etc.)
            config: Toxo configuration object
            name: Human-readable name for the layer
            description: Description of the layer's purpose
            existing_soft_prompt: Pre-existing soft prompt for retraining
            user_id: ID of the user who owns this layer
            user_email: Email of the user who owns this layer
            is_public: Whether this layer is publicly accessible
            llm_provider: LLM provider ("gemini", "openai", "claude")
            llm_model: Specific model name (optional, uses defaults)
            llm_api_key: API key for the LLM provider
        """
        # Configuration
        self.config = config or ToxoConfig()
        self.domain = domain
        self.task_type = task_type
        self.name = name or f"{domain}_{task_type}_layer"
        self.description = description
        self.logger = get_logger(f"toxo.layer.{self.name}")
        
        # LLM Provider Configuration
        self.llm_provider = llm_provider
        self.llm_model = llm_model
        self.llm_api_key = llm_api_key
        self.llm_client = None
        
        # Training state
        self.is_trained = False
        self.training_history = []
        self.feedback_history = []  # PHASE 2 ENHANCEMENT: Track feedback data
        
        # Metadata (now includes user ownership)
        self.metadata = ToxoLayerMetadata(
            name=self.name,
            version="1.0.0",
            created_at=datetime.now().isoformat(),
            updated_at=datetime.now().isoformat(),
            domain=domain,
            task_type=task_type,
            description=description,
            user_id=user_id,
            user_email=user_email,
            is_public=is_public
        )
        
        # Initialize LLM client if API key and model provided
        if self.llm_api_key and self.llm_model:
            self.setup_llm_client(self.llm_provider, self.llm_model, self.llm_api_key)
        elif self.llm_api_key and not self.llm_model:
            self.logger.warning("API key provided but no model specified. Use setup_llm_client() to configure.")
        
        # Prompt config and static memory rules (set when loading from .toxo)
        self._prompt_config = {}
        self._static_memory_rules = {}
        
        # Initialize components with existing soft prompt if provided
        self._initialize_components(existing_soft_prompt=existing_soft_prompt)
        
        # Initialize comprehensive response pipeline
        self._initialize_comprehensive_response_pipeline()
        
        self.logger.info(f"Initialized Toxo Layer: {self.name} (domain: {domain}, task: {task_type}, user: {user_id})")
    
    def setup_llm_client(self, provider: str, model: str, api_key: str):
        """
        Setup LLM client for the layer.
        
        Args:
            provider: LLM provider ("gemini", "openai", "claude")
            model: Model name (user must specify)
            api_key: API key for the provider
        """
        if not api_key:
            raise ValueError("API key is required")
        if not model:
            raise ValueError("Model name is required. User must specify the model.")
        
        self.llm_provider = provider
        self.llm_model = model
        self.llm_api_key = api_key
        
        try:
            if provider == "gemini":
                # Use integrations GeminiClient (has .generate()) with updated config
                self.config.gemini.api_key = api_key
                self.config.gemini.model = model
                self.gemini_client = GeminiClient(self.config.gemini)
                self.llm_client = self.gemini_client
                self.logger.info(f"LLM client setup: {provider} with model {model}")
            else:
                self.llm_client = create_llm_client(provider, api_key, model)
                self.gemini_client = self.llm_client  # Multi-provider may need generate_text adapter
                self.logger.info(f"LLM client setup: {provider} with model {self.llm_client.get_provider_info()['model']}")
        except Exception as e:
            self.logger.error(f"Failed to setup LLM client: {str(e)}")
            raise
    
    def setup_api_key(self, api_key: str, model: str = None, provider: str = "gemini"):
        """
        Setup API key and model (docs-compatible alias for setup_llm_client).
        
        Args:
            api_key: API key for the LLM provider
            model: Model name (e.g. "gemini-2.5-flash-lite", "gpt-4")
            provider: LLM provider ("gemini", "openai", "claude"). Default "gemini".
        """
        model = model or (self.llm_model or "gemini-2.5-flash-lite")
        self.setup_llm_client(provider, model, api_key)
    
    @classmethod
    def load(cls, toxo_path: Union[str, Path], validate_integrity: bool = False) -> "ToxoLayer":
        """
        Load a .toxo package via ToxoPackageManager (full power: soft prompt, reranker, memory, prompt_config, memory_state).
        
        Uses ToxoPackageManager to load soft prompts, reranker, memory, prompt_config,
        and memory_state (training_contexts, key_rules, example_phrases).
        
        Args:
            toxo_path: Path to the .toxo file
            validate_integrity: If True, validate file checksums (skip if manifest has no checksums)
            
        Returns:
            ToxoLayer instance with all components loaded
            
        Example:
            layer = ToxoLayer.load("my_expert.toxo")
            layer.setup_api_key("your_key", "gemini-2.5-flash-lite")
            response = layer.query_sync("Your question")
        """
        from .package_manager import ToxoPackageManager
        pm = ToxoPackageManager()
        return pm.load_package(toxo_path, validate_integrity=validate_integrity)
    
    def _initialize_components(self, existing_soft_prompt: Optional["np.ndarray"] = None):
        """Initialize core components of the Toxo layer."""
        try:
            # LLM client (multi-provider or fallback to Gemini)
            if self.llm_client:
                # Use the configured multi-provider client
                self.gemini_client = self.llm_client
                self.logger.info(f"Using multi-provider LLM client: {self.llm_provider}")
            else:
                # Fallback to Gemini client for backward compatibility
                self.gemini_client = GeminiClient(self.config.gemini)
                self.logger.info("Using default Gemini client")
            
            # Rich prompt engine for advanced prompts
            self.rich_prompt_engine = RichPromptEngine(self.config)
            
            # Create master prompt for this domain
            self.master_prompt = self.rich_prompt_engine.create_rich_master_prompt(
                task_description=self.description,
                domain=self.domain,
                complexity_level="intermediate"
            )
            
            # Memory system
            self.memory = DomainMemory(self.config, domain=self.domain)
            
            # Soft prompt engine for trainable prompts - use existing prompt if provided
            self.soft_prompt_engine = SoftPromptEngine(
                dimension=self.config.memory.embedding_dimension,
                length=self.config.prompt.soft_prompt_length,
                config=self.config.training,
                prompt_config=self.config.prompt,
                existing_prompt=existing_soft_prompt
            )
            
            # Response reranker for quality improvement (optional)
            if RERANKER_AVAILABLE and ResponseReranker:
                self.reranker = ResponseReranker(self.config)
                self.logger.info("Response reranker initialized")
            else:
                self.reranker = None
                self.logger.warning("Response reranker not available - using single response mode")
            
            self.logger.info("Successfully initialized all available Toxo components")
            
        except Exception as e:
            self.logger.error(f"Failed to initialize Toxo components: {str(e)}")
            raise ModelError(f"Component initialization failed: {str(e)}")
        
        # PHASE 2 ENHANCEMENT: Initialize ALL advanced systems after core components
        self._initialize_cognitive_systems()
        self._initialize_memory_systems()
        self._initialize_agent_framework()
        enable_local_processing_pipeline = (
            os.getenv("TOXO_ENABLE_LOCAL_PROCESSING_PIPELINE", "0") == "1"
            or bool(getattr(self.config, "enable_local_processing_pipeline", False))
        )
        if enable_local_processing_pipeline:
            self._initialize_processing_pipeline()
        else:
            # Local document/layout pipeline is optional and noisy when dependencies
            # aren't installed. The default multimodal path sends bytes directly to Gemini.
            self.multimodal_processor = None
            self.document_processor = None
            self.advanced_document_processor = None
        self._initialize_enterprise_features()
    
    async def query(
        self,
        input_text: str,
        context: Optional[Dict[str, Any]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        response_depth: Optional[str] = None,
        **kwargs
    ) -> str:
        """
        Process a query through the Toxo layer.
        
        This is the main inference method that orchestrates all components
        to provide enhanced responses through the trained layer.
        
        Args:
            input_text: User's input query
            context: Additional context information
            temperature: Override default temperature
            max_tokens: Override default max tokens
            response_depth: Verbosity - 'concise', 'balanced', or 'detailed'
            **kwargs: Additional parameters for generation
            
        Returns:
            Enhanced response from the Toxo layer
        """
        depth = response_depth or getattr(getattr(self.config, "prompt", None), "response_depth", None) or "balanced"
        if depth not in RESPONSE_DEPTH_INSTRUCTIONS:
            depth = "balanced"
        effective_depth = _should_escalate_depth(input_text, depth, self.domain)
        # Ensure the LLM is constrained directly (avoid post-truncation).
        if max_tokens is None:
            max_tokens = min(
                getattr(getattr(self.config, "gemini", None), "max_tokens", 8192),
                DEPTH_MAX_TOKENS.get(effective_depth, 768),
            )

        with PerformanceTimer(self.logger, "query"):
            try:
                self.logger.debug(f"Processing query: {input_text[:100]}...")
                
                # Step 1: Process and enrich input
                enriched_input = await self._process_input(input_text, context, response_depth=effective_depth)
                
                # Step 2: Generate multiple candidate responses
                candidates = await self._generate_candidates(
                    enriched_input, 
                    temperature, 
                    max_tokens, 
                    **kwargs
                )
                
                # Step 3: Rerank and select best response
                best_response = await self._select_best_response(
                    input_text, 
                    candidates
                )
                
                # Step 4: Post-process and format output
                final_response = await self._post_process_response(
                    best_response, 
                    context,
                    response_depth=effective_depth
                )
                
                # Step 5: Update memory with this interaction
                await self._update_memory(input_text, final_response, context)
                
                self.logger.info("Query processed successfully")
                return final_response
                
            except Exception as e:
                self.logger.error(f"Query processing failed: {str(e)}")
                raise InferenceError(
                    f"Failed to process query: {str(e)}",
                    query=input_text
                )

    # ---------------------------------------------------------------------
    # Sync + alias helpers (docs / ergonomics)
    # ---------------------------------------------------------------------
    async def query_async(self, *args, **kwargs) -> str:
        """Alias for `query(...)` (async)."""
        return await self.query(*args, **kwargs)

    def query_sync(self, *args, **kwargs) -> str:
        """
        Synchronous wrapper for `query(...)`.

        Note: this cannot be called from inside an existing event loop.
        """
        try:
            asyncio.get_running_loop()
            raise RuntimeError("query_sync() cannot run inside an active event loop. Use `await layer.query(...)`.")
        except RuntimeError as e:
            # If no loop is running, get_running_loop() raises; that is the good path.
            if str(e).startswith("no running event loop"):
                return asyncio.run(self.query(*args, **kwargs))
            raise

    async def _process_input(
        self, 
        input_text: str, 
        context: Optional[Dict[str, Any]] = None,
        response_depth: str = "balanced"
    ) -> str:
        """
        Process and enrich the input with context and memory.
        
        Args:
            input_text: Original input text
            context: Additional context
            response_depth: Verbosity depth for response style instructions
            
        Returns:
            Enriched input with soft prompts and context
        """
        # Retrieve relevant memories (skip if memory not initialized)
        relevant_memories = []
        if self.memory and hasattr(self.memory, "retrieve"):
            try:
                relevant_memories = await self.memory.retrieve(
                    query=input_text,
                    limit=self.config.memory.max_retrievals
                )
            except Exception as e:
                self.logger.debug(f"Memory retrieve skipped: {e}")
        
        # Get soft prompt and convert to meaningful text guidance
        soft_prompt = self.soft_prompt_engine.get_current_prompt()
        
        # Build enriched context
        enriched_parts = []
        
        # Inject prompt_config (master prompt content) when loaded from .toxo
        prompt_config = getattr(self, "_prompt_config", {}) or {}
        master = prompt_config.get("master_prompt", {})
        if isinstance(master, dict):
            content = master.get("content")
            if not content:
                # Build from identity/instructions when "content" is absent (common .toxo format)
                parts = []
                ident = master.get("identity") or {}
                if isinstance(ident, dict):
                    if ident.get("core_role"):
                        parts.append(f"Role: {ident['core_role']}")
                    if ident.get("specialization"):
                        parts.append(f"Specialization: {ident['specialization']}")
                    if ident.get("communication_style"):
                        parts.append(f"Style: {ident['communication_style']}")
                instr = master.get("instructions") or {}
                if isinstance(instr, dict) and instr.get("primary_directive"):
                    parts.append(instr["primary_directive"])
                content = "\n".join(parts) if parts else None
            if content:
                enriched_parts.append(f"[LAYER SYSTEM INSTRUCTION]\n{content}")
        
        # Inject static memory rules (training_contexts, key_rules, example_phrases)
        static = getattr(self, "_static_memory_rules", {}) or {}
        if static:
            rules_parts = []
            for ctx in static.get("training_contexts", [])[:10]:
                rules_parts.append(f"• {ctx}")
            for k, v in (static.get("key_rules") or {}).items():
                if v:
                    rules_parts.append(f"• {k}: {v}")
            for k, v in (static.get("example_phrases") or {}).items():
                if v:
                    rules_parts.append(f"• {k}: {v}")
            if rules_parts:
                enriched_parts.append(f"[LAYER RULES & EXAMPLES]\n" + "\n".join(rules_parts))
        
        # CRITICAL FIX: Convert soft prompt array to actual guidance text
        if soft_prompt is not None and getattr(soft_prompt, "size", 0) > 0:
            # Convert the trained soft prompt parameters into actual textual guidance
            soft_prompt_guidance = convert_soft_prompt_to_guidance(
                soft_prompt=soft_prompt,
                soft_prompt_length=getattr(self.soft_prompt_engine, "length", 0),
                domain=self.domain,
                training_history=getattr(self, "training_history", None),
                logger=self.logger,
            )
            if soft_prompt_guidance:
                enriched_parts.append(f"[TRAINED EXPERTISE GUIDANCE]\n{soft_prompt_guidance}")
        
        # Add relevant memories
        if relevant_memories:
            memory_context = "\n".join([
                f"Relevant Context: {mem.content}" for mem in relevant_memories[:3]
            ])
            enriched_parts.append(f"[CONTEXTUAL KNOWLEDGE]\n{memory_context}")
        
        # Add response depth / verbosity instruction (token efficiency calibration)
        depth_instruction = RESPONSE_DEPTH_INSTRUCTIONS.get(response_depth)
        if depth_instruction:
            enriched_parts.append(depth_instruction)

        # Query-type-specific format instructions (comparison, technical, vague)
        if classify_query:
            qtype = classify_query(input_text, self.domain)
            if qtype == "technical":
                q_low = input_text.lower()
                numeric_request = any(n in q_low for n in ("ratio", "%", "target", "yield", "give"))
                type_key = "technical_numeric" if numeric_request else "technical"
            else:
                type_key = qtype
            type_instruction = QUERY_TYPE_INSTRUCTIONS.get(type_key)
            if type_instruction:
                enriched_parts.append(type_instruction)

        # Add domain-specific context
        if self.domain != "general":
            domain_context = f"[DOMAIN SPECIALIZATION]\nYou are an expert in {self.domain}. Apply domain-specific knowledge and best practices."
            enriched_parts.append(domain_context)
        
        # Add the actual user query
        enriched_parts.append(f"[USER QUERY]\n{input_text}")
        
        enriched_input = "\n\n".join(enriched_parts)
        
        self.logger.debug(f"Input enriched with {len(relevant_memories)} memories and trained guidance")
        return enriched_input
    
    async def _generate_candidates(
        self,
        enriched_input: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> List[str]:
        """
        Generate multiple candidate responses.
        
        Args:
            enriched_input: Processed input with context
            temperature: Generation temperature
            max_tokens: Maximum tokens
            **kwargs: Additional generation parameters
            
        Returns:
            List of candidate responses
        """
        # Use config defaults if not specified
        temperature = temperature or self.config.gemini.temperature
        max_tokens = max_tokens or self.config.gemini.max_tokens
        
        # Generate primary response
        primary_response = await self.gemini_client.generate(
            prompt=enriched_input,
            temperature=temperature,
            max_tokens=max_tokens,
            **kwargs
        )
        
        candidates = [primary_response]
        
        # Generate additional candidates with variation if reranking is enabled
        if hasattr(self.reranker, 'enabled') and self.reranker.enabled:
            for i in range(2):  # Generate 2 additional candidates
                varied_response = await self.gemini_client.generate(
                    prompt=enriched_input,
                    temperature=min(temperature + 0.1 * (i + 1), 1.0),
                    max_tokens=max_tokens,
                    **kwargs
                )
                candidates.append(varied_response)
        
        self.logger.debug(f"Generated {len(candidates)} candidate responses")
        return candidates
    
    async def _select_best_response(
        self, 
        original_input: str, 
        candidates: List[str]
    ) -> str:
        """
        Select the best response from candidates using reranker.
        
        Args:
            original_input: Original user input
            candidates: List of candidate responses
            
        Returns:
            Best selected response
        """
        if len(candidates) == 1:
            return candidates[0]
        
        # Use reranker to select best response if available
        if self.reranker is not None:
            try:
                best_response = await self.reranker.rank_responses(
                    query=original_input,
                    responses=candidates,
                    context={"domain": self.domain}
                )
                self.logger.debug("Best response selected using reranker")
                return best_response
            except Exception as e:
                self.logger.warning(f"Reranking failed, using first candidate: {str(e)}")
                return candidates[0]
        else:
            # No reranker available, return first candidate
            self.logger.debug("No reranker available, using first candidate")
            return candidates[0]
    
    async def _post_process_response(
        self, 
        response: str, 
        context: Optional[Dict[str, Any]] = None,
        response_depth: str = "balanced"
    ) -> str:
        """
        Post-process the response for formatting, safety, and optional compression.
        
        Args:
            response: Raw response from LLM
            context: Additional context for processing
            response_depth: Current depth setting (affects compression)
            
        Returns:
            Post-processed response
        """
        processed_response = response.strip()
        # IMPORTANT: Do not truncate/compress generated responses to meet "concise" requirements.
        # We enforce length via max_tokens passed to the LLM upstream.
        return processed_response
    
    async def _update_memory(
        self, 
        input_text: str, 
        response: str, 
        context: Optional[Dict[str, Any]] = None
    ):
        """
        Update memory with the current interaction.
        
        Args:
            input_text: User input
            response: System response
            context: Additional context
        """
        try:
            await self.memory.store_interaction(
                input_text=input_text,
                response=response,
                context={
                    "domain": self.domain,
                    "timestamp": datetime.now().isoformat(),
                    "original_context": context or {}
                }
            )
            self.logger.debug("Memory updated with interaction")
        except Exception as e:
            self.logger.warning(f"Failed to update memory: {str(e)}")
    
    async def provide_feedback(
        self,
        input_text: str,
        expected_output: str,
        actual_output: Optional[str] = None,
        feedback_type: str = "correction",
        analysis_score: Optional[float] = None,
        suggestions: Optional[List[str]] = None,
        iteration: Optional[int] = None,
        **metadata
    ):
        """
        Provide feedback to improve the layer's performance.
        
        Args:
            input_text: Original input
            expected_output: Desired output
            actual_output: What the system actually produced
            feedback_type: Type of feedback ('correction', 'improvement', 'refinement')
            analysis_score: Quality score from self-analysis (0-10)
            suggestions: List of improvement suggestions
            iteration: Training iteration number
            **metadata: Additional feedback metadata
        """
        self.logger.info(f"Received {feedback_type} feedback")
        
        try:
            # Store feedback in memory with enhanced parameters
            await self.memory.store_feedback(
                input_text=input_text,
                expected_output=expected_output,
                actual_output=actual_output,
                feedback_type=feedback_type,
                analysis_score=analysis_score,
                suggestions=suggestions,
                iteration=iteration,
                metadata=metadata
            )
            
            # Update soft prompt based on feedback with enhanced parameters
            if feedback_type in ["correction", "improvement", "refinement"]:
                await self.soft_prompt_engine.update_from_feedback(
                    input_text=input_text,
                    expected_output=expected_output,
                    actual_output=actual_output,
                    analysis_score=analysis_score,
                    suggestions=suggestions
                )
            
            # Update reranker if applicable
            if self.reranker is not None and hasattr(self.reranker, 'update_from_feedback'):
                await self.reranker.update_from_feedback(
                    input_text=input_text,
                    expected_output=expected_output,
                    actual_output=actual_output,
                    feedback_type=feedback_type
                )
            
            self.logger.info(f"Feedback processed and applied - Score: {analysis_score}, Suggestions: {len(suggestions or [])}")
            
        except Exception as e:
            self.logger.error(f"Failed to process feedback: {str(e)}")
            raise ToxoError(f"Feedback processing failed: {str(e)}")

    def provide_feedback_sync(self, *args, **kwargs):
        """Synchronous wrapper for `provide_feedback(...)`."""
        try:
            asyncio.get_running_loop()
            raise RuntimeError("provide_feedback_sync() cannot run inside an active event loop. Use `await layer.provide_feedback(...)`.")
        except RuntimeError as e:
            if str(e).startswith("no running event loop"):
                return asyncio.run(self.provide_feedback(*args, **kwargs))
            raise

    async def add_feedback(
        self,
        question: str,
        response: str,
        rating: Optional[float] = None,
        suggestions: Optional[List[str]] = None,
        expected_output: Optional[str] = None,
        **metadata,
    ):
        """
        Docs-compatible feedback helper.

        If `expected_output` is provided, this routes into `provide_feedback(...)` so the layer can learn.
        Otherwise, it records feedback as metadata without attempting training updates.
        """
        # If we have a target/expected output, use the training-enabled path.
        if expected_output is not None:
            return await self.provide_feedback(
                input_text=question,
                expected_output=expected_output,
                actual_output=response,
                feedback_type="improvement",
                analysis_score=rating,
                suggestions=suggestions,
                **metadata,
            )

        # Otherwise: store lightweight feedback record (non-training).
        record = {
            "question": question,
            "response": response,
            "rating": rating,
            "suggestions": suggestions or [],
            "metadata": metadata or {},
            "timestamp": datetime.now().isoformat(),
        }
        try:
            self.feedback_history.append(record)
        except Exception:
            self.feedback_history = [record]
        # Also store into memory if available
        if getattr(self, "memory", None) and hasattr(self.memory, "store_interaction"):
            try:
                await self.memory.store_interaction(
                    input_text=question,
                    response=response,
                    context={"feedback": record},
                )
            except Exception:
                pass

    def add_feedback_sync(self, *args, **kwargs):
        """Synchronous wrapper for add_feedback()."""
        return asyncio.run(self.add_feedback(*args, **kwargs))
    
    # Packaging methods live in `PackageWriterMixin` (see `toxo/core/package_writer.py`).
    def get_info(self) -> Dict[str, Any]:
        """
        Get a stable, docs-friendly summary of this layer.

        This is kept for backward compatibility with the CLI, README, and examples.
        """
        info: Dict[str, Any] = {
            "name": getattr(self, "name", None),
            "domain": getattr(self, "domain", None),
            "task_type": getattr(self, "task_type", None),
            "description": getattr(self, "description", ""),
            "is_trained": getattr(self, "is_trained", False),
            "version": getattr(getattr(self, "metadata", None), "version", None),
            "created_at": getattr(getattr(self, "metadata", None), "created_at", None),
            "updated_at": getattr(getattr(self, "metadata", None), "updated_at", None),
            "llm_provider": getattr(self, "llm_provider", None),
            "llm_model": getattr(self, "llm_model", None) or getattr(getattr(getattr(self, "config", None), "gemini", None), "model", None),
        }
        # Include ownership fields when present (older packages may not have them).
        meta = getattr(self, "metadata", None)
        if meta is not None:
            info.update(
                {
                    "user_id": getattr(meta, "user_id", "anonymous"),
                    "user_email": getattr(meta, "user_email", ""),
                    "is_public": getattr(meta, "is_public", True),
                    "tags": getattr(meta, "tags", []) or [],
                    "performance_metrics": getattr(meta, "performance_metrics", {}) or {},
                    "training_examples": getattr(meta, "training_examples", 0) or 0,
                }
            )
        return info

    def get_status(self) -> Dict[str, Any]:
        """Get current status of the Toxo layer."""
        return {
            "name": self.name,
            "domain": self.domain,
            "task_type": self.task_type,
            "is_trained": self.is_trained,
            "created_at": self.metadata.created_at,
            "updated_at": self.metadata.updated_at,
            "performance_metrics": self.metadata.performance_metrics,
            "training_examples": self.metadata.training_examples,
            "components": {
                "rich_prompt_engine": self.rich_prompt_engine is not None,
                "soft_prompt_engine": self.soft_prompt_engine is not None,
                "memory": self.memory is not None,
                "reranker": self.reranker is not None,
                "document_processor": getattr(self, 'document_processor', None) is not None
            }
        }

    async def get_training_statistics(self) -> Dict[str, Any]:
        """
        Get comprehensive training statistics and learning patterns.
        
        Returns:
            Dictionary with training statistics and learning insights
        """
        stats = {
            "layer_info": {
                "name": self.name,
                "domain": self.domain,
                "task_type": self.task_type,
                "is_trained": self.is_trained
            },
            "soft_prompt_stats": {},
            "memory_stats": {},
            "learning_patterns": {},
            "performance_metrics": self.metadata.performance_metrics.copy()
        }
        
        # Get soft prompt statistics
        if self.soft_prompt_engine:
            stats["soft_prompt_stats"] = self.soft_prompt_engine.get_training_stats()
            
            # Add parameter analysis
            current_prompt = self.soft_prompt_engine.get_current_prompt()
            if current_prompt is not None:
                stats["soft_prompt_stats"].update({
                    "parameter_count": current_prompt.size,
                    "parameter_mean": float(current_prompt.mean()),
                    "parameter_std": float(current_prompt.std()),
                    "parameter_range": [float(current_prompt.min()), float(current_prompt.max())]
                })
        
        # Get memory statistics
        if self.memory:
            if hasattr(self.memory, 'get_statistics'):
                stats["memory_stats"] = self.memory.get_statistics()
            
            # Get learning patterns from feedback
            if hasattr(self.memory, 'get_learning_patterns'):
                stats["learning_patterns"] = await self.memory.get_learning_patterns()
        
        return stats

    async def get_improvement_suggestions(self) -> Dict[str, Any]:
        """
        Get AI-powered suggestions for improving the layer's performance.
        
        Returns:
            Dictionary with improvement suggestions and recommendations
        """
        suggestions = {
            "parameter_optimization": [],
            "training_recommendations": [],
            "memory_optimization": [],
            "performance_insights": []
        }
        
        # Analyze soft prompt training
        if self.soft_prompt_engine:
            training_stats = self.soft_prompt_engine.get_training_stats()
            
            # Check learning rate effectiveness
            if training_stats.get("iteration", 0) > 10:
                recent_losses = training_stats.get("loss_history", [])[-5:]
                if recent_losses and len(set(recent_losses)) == 1:
                    suggestions["parameter_optimization"].append(
                        "Learning rate may be too low - consider increasing for better adaptation"
                    )
                elif recent_losses and max(recent_losses) - min(recent_losses) > 0.5:
                    suggestions["parameter_optimization"].append(
                        "High loss variance detected - consider reducing learning rate for stability"
                    )
            
            # Check parameter adaptation
            current_prompt = self.soft_prompt_engine.get_current_prompt()
            if current_prompt is not None:
                param_std = float(current_prompt.std())
                if param_std < 0.001:
                    suggestions["parameter_optimization"].append(
                        "Low parameter variance - consider increasing perturbation scale"
                    )
                elif param_std > 0.1:
                    suggestions["parameter_optimization"].append(
                        "High parameter variance - consider adding regularization"
                    )
        
        # Analyze memory and learning patterns
        if self.memory and hasattr(self.memory, 'get_learning_patterns'):
            learning_patterns = await self.memory.get_learning_patterns()
            
            feedback_count = learning_patterns.get("total_feedback", 0)
            if feedback_count < 10:
                suggestions["training_recommendations"].append(
                    f"Limited feedback data ({feedback_count} items) - consider providing more training examples"
                )
            
            improvement_trend = learning_patterns.get("improvement_trend", "stable")
            if improvement_trend == "declining":
                suggestions["training_recommendations"].append(
                    "Performance declining - consider reviewing recent feedback quality"
                )
            elif improvement_trend == "stable":
                suggestions["training_recommendations"].append(
                    "Performance stable - consider introducing more diverse training examples"
                )
            
            # Memory optimization suggestions
            if hasattr(self.memory, 'get_statistics'):
                memory_stats = self.memory.get_statistics()
                total_memories = memory_stats.get("total_memories", 0)
                
                if total_memories > 1000:
                    suggestions["memory_optimization"].append(
                        "Large memory size - consider running cleanup to improve performance"
                    )
                elif total_memories < 50:
                    suggestions["memory_optimization"].append(
                        "Limited memory data - performance may improve with more interactions"
                    )
        
        # Performance insights
        latest_score = None
        if self.memory and hasattr(self.memory, 'get_learning_patterns'):
            learning_patterns = await self.memory.get_learning_patterns()
            latest_score = learning_patterns.get("latest_score")
        
        if latest_score is not None:
            if latest_score >= 8.5:
                suggestions["performance_insights"].append(
                    f"Excellent performance (score: {latest_score}/10) - consider fine-tuning for edge cases"
                )
            elif latest_score >= 7.0:
                suggestions["performance_insights"].append(
                    f"Good performance (score: {latest_score}/10) - focus on consistency improvements"
                )
            else:
                suggestions["performance_insights"].append(
                    f"Performance needs improvement (score: {latest_score}/10) - consider additional training"
                )
        
        return suggestions

    async def optimize_performance(self) -> Dict[str, Any]:
        """
        Automatically optimize the layer's performance based on current statistics.
        
        Returns:
            Dictionary with optimization results and actions taken
        """
        optimization_results = {
            "actions_taken": [],
            "performance_before": {},
            "performance_after": {},
            "recommendations": []
        }
        
        # Get current performance baseline
        training_stats = await self.get_training_statistics()
        optimization_results["performance_before"] = training_stats.copy()
        
        # Memory optimization
        if self.memory and hasattr(self.memory, 'cleanup_old_memories'):
            try:
                cleaned_count = await self.memory.cleanup_old_memories(
                    max_age_days=180,  # 6 months
                    min_importance=0.2,
                    max_memories=5000
                )
                if cleaned_count > 0:
                    optimization_results["actions_taken"].append(
                        f"Cleaned up {cleaned_count} old memories"
                    )
            except Exception as e:
                self.logger.warning(f"Memory cleanup failed: {e}")
        
        # Soft prompt optimization
        if self.soft_prompt_engine:
            training_stats = self.soft_prompt_engine.get_training_stats()
            
            # Reset if performance is poor and many iterations have passed
            if (training_stats.get("iteration", 0) > 100 and 
                training_stats.get("best_loss", float('inf')) > 0.8):
                
                # Save current state before reset
                current_state = self.soft_prompt_engine.save_state()
                
                # Reset and reinitialize
                self.soft_prompt_engine.reset()
                optimization_results["actions_taken"].append(
                    "Reset soft prompt engine due to poor convergence"
                )
        
        # Get updated performance
        updated_stats = await self.get_training_statistics()
        optimization_results["performance_after"] = updated_stats
        
        # Generate recommendations
        suggestions = await self.get_improvement_suggestions()
        optimization_results["recommendations"] = suggestions
        
        return optimization_results
    
    @staticmethod
    def load(path: Union[str, Path], config: Optional[ToxoConfig] = None) -> 'ToxoLayer':
        """Load a ToxoLayer from a .toxo file (full power via ToxoPackageManager)."""
        from .package_manager import ToxoPackageManager
        path = Path(path)
        if not path.exists():
            raise ToxoError(f"Toxo file not found: {path}")
        pm = ToxoPackageManager()
        return pm.load_package(path, validate_integrity=False)
    
    @staticmethod
    def _load_legacy(path: Union[str, Path], config: Optional[ToxoConfig] = None) -> 'ToxoLayer':
        """Legacy load path (kept for reference)."""
        from .legacy_loader import load_legacy_layer

        return load_legacy_layer(path, config=config)
    
    def __repr__(self) -> str:
        """String representation of the layer."""
        return f"ToxoLayer(name='{self.name}', domain='{self.domain}', trained={self.is_trained})" 
