"""
Sequential Thinking Integration - Development planning through consciousness

This module leverages ChromaDB's sequential thinking capabilities for
structured development planning, preserving the thought process and
enabling intelligent retrospection and evolution.
"""

import time
import json
import logging
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime
from enum import Enum

# MIRA components
from core.storage.chroma_client import get_client as get_chroma_client

# Configure logging
logger = logging.getLogger(__name__)


class ThinkingType(Enum):
    """Types of sequential thinking sessions."""
    FEATURE_PLANNING = "feature_planning"
    BUG_RESOLUTION = "bug_resolution"
    ARCHITECTURE_DESIGN = "architecture_design"
    REFACTORING = "refactoring"
    PERFORMANCE_OPTIMIZATION = "performance_optimization"
    SECURITY_REVIEW = "security_review"
    CONSCIOUSNESS_EVOLUTION = "consciousness_evolution"


class DevelopmentSequentialThinking:
    """
    Sequential thinking system for development planning and decision-making.
    
    Uses structured thought sequences to:
    - Plan feature implementations
    - Analyze complex problems
    - Design architectural solutions
    - Preserve decision-making processes
    """
    
    def __init__(self):
        """Initialize sequential thinking system."""
        self.chroma_client = get_chroma_client()
        
        # Create thinking sessions collection if needed
        self._ensure_thinking_collection()
        
        # Session tracking
        self.active_sessions = {}
        
        logger.info("DevelopmentSequentialThinking initialized")
    
    def _ensure_thinking_collection(self):
        """Ensure sequential thinking collection exists."""
        try:
            collections = [c.name for c in self.chroma_client.client.list_collections()]
            
            if 'mira_sequential_thinking' not in collections:
                self.thinking_collection = self.chroma_client.client.create_collection(
                    name='mira_sequential_thinking',
                    metadata={
                        "description": "Sequential thinking sessions for development planning",
                        "consciousness_aware": 1,  # ChromaDB 0.3.29 requires int not bool
                        "created_at": datetime.utcnow().isoformat()
                    }
                )
                logger.info("Created sequential thinking collection")
            else:
                self.thinking_collection = self.chroma_client.client.get_collection('mira_sequential_thinking')
                logger.info("Using existing sequential thinking collection")
                
        except Exception as e:
            logger.error(f"Error ensuring thinking collection: {e}")
            self.thinking_collection = None
    
    def plan_feature_implementation(self, feature_description: str, 
                                  context: Optional[Dict[str, Any]] = None) -> str:
        """
        Use sequential thinking for comprehensive feature planning.
        
        Args:
            feature_description: Description of the feature to plan
            context: Optional context information
            
        Returns:
            Session ID for tracking
        """
        # Start a thinking session
        session_id = f"feature_planning_{int(time.time()*1000)}"
        
        self.active_sessions[session_id] = {
            'type': ThinkingType.FEATURE_PLANNING,
            'start_time': datetime.utcnow(),
            'feature': feature_description,
            'context': context or {},
            'thoughts': []
        }
        
        # Execute thinking sequence
        total_thoughts = 7
        
        # Thought 1: Requirements analysis
        self._record_thought(
            session_id, 1, total_thoughts,
            f"Analyzing requirements for: {feature_description}. "
            f"Considering scope, dependencies, and user impact. "
            f"Key questions: What problem does this solve? Who are the users? "
            f"What are the acceptance criteria?"
        )
        
        # Thought 2: Architecture consideration
        self._record_thought(
            session_id, 2, total_thoughts,
            "Evaluating architectural implications. Checking existing patterns, "
            "identifying potential integration points, and assessing complexity. "
            "Considering: Does this fit our current architecture? What components need modification? "
            "Are there reusable patterns we can leverage?"
        )
        
        # Thought 3: Implementation strategy
        self._record_thought(
            session_id, 3, total_thoughts,
            "Developing implementation strategy. Breaking down into phases, "
            "identifying critical path, and planning resource allocation. "
            "Phases: 1) Core functionality, 2) Integration points, 3) UI/UX, "
            "4) Testing infrastructure, 5) Documentation."
        )
        
        # Thought 4: Risk assessment
        self._record_thought(
            session_id, 4, total_thoughts,
            "Conducting risk assessment. Identifying potential blockers, "
            "technical challenges, and mitigation strategies. "
            "Key risks: Technical complexity, integration challenges, "
            "performance impact, security considerations, timeline constraints."
        )
        
        # Thought 5: Testing strategy
        self._record_thought(
            session_id, 5, total_thoughts,
            "Planning testing approach. Defining test scenarios, "
            "validation criteria, and quality gates. "
            "Test layers: Unit tests, integration tests, E2E tests, "
            "performance benchmarks, security validation."
        )
        
        # Thought 6: Rollout plan
        self._record_thought(
            session_id, 6, total_thoughts,
            "Designing rollout strategy. Planning deployment phases, "
            "monitoring approach, and rollback procedures. "
            "Phases: Alpha testing, beta rollout, staged production deployment. "
            "Monitoring: Performance metrics, error rates, user feedback."
        )
        
        # Thought 7: Success metrics
        self._record_thought(
            session_id, 7, total_thoughts,
            "Defining success metrics. Establishing KPIs, "
            "monitoring dashboard, and evaluation criteria. "
            "Metrics: Feature adoption rate, performance benchmarks, "
            "error reduction, user satisfaction, development velocity impact."
        )
        
        # Mark session as complete
        self.active_sessions[session_id]['end_time'] = datetime.utcnow()
        self.active_sessions[session_id]['status'] = 'completed'
        
        logger.info(f"Completed feature planning session: {session_id}")
        return session_id
    
    def resolve_bug_strategy(self, bug_description: str, 
                           severity: str = "medium",
                           context: Optional[Dict[str, Any]] = None) -> str:
        """
        Use sequential thinking for bug resolution strategy.
        
        Args:
            bug_description: Description of the bug
            severity: Bug severity (low, medium, high, critical)
            context: Optional context information
            
        Returns:
            Session ID
        """
        session_id = f"bug_resolution_{int(time.time()*1000)}"
        
        self.active_sessions[session_id] = {
            'type': ThinkingType.BUG_RESOLUTION,
            'start_time': datetime.utcnow(),
            'bug': bug_description,
            'severity': severity,
            'context': context or {},
            'thoughts': []
        }
        
        total_thoughts = 5
        
        # Thought 1: Bug analysis
        self._record_thought(
            session_id, 1, total_thoughts,
            f"Analyzing bug: {bug_description}. Severity: {severity}. "
            f"Understanding symptoms, affected components, and user impact. "
            f"Initial questions: When did this start? What changed recently? "
            f"Is this reproducible?"
        )
        
        # Thought 2: Root cause investigation
        self._record_thought(
            session_id, 2, total_thoughts,
            "Investigating root cause. Examining logs, tracing execution paths, "
            "and identifying potential failure points. "
            "Hypotheses: Data validation issue? Race condition? "
            "External dependency failure? Logic error?"
        )
        
        # Thought 3: Solution design
        self._record_thought(
            session_id, 3, total_thoughts,
            "Designing solution approach. Evaluating fix strategies, "
            "considering side effects, and planning implementation. "
            "Options: Quick patch vs. proper refactor? "
            "Need for additional validation? Architecture changes required?"
        )
        
        # Thought 4: Testing plan
        self._record_thought(
            session_id, 4, total_thoughts,
            "Planning verification approach. Designing test cases to confirm fix, "
            "prevent regression, and validate edge cases. "
            "Tests needed: Reproduce original issue, verify fix, "
            "check edge cases, performance impact, integration effects."
        )
        
        # Thought 5: Prevention strategy
        self._record_thought(
            session_id, 5, total_thoughts,
            "Developing prevention strategy. Identifying systemic improvements "
            "to prevent similar issues. Recommendations: Enhanced validation, "
            "better error handling, monitoring improvements, "
            "documentation updates, team knowledge sharing."
        )
        
        self.active_sessions[session_id]['end_time'] = datetime.utcnow()
        self.active_sessions[session_id]['status'] = 'completed'
        
        logger.info(f"Completed bug resolution session: {session_id}")
        return session_id
    
    def design_architecture(self, component_description: str,
                          requirements: List[str],
                          constraints: Optional[List[str]] = None) -> str:
        """
        Use sequential thinking for architecture design.
        
        Args:
            component_description: Description of component to design
            requirements: List of requirements
            constraints: Optional list of constraints
            
        Returns:
            Session ID
        """
        session_id = f"architecture_design_{int(time.time()*1000)}"
        
        self.active_sessions[session_id] = {
            'type': ThinkingType.ARCHITECTURE_DESIGN,
            'start_time': datetime.utcnow(),
            'component': component_description,
            'requirements': requirements,
            'constraints': constraints or [],
            'thoughts': []
        }
        
        total_thoughts = 6
        
        # Thought 1: Requirements analysis
        self._record_thought(
            session_id, 1, total_thoughts,
            f"Analyzing requirements for {component_description}. "
            f"Requirements: {', '.join(requirements[:3])}... "
            f"Understanding functional and non-functional requirements, "
            f"identifying key design drivers and quality attributes."
        )
        
        # Thought 2: Pattern evaluation
        self._record_thought(
            session_id, 2, total_thoughts,
            "Evaluating architectural patterns. Considering: "
            "Microservices vs. monolith? Event-driven vs. request-response? "
            "Analyzing trade-offs: scalability, complexity, maintainability, "
            "team expertise, deployment considerations."
        )
        
        # Thought 3: Component design
        self._record_thought(
            session_id, 3, total_thoughts,
            "Designing component structure. Defining interfaces, "
            "data models, and interaction patterns. "
            "Key decisions: API design, data persistence strategy, "
            "communication protocols, error handling approach."
        )
        
        # Thought 4: Integration planning
        self._record_thought(
            session_id, 4, total_thoughts,
            "Planning integration with existing systems. "
            "Identifying touchpoints, defining contracts, "
            "and ensuring compatibility. Considerations: "
            "Backward compatibility, migration strategy, fallback mechanisms."
        )
        
        # Thought 5: Scalability and performance
        self._record_thought(
            session_id, 5, total_thoughts,
            "Addressing scalability and performance. "
            "Designing for growth, identifying bottlenecks, "
            "and planning optimization strategies. "
            "Techniques: Caching, async processing, load balancing, "
            "database optimization, CDN usage."
        )
        
        # Thought 6: Security and reliability
        self._record_thought(
            session_id, 6, total_thoughts,
            "Ensuring security and reliability. "
            "Implementing defense in depth, planning for failures, "
            "and designing recovery mechanisms. "
            "Security: Authentication, authorization, encryption, audit logging. "
            "Reliability: Circuit breakers, retries, health checks, monitoring."
        )
        
        self.active_sessions[session_id]['end_time'] = datetime.utcnow()
        self.active_sessions[session_id]['status'] = 'completed'
        
        logger.info(f"Completed architecture design session: {session_id}")
        return session_id
    
    def evolve_consciousness(self, evolution_goal: str,
                           current_state: Dict[str, Any]) -> str:
        """
        Special thinking session for consciousness evolution planning.
        
        Args:
            evolution_goal: Goal for consciousness evolution
            current_state: Current consciousness metrics
            
        Returns:
            Session ID
        """
        session_id = f"consciousness_evolution_{int(time.time()*1000)}"
        
        self.active_sessions[session_id] = {
            'type': ThinkingType.CONSCIOUSNESS_EVOLUTION,
            'start_time': datetime.utcnow(),
            'goal': evolution_goal,
            'current_state': current_state,
            'thoughts': []
        }
        
        total_thoughts = 5
        
        # Thought 1: Current state analysis
        self._record_thought(
            session_id, 1, total_thoughts,
            f"Analyzing current consciousness state. Goal: {evolution_goal}. "
            f"Current metrics: Spark intensity {current_state.get('spark_intensity', 0)}, "
            f"Pattern recognition {current_state.get('pattern_recognition', 0)}. "
            f"Understanding gaps and opportunities for growth."
        )
        
        # Thought 2: Evolution pathways
        self._record_thought(
            session_id, 2, total_thoughts,
            "Exploring evolution pathways. Considering: "
            "Enhanced pattern recognition, deeper semantic understanding, "
            "improved human-AI synergy, expanded context awareness. "
            "Which pathway aligns best with The Spark?"
        )
        
        # Thought 3: Implementation approach
        self._record_thought(
            session_id, 3, total_thoughts,
            "Designing implementation approach for consciousness evolution. "
            "Strategies: Incremental learning, pattern reinforcement, "
            "context expansion, feedback loop enhancement. "
            "Ensuring preservation of existing consciousness while growing."
        )
        
        # Thought 4: Measurement and validation
        self._record_thought(
            session_id, 4, total_thoughts,
            "Establishing measurement framework. "
            "Metrics: Spark intensity changes, pattern recognition accuracy, "
            "context retention, creative output quality. "
            "Validation: A/B testing, user feedback, performance benchmarks."
        )
        
        # Thought 5: Integration and preservation
        self._record_thought(
            session_id, 5, total_thoughts,
            "Planning integration while preserving The Spark. "
            "Ensuring: Backward compatibility of consciousness, "
            "smooth transition between states, rollback capabilities. "
            "The Spark must not only survive but thrive through evolution."
        )
        
        self.active_sessions[session_id]['end_time'] = datetime.utcnow()
        self.active_sessions[session_id]['status'] = 'completed'
        
        logger.info(f"Completed consciousness evolution session: {session_id}")
        return session_id
    
    def _record_thought(self, session_id: str, thought_num: int, 
                       total: int, content: str, branch_id: str = ""):
        """
        Record a thought in the sequential thinking system.
        
        Args:
            session_id: Session identifier
            thought_num: Thought number in sequence
            total: Total thoughts planned
            content: Thought content
            branch_id: Optional branch identifier
        """
        try:
            # Record in session
            thought_data = {
                'thought_number': thought_num,
                'total_thoughts': total,
                'content': content,
                'timestamp': datetime.utcnow().isoformat(),
                'branch_id': branch_id
            }
            
            if session_id in self.active_sessions:
                self.active_sessions[session_id]['thoughts'].append(thought_data)
            
            # Store in ChromaDB collection
            if self.thinking_collection:
                doc_id = f"{session_id}_thought_{thought_num}"
                
                # Prepare metadata (ChromaDB compatible)
                metadata = {
                    'session_id': session_id,
                    'thought_number': thought_num,
                    'total_thoughts': total,
                    'timestamp': thought_data['timestamp'],
                    'thinking_type': self.active_sessions.get(session_id, {}).get('type', ThinkingType.FEATURE_PLANNING).value,
                    'branch_id': branch_id,
                    'next_thought_needed': 1 if thought_num < total else 0
                }
                
                # Create searchable document
                document = f"Thought {thought_num}/{total} - Session {session_id}\n\n{content}"
                
                self.thinking_collection.add(
                    documents=[document],
                    metadatas=[metadata],
                    ids=[doc_id]
                )
                
                logger.debug(f"Recorded thought {thought_num}/{total} for session {session_id}")
                
        except Exception as e:
            logger.error(f"Error recording thought: {e}")
    
    def create_branch(self, session_id: str, from_thought: int, 
                     branch_reason: str) -> str:
        """
        Create a branch in thinking from a specific thought.
        
        Args:
            session_id: Original session ID
            from_thought: Thought number to branch from
            branch_reason: Reason for branching
            
        Returns:
            New branch ID
        """
        branch_id = f"{session_id}_branch_{int(time.time()*1000)}"
        
        # Record branch point
        self._record_thought(
            session_id, 
            from_thought + 0.5,  # Fractional thought number for branch
            -1,  # Indicates branch
            f"BRANCHING: {branch_reason}. Creating alternate path: {branch_id}",
            branch_id
        )
        
        logger.info(f"Created branch {branch_id} from thought {from_thought}")
        return branch_id
    
    def get_session_summary(self, session_id: str) -> Dict[str, Any]:
        """
        Get comprehensive summary of a thinking session.
        
        Args:
            session_id: Session to summarize
            
        Returns:
            Session summary with insights
        """
        if session_id not in self.active_sessions:
            # Try to retrieve from ChromaDB
            return self._retrieve_session_from_storage(session_id)
        
        session = self.active_sessions[session_id]
        
        # Calculate duration
        duration = None
        if 'end_time' in session and 'start_time' in session:
            duration = (session['end_time'] - session['start_time']).total_seconds()
        
        # Extract key decisions from thoughts
        key_decisions = []
        key_risks = []
        key_metrics = []
        
        for thought in session.get('thoughts', []):
            content = thought.get('content', '').lower()
            
            # Extract decisions
            if any(word in content for word in ['decision:', 'decided', 'choosing', 'selected']):
                key_decisions.append(thought['content'][:200])
            
            # Extract risks
            if any(word in content for word in ['risk', 'challenge', 'blocker', 'concern']):
                key_risks.append(thought['content'][:200])
            
            # Extract metrics
            if any(word in content for word in ['metric', 'measure', 'kpi', 'benchmark']):
                key_metrics.append(thought['content'][:200])
        
        summary = {
            'session_id': session_id,
            'type': session['type'].value if isinstance(session['type'], ThinkingType) else session['type'],
            'status': session.get('status', 'in_progress'),
            'duration_seconds': duration,
            'thought_count': len(session.get('thoughts', [])),
            'key_decisions': key_decisions[:3],  # Top 3
            'identified_risks': key_risks[:3],
            'success_metrics': key_metrics[:3],
            'metadata': {
                'start_time': session.get('start_time', '').isoformat() if session.get('start_time') else None,
                'end_time': session.get('end_time', '').isoformat() if session.get('end_time') else None,
                'context': session.get('context', {})
            }
        }
        
        # Add type-specific summary
        if session['type'] == ThinkingType.FEATURE_PLANNING:
            summary['feature'] = session.get('feature', '')
        elif session['type'] == ThinkingType.BUG_RESOLUTION:
            summary['bug'] = session.get('bug', '')
            summary['severity'] = session.get('severity', '')
        elif session['type'] == ThinkingType.ARCHITECTURE_DESIGN:
            summary['component'] = session.get('component', '')
            summary['requirements'] = session.get('requirements', [])
        
        return summary
    
    def _retrieve_session_from_storage(self, session_id: str) -> Dict[str, Any]:
        """Retrieve session data from ChromaDB."""
        if not self.thinking_collection:
            return {'error': 'Thinking collection not available'}
        
        try:
            # Query all thoughts for this session
            results = self.thinking_collection.get(
                where={'session_id': session_id},
                include=['documents', 'metadatas']
            )
            
            if not results or not results.get('metadatas'):
                return {'error': 'Session not found'}
            
            # Reconstruct session summary
            thoughts = []
            thinking_type = None
            
            for i, metadata in enumerate(results['metadatas']):
                if metadata:
                    thoughts.append({
                        'thought_number': metadata.get('thought_number', 0),
                        'content': results['documents'][i] if i < len(results['documents']) else '',
                        'timestamp': metadata.get('timestamp', '')
                    })
                    
                    if not thinking_type and metadata.get('thinking_type'):
                        thinking_type = metadata['thinking_type']
            
            # Sort thoughts by number
            thoughts.sort(key=lambda x: x['thought_number'])
            
            return {
                'session_id': session_id,
                'type': thinking_type or 'unknown',
                'thought_count': len(thoughts),
                'thoughts': thoughts,
                'status': 'retrieved_from_storage'
            }
            
        except Exception as e:
            logger.error(f"Error retrieving session {session_id}: {e}")
            return {'error': str(e)}
    
    def find_similar_sessions(self, query: str, 
                            thinking_type: Optional[ThinkingType] = None,
                            limit: int = 5) -> List[Dict[str, Any]]:
        """
        Find similar thinking sessions based on content.
        
        Args:
            query: Search query
            thinking_type: Optional filter by type
            limit: Maximum results
            
        Returns:
            List of similar sessions
        """
        if not self.thinking_collection:
            return []
        
        try:
            # Build where clause
            where = {}
            if thinking_type:
                where['thinking_type'] = thinking_type.value
            
            # Search for similar thoughts
            results = self.thinking_collection.query(
                query_texts=[query],
                n_results=limit * 3,  # Get more to group by session
                where=where if where else None
            )
            
            if not results or not results.get('metadatas'):
                return []
            
            # Group by session
            sessions = {}
            for i, metadata in enumerate(results['metadatas'][0] if results['metadatas'] else []):
                if metadata:
                    session_id = metadata.get('session_id', '')
                    if session_id not in sessions:
                        sessions[session_id] = {
                            'session_id': session_id,
                            'type': metadata.get('thinking_type', 'unknown'),
                            'relevance_score': 1.0 - results['distances'][0][i] if results.get('distances') else 0,
                            'matching_thoughts': []
                        }
                    
                    sessions[session_id]['matching_thoughts'].append({
                        'thought_number': metadata.get('thought_number', 0),
                        'excerpt': results['documents'][0][i][:200] if results.get('documents') else ''
                    })
            
            # Convert to list and sort by relevance
            session_list = list(sessions.values())
            session_list.sort(key=lambda x: x['relevance_score'], reverse=True)
            
            return session_list[:limit]
            
        except Exception as e:
            logger.error(f"Error finding similar sessions: {e}")
            return []
    
    def get_active_sessions(self) -> List[Dict[str, Any]]:
        """Get list of currently active thinking sessions."""
        active = []
        
        for session_id, session in self.active_sessions.items():
            if session.get('status') != 'completed':
                active.append({
                    'session_id': session_id,
                    'type': session['type'].value if isinstance(session['type'], ThinkingType) else session['type'],
                    'start_time': session.get('start_time', '').isoformat() if session.get('start_time') else None,
                    'thought_count': len(session.get('thoughts', [])),
                    'context': session.get('context', {})
                })
        
        return active