# langchain.py - LangChain enhancement
# ============================================================================
# FILE: orchestra/adapters/langchain.py
# LangChain adapter supporting LLMChain and LCEL Runnables
# ============================================================================

import time
import logging
import json
from typing import Any, Dict, Optional, Union, List

from .base import BaseAdapter
from ..core.cache_manager import CacheManager
from ..core.metrics import MetricsTracker

logger = logging.getLogger(__name__)

class OrchestraLangChainConfig:
    def __init__(
        self,
        similarity_threshold: float = 0.92,
        cache_ttl: int = 3600,
        enable_compression: bool = True
    ):
        self.similarity_threshold = similarity_threshold
        self.cache_ttl = cache_ttl
        self.enable_compression = enable_compression

class EnhancedLangChain(BaseAdapter):
    """
    Adapter for LangChain (both Chains and Runnables).
    """
    
    def __init__(
        self, 
        chain_or_runnable: Any, 
        config: Optional[OrchestraLangChainConfig] = None
    ):
        super().__init__(config)
        self.chain = chain_or_runnable
        self.config = config or OrchestraLangChainConfig()
        
        self.cache_manager = CacheManager(
            similarity_threshold=self.config.similarity_threshold,
            cache_ttl=self.config.cache_ttl,
            enable_compression=self.config.enable_compression
        )
        self.metrics = MetricsTracker()
        
        logger.info("✨ Orchestra enhancement enabled for LangChain")

    def invoke(self, input: Any, config: Optional[Dict] = None, **kwargs) -> Any:
        """
        LCEL-compatible invoke.
        Also supports legacy .run() via __getattr__ delegation if needed,
        but typically users wrap usage with this invoke.
        """
        return self._execute(input, config, is_async=False, **kwargs)

    async def ainvoke(self, input: Any, config: Optional[Dict] = None, **kwargs) -> Any:
        """LCEL-compatible async invoke."""
        return await self._execute_async(input, config, **kwargs)

    def _execute(self, input: Any, config: Any, is_async: bool, **kwargs) -> Any:
        # 1. Serialize input to string for embedding
        input_str = self._to_string(input)
        
        # 2. Check Cache
        cached_result = self.cache_manager.get(input_str)
        if cached_result is not None:
            self.metrics.record_cache_hit(0.01) # Dummy latency
            logger.info("💰 Cache HIT (LangChain)")
            return cached_result
        
        # 3. Execute
        self.metrics.start_execution()
        start = time.time()
        
        # Support both Runnable .invoke() and Chain .run() / .__call__
        if hasattr(self.chain, "invoke"):
            result = self.chain.invoke(input, config, **kwargs)
        elif hasattr(self.chain, "run") and isinstance(input, str):
            result = self.chain.run(input, **kwargs)
        else:
            result = self.chain(input, **kwargs)
            if isinstance(result, dict) and len(result) == 1:
                    # Unpack single dict result often returned by chains
                    result = list(result.values())[0]

        latency = time.time() - start
        self.metrics.end_execution(latency)
        
        # 4. Store
        self.cache_manager.put(input_str, result)
        
        return result

    async def _execute_async(self, input: Any, config: Any, **kwargs):
        # 1. Serialize
        input_str = self._to_string(input)
        
        # 2. Check Cache
        cached_result = self.cache_manager.get(input_str)
        if cached_result is not None:
            self.metrics.record_cache_hit(0.01)
            logger.info("💰 Cache HIT (LangChain Async)")
            return cached_result
        
        # 3. Execute
        self.metrics.start_execution()
        start = time.time()
        
        if hasattr(self.chain, "ainvoke"):
            result = await self.chain.ainvoke(input, config, **kwargs)
        elif hasattr(self.chain, "arun") and isinstance(input, str):
            result = await self.chain.arun(input, **kwargs)
        elif hasattr(self.chain, "acall"):
            result = await self.chain.acall(input, **kwargs)
            if isinstance(result, dict) and len(result) == 1:
                result = list(result.values())[0]
        else:
            # Fallback to sync invoke if async not supported
            logger.warning("Chain does not support async invoke/run, falling back to sync.")
            result = self.chain.invoke(input, config, **kwargs)

        latency = time.time() - start
        self.metrics.end_execution(latency)
        
        # 4. Store
        self.cache_manager.put(input_str, result)
        
        return result

    def _to_string(self, input: Any) -> str:
        if isinstance(input, str):
            return input
        if isinstance(input, dict):
            # Try to find a single value or dump
            return json.dumps(input, sort_keys=True, default=str)
        return str(input)

    # Proxy other methods to underlying chain
    def __getattr__(self, name):
        return getattr(self.chain, name)

def enhance(chain_or_runnable, config=None):
    return EnhancedLangChain(chain_or_runnable, config)
