#!/usr/bin/env python3
"""
DcisionAI MCP Tools - Production Version 2.0
============================================
SECURITY: No eval(), uses AST parsing
VALIDATION: Comprehensive result validation  
RELIABILITY: Multi-region failover, rate limiting
"""

import asyncio
import json
import logging
import re
import os
import hashlib
import time
from datetime import datetime
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, asdict
from functools import lru_cache

# Import MathOpt model builder
try:
    from .mathopt_model_builder import MathOptModelBuilder, HAS_MATHOPT
except ImportError:
    HAS_MATHOPT = False
    logging.warning("MathOpt model builder not available")
from collections import deque
import ast
import operator

try:
    import numpy as np
    HAS_MONTE_CARLO = True
except ImportError:
    HAS_MONTE_CARLO = False

import boto3
from botocore.exceptions import ClientError

from .workflows import WorkflowManager
from .config import Config
from .optimization_engine import solve_real_optimization
from .solver_selector import SolverSelector

logger = logging.getLogger(__name__)

# ============================================================================
# DATA STRUCTURES
# ============================================================================

@dataclass
class Variable:
    name: str
    type: str
    bounds: str
    description: str
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

@dataclass
class Constraint:
    expression: str
    description: str
    type: str = "inequality"
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

@dataclass
class Objective:
    type: str
    expression: str
    description: str
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

@dataclass
class ModelSpec:
    variables: List[Variable]
    constraints: List[Constraint]
    objective: Objective
    model_type: str
    model_complexity: str = "medium"
    estimated_solve_time: float = 1.0
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ModelSpec':
        if 'result' in data:
            data = data['result']
        if 'raw_response' in data:
            data = json.loads(data['raw_response'])
        
        variables = [Variable(**v) if isinstance(v, dict) else v for v in data.get('variables', [])]
        constraints = [Constraint(**c) if isinstance(c, dict) else c for c in data.get('constraints', [])]
        obj = data.get('objective', {})
        objective = Objective(**obj) if isinstance(obj, dict) else obj
        
        return cls(
            variables=variables,
            constraints=constraints,
            objective=objective,
            model_type=data.get('model_type', 'linear_programming'),
            model_complexity=data.get('model_complexity', 'medium'),
            estimated_solve_time=data.get('estimated_solve_time', 1.0)
        )
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'variables': [v.to_dict() for v in self.variables],
            'constraints': [c.to_dict() for c in self.constraints],
            'objective': self.objective.to_dict(),
            'model_type': self.model_type,
            'model_complexity': self.model_complexity,
            'estimated_solve_time': self.estimated_solve_time
        }

# ============================================================================
# KNOWLEDGE BASE
# ============================================================================

class KnowledgeBase:
    def __init__(self, path: str):
        self.path = path
        self.kb = self._load()
    
    def _load(self) -> Dict[str, Any]:
        try:
            if os.path.exists(self.path):
                with open(self.path) as f:
                    return json.load(f)
        except Exception as e:
            logger.warning(f"KB load failed: {e}")
        return {'examples': []}
    
    @lru_cache(maxsize=500)
    def search(self, query: str, top_k: int = 2) -> str:
        query_lower = query.lower()
        results = []
        
        for ex in self.kb.get('examples', []):
            score = sum(2 for w in query_lower.split() if w in ex.get('problem_description', '').lower())
            score += sum(3 for kw in ex.get('keywords', []) if kw.lower() in query_lower)
            if score > 0:
                results.append((score, ex))
        
        results.sort(key=lambda x: x[0], reverse=True)
        if not results[:top_k]:
            return "No similar examples."
        
        context = "Similar:\n"
        for _, ex in results[:top_k]:
            context += f"- {ex.get('problem_type', '')}: {ex.get('solution', '')[:80]}...\n"
        return context[:300]
    
    def get_problem_type_guidance(self, problem_description: str) -> str:
        """Get specific guidance based on problem type."""
        query_lower = problem_description.lower()
        
        if 'portfolio' in query_lower or 'investment' in query_lower or 'asset' in query_lower:
            return """
**Portfolio Optimization Guidance:**
- Decision variables represent investment allocations or asset weights
- Constraints include budget limits, risk limits, and diversification requirements
- Objective balances return maximization with risk minimization
- Consider correlation matrices, expected returns, and risk measures
- If individual stocks are mentioned, create individual stock variables
- If sector constraints exist, create sector-level constraints
"""
        elif 'production' in query_lower or 'factory' in query_lower or 'manufacturing' in query_lower:
            return """
**Production Planning Guidance:**
- Decision variables typically represent production quantities or resource allocation
- Constraints often include capacity limits, demand requirements, and resource availability
- Objective is usually cost minimization or profit maximization
- Consider time periods, production lines, and inventory constraints
"""
        elif 'schedule' in query_lower or 'task' in query_lower or 'employee' in query_lower:
            return """
**Scheduling Optimization Guidance:**
- Decision variables represent task assignments, start times, or resource allocations
- Constraints include precedence relationships, resource capacity, and deadlines
- Objective is usually makespan minimization or cost optimization
- Consider task dependencies, resource constraints, and time windows
"""
        else:
            return """
**Generic Optimization Guidance:**
- Identify the key decisions to be made (decision variables)
- Determine the limitations and requirements (constraints)
- Define the optimization goal (objective function)
- Ensure all variables are used and constraints are mathematically sound
"""

# ============================================================================
# RATE LIMITING
# ============================================================================

class RateLimiter:
    def __init__(self, max_calls: int, period: float):
        self.max_calls = max_calls
        self.period = period
        self.calls = deque()
        self.lock = asyncio.Lock()
    
    async def acquire(self):
        async with self.lock:
            now = time.time()
            while self.calls and self.calls[0] < now - self.period:
                self.calls.popleft()
            while len(self.calls) >= self.max_calls:
                await asyncio.sleep(0.1)
                now = time.time()
                while self.calls and self.calls[0] < now - self.period:
                    self.calls.popleft()
            self.calls.append(now)

# ============================================================================
# SAFE EXPRESSION EVALUATOR (NO eval()!)
# ============================================================================

class SafeEvaluator:
    OPS = {
        ast.Add: operator.add,
        ast.Sub: operator.sub,
        ast.Mult: operator.mul,
        ast.Div: operator.truediv,
        ast.Pow: operator.pow,
        ast.USub: operator.neg,
    }
    
    @classmethod
    def evaluate(cls, expr: str, vars: Dict[str, float]) -> float:
        for name, val in sorted(vars.items(), key=lambda x: -len(x[0])):
            expr = re.sub(r'\b' + re.escape(name) + r'\b', str(val), expr)
        try:
            node = ast.parse(expr, mode='eval')
            return cls._eval(node.body)
        except Exception as e:
            raise ValueError(f"Expression evaluation failed: {e}")
    
    @classmethod
    def _eval(cls, node):
        if isinstance(node, (ast.Constant, ast.Num)):
            return node.value if hasattr(node, 'value') else node.n
        elif isinstance(node, ast.BinOp):
            return cls.OPS[type(node.op)](cls._eval(node.left), cls._eval(node.right))
        elif isinstance(node, ast.UnaryOp):
            return cls.OPS[type(node.op)](cls._eval(node.operand))
        raise ValueError(f"Unsupported: {type(node)}")

# ============================================================================
# VALIDATION ENGINE
# ============================================================================

class Validator:
    def __init__(self):
        self.eval = SafeEvaluator()
    
    def validate(self, result: Dict, model: ModelSpec) -> Dict[str, Any]:
        errors = []
        warnings = []
        
        status = result.get('status')
        values = result.get('optimal_values', {})
        obj_val = result.get('objective_value', 0)
        
        if status == 'optimal' and values:
            try:
                calc = self.eval.evaluate(model.objective.expression, values)
                err = abs(calc - obj_val) / max(abs(calc), 1e-10)
                if err > 0.001:
                    errors.append(f"Objective mismatch: calc={calc:.4f}, reported={obj_val:.4f}")
            except Exception as e:
                warnings.append(f"Could not validate objective: {e}")
        
        if status == 'optimal' and values:
            for c in model.constraints:
                try:
                    if not self._check_constraint(c.expression, values):
                        errors.append(f"Violated: {c.expression}")
                except Exception as e:
                    warnings.append(f"Could not check {c.expression}: {e}")
        
        return {
            'is_valid': len(errors) == 0,
            'errors': errors,
            'warnings': warnings
        }
    
    def _check_constraint(self, expr: str, vars: Dict[str, float]) -> bool:
        if '<=' in expr:
            left, right = expr.split('<=', 1)
            return self.eval.evaluate(left, vars) <= self.eval.evaluate(right, vars) + 1e-6
        elif '>=' in expr:
            left, right = expr.split('>=', 1)
            return self.eval.evaluate(left, vars) >= self.eval.evaluate(right, vars) - 1e-6
        elif '==' in expr or '=' in expr:
            left, right = expr.split('==' if '==' in expr else '=', 1)
            return abs(self.eval.evaluate(left, vars) - self.eval.evaluate(right, vars)) < 1e-6
        return True

# ============================================================================
# BEDROCK CLIENT WITH FAILOVER
# ============================================================================

class BedrockClient:
    def __init__(self, regions: List[str] = None):
        self.regions = regions or ['us-east-1', 'us-west-2']
        self.current = 0
        self.clients = {}
        for r in self.regions:
            try:
                self.clients[r] = boto3.client('bedrock-runtime', region_name=r)
            except Exception as e:
                logger.error(f"Failed to init {r}: {e}")
        self.limiters = {
            'haiku': RateLimiter(10, 1.0),
            'sonnet': RateLimiter(5, 1.0)
        }
    
    async def invoke(self, model_id: str, prompt: str, max_tokens: int = 4000) -> str:
        limiter = self.limiters['haiku' if 'haiku' in model_id.lower() else 'sonnet']
        await limiter.acquire()
        
        for attempt in range(len(self.regions)):
            region = self.regions[self.current]
            client = self.clients.get(region)
            if not client:
                self.current = (self.current + 1) % len(self.regions)
                continue
            
            try:
                body = json.dumps({
                    "anthropic_version": "bedrock-2023-05-31",
                    "max_tokens": max_tokens,
                    "temperature": 0.1,
                    "messages": [{"role": "user", "content": prompt}]
                })
                
                resp = await asyncio.to_thread(
                    client.invoke_model,
                    modelId=model_id,
                    body=body,
                    contentType="application/json"
                )
                
                data = json.loads(resp['body'].read())
                if 'content' in data:
                    return data['content'][0]['text']
                elif 'completion' in data:
                    return data['completion']
                raise ValueError("Unexpected response")
                
            except Exception as e:
                if "ServiceUnavailable" in str(e) or "Throttling" in str(e):
                    self.current = (self.current + 1) % len(self.regions)
                    if attempt < len(self.regions) - 1:
                        continue
                raise
        
        raise RuntimeError("All regions failed")

# ============================================================================
# MAIN TOOLS CLASS
# ============================================================================

class DcisionAITools:
    def __init__(self, config: Optional[Config] = None):
        self.config = config or Config()
        self.bedrock = BedrockClient()
        self.validator = Validator()
        self.solver_selector = SolverSelector()
        self.workflow_manager = WorkflowManager()
        
        kb_path = os.path.join(os.path.dirname(__file__), '..', 'dcisionai_kb.json')
        self.kb = KnowledgeBase(kb_path)
        self.cache = {}
        
        logger.info("DcisionAI Tools v2.0 initialized")
    
    def _parse_json(self, text: str) -> Dict[str, Any]:
        """Enhanced JSON parsing with comprehensive error handling and recovery"""
        if not text:
            return {"raw_response": ""}
        
        # Enhanced text cleaning
        text = re.sub(r'```json\s*', '', text)
        text = re.sub(r'```\s*', '', text)
        text = re.sub(r',(\s*[}\]])', r'\1', text)  # Remove trailing commas
        text = re.sub(r'([{,]\s*)([a-zA-Z_][a-zA-Z0-9_]*)\s*:', r'\1"\2":', text)  # Quote unquoted keys
        text = re.sub(r':\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*([,}])', r': "\1"\2', text)  # Quote unquoted string values
        text = text.strip()
        
        # Try direct parsing first
        try:
            result = json.loads(text)
            if isinstance(result, dict):
                return result
        except json.JSONDecodeError as e:
            logger.debug(f"Direct JSON parsing failed: {e}")
        
        # Try to find JSON object in the text - look for the first complete JSON object
        try:
            # Find the first opening brace
            start_pos = text.find('{')
            if start_pos == -1:
                raise ValueError("No opening brace found")
            
            # Count braces to find the matching closing brace
            brace_count = 0
            in_string = False
            escape_next = False
            
            for i in range(start_pos, len(text)):
                char = text[i]
                
                if escape_next:
                    escape_next = False
                    continue
                
                if char == '\\':
                    escape_next = True
                    continue
                
                if char == '"' and not escape_next:
                    in_string = not in_string
                    continue
                
                if not in_string:
                    if char == '{':
                        brace_count += 1
                    elif char == '}':
                        brace_count -= 1
                        if brace_count == 0:
                            # Found matching closing brace
                            json_text = text[start_pos:i+1]
                            
                            # Enhanced cleaning
                            json_text = re.sub(r'[\x00-\x1f\x7f-\x9f]', ' ', json_text)  # Control characters
                            json_text = re.sub(r'\\n', ' ', json_text)  # Newlines in strings
                            json_text = re.sub(r'\\t', ' ', json_text)  # Tabs in strings
                            json_text = re.sub(r'\\r', ' ', json_text)  # Carriage returns
                            
                            result = json.loads(json_text)
                            if isinstance(result, dict):
                                return result
                            break
            
            # If we get here, no matching closing brace was found
            raise ValueError("No matching closing brace found")
            
        except (json.JSONDecodeError, ValueError) as e:
            logger.debug(f"Brace counting JSON parsing failed: {e}")
        
        # Try regex fallback with improved patterns
        try:
            # Try multiple regex patterns for different JSON structures
            patterns = [
                r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}',  # Original pattern
                r'\{[^{}]*"[^"]*"[^{}]*\}',  # Pattern for quoted strings
                r'\{[^{}]*:[^{}]*\}',  # Pattern for key-value pairs
            ]
            
            for pattern in patterns:
                matches = re.finditer(pattern, text, re.DOTALL)
                for match in matches:
                    try:
                        json_text = match.group(0)
                        # Enhanced cleaning
                        json_text = re.sub(r'[\x00-\x1f\x7f-\x9f]', ' ', json_text)
                        json_text = re.sub(r'\\n', ' ', json_text)
                        json_text = re.sub(r'\\t', ' ', json_text)
                        json_text = re.sub(r'\\r', ' ', json_text)
                        
                        result = json.loads(json_text)
                        if isinstance(result, dict):
                            return result
                    except json.JSONDecodeError:
                        continue
        except Exception as e:
            logger.debug(f"Regex JSON parsing failed: {e}")
        
        # Try to extract partial JSON and fix common issues
        try:
            # Find the largest potential JSON object
            start_pos = text.find('{')
            if start_pos != -1:
                # Try to find a reasonable end point
                end_pos = text.rfind('}')
                if end_pos > start_pos:
                    json_text = text[start_pos:end_pos+1]
                    
                    # Try to fix common JSON issues
                    json_text = re.sub(r'([{,]\s*)([a-zA-Z_][a-zA-Z0-9_]*)\s*:', r'\1"\2":', json_text)
                    json_text = re.sub(r':\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*([,}])', r': "\1"\2', json_text)
                    json_text = re.sub(r'[\x00-\x1f\x7f-\x9f]', ' ', json_text)
                    
                    result = json.loads(json_text)
                    if isinstance(result, dict):
                        return result
        except Exception as e:
            logger.debug(f"Partial JSON parsing failed: {e}")
        
        # If all else fails, return raw response with debug info
        logger.warning(f"JSON parsing failed for text: {text[:200]}...")
        return {"raw_response": text}
    
    async def classify_intent(self, problem_description: str, context: Optional[str] = None) -> Dict[str, Any]:
        try:
            cache_key = hashlib.md5(f"intent:{problem_description}".encode()).hexdigest()
            if cache_key in self.cache:
                return self.cache[cache_key]
            
            kb_ctx = self.kb.search(problem_description)
            
            prompt = f"""Classify this optimization problem.

PROBLEM: {problem_description}
SIMILAR: {kb_ctx}

JSON only:
{{
  "intent": "resource_allocation|production_planning|portfolio_optimization|scheduling",
  "industry": "manufacturing|finance|healthcare|retail|logistics|general",
  "optimization_type": "linear_programming|quadratic_programming|mixed_integer_linear_programming",
  "complexity": "low|medium|high",
  "confidence": 0.85
}}"""
            
            resp = await self.bedrock.invoke("anthropic.claude-3-haiku-20240307-v1:0", prompt, 1000)
            result = self._parse_json(resp)
            result.setdefault('intent', 'unknown')
            result.setdefault('optimization_type', 'linear_programming')
            result.setdefault('confidence', 0.7)
            
            response = {
                "status": "success",
                "step": "intent_classification",
                "timestamp": datetime.now().isoformat(),
                "result": result,
                "message": f"Intent: {result['intent']}"
            }
            self.cache[cache_key] = response
            return response
            
        except Exception as e:
            logger.error(f"Intent error: {e}")
            return {"status": "error", "step": "intent_classification", "error": str(e)}
    
    async def analyze_data(self, problem_description: str, intent_data: Optional[Dict] = None) -> Dict[str, Any]:
        try:
            intent = intent_data.get('intent', 'unknown') if intent_data else 'unknown'
            kb_ctx = self.kb.search(problem_description)
            
            prompt = f"""Analyze data for optimization.

PROBLEM: {problem_description}
INTENT: {intent}
SIMILAR: {kb_ctx}

JSON only:
{{
  "readiness_score": 0.85,
  "entities": 10,
  "data_quality": "high|medium|low",
  "variables_identified": ["x1", "x2"],
  "constraints_identified": ["capacity", "demand"]
}}"""
            
            resp = await self.bedrock.invoke("anthropic.claude-3-haiku-20240307-v1:0", prompt, 2000)
            result = self._parse_json(resp)
            result.setdefault('readiness_score', 0.8)
            result.setdefault('entities', 0)
            
            return {
                "status": "success",
                "step": "data_analysis",
                "timestamp": datetime.now().isoformat(),
                "result": result,
                "message": f"Ready: {result['readiness_score']:.1%}"
            }
        except Exception as e:
            logger.error(f"Data error: {e}")
            return {"status": "error", "step": "data_analysis", "error": str(e)}
    
    async def build_model(
        self,
        problem_description: str,
        intent_data: Optional[Dict] = None,
        data_analysis: Optional[Dict] = None,
        solver_selection: Optional[Dict] = None,
        max_retries: int = 2
    ) -> Dict[str, Any]:
        try:
            opt_type = intent_data.get('optimization_type', 'linear_programming') if intent_data else 'linear_programming'
            industry = intent_data.get('industry', 'unknown') if intent_data else 'unknown'
            selected_solver = solver_selection.get('result', {}).get('selected_solver', 'GLOP') if solver_selection else 'GLOP'
            solver_capabilities = ['linear', 'quadratic', 'integer']  # Default capabilities
            
            kb_ctx = self.kb.search(problem_description)
            kb_guidance = self.kb.get_problem_type_guidance(problem_description)
            
            for attempt in range(max_retries + 1):
                retry_note = f"\nPREVIOUS FAILED. Ensure all variables used.\n" if attempt > 0 else ""
                
                prompt = f"""{retry_note}You are a PhD-level optimization expert. Build a mathematical optimization model.

# CRITICAL RULES FOR MODEL BUILDING

## RULE 1: PROBLEM-SPECIFIC FORMULATION
- Read the problem description CAREFULLY
- Identify the SPECIFIC decisions to be made
- Formulate based on THESE specifics, not on general patterns

## RULE 2: VARIABLE DESIGN PRINCIPLES
- Define variables that represent the ACTUAL decisions
- For portfolio problems: If individual stocks are mentioned, create individual stock variables
- For production problems: If individual products are mentioned, create individual product variables
- NEVER oversimplify by grouping when individual items have different constraints

         ## RULE 2A: VARIABLE EXPANSION FOR COMPLEX PROBLEMS
         - **Multi-dimensional problems**: If problem has multiple dimensions (e.g., sites × seasons × archaeologists), create variables for EACH combination
         - **Time-based problems**: If problem spans multiple time periods, create variables for EACH time period
         - **Resource allocation**: If problem involves multiple resources and multiple tasks, create variables for EACH resource-task combination
         - **Scheduling problems**: If problem involves multiple entities (nurses, shifts, days), create variables for EACH entity-shift-day combination
         - **Routing problems**: If problem involves multiple vehicles and multiple locations, create variables for EACH vehicle-location combination
         - **Matrix problems**: If problem involves matrices (e.g., 5 vehicles × 20 customers), create variables for EACH matrix element
         - **Example**: For "10 nurses × 7 days × 3 shifts", create 210 variables (x_nurse_day_shift), not 1 generic variable
         
         ## CRITICAL: NO MATHEMATICAL NOTATION IN VARIABLES
         - **NEVER use Σ (summation) or mathematical notation in variable names**
         - **NEVER use generic variables like x_n_d_s for multi-dimensional problems**
         - **ALWAYS create individual variables for each combination**
         - **Example**: For 3 nurses × 2 days × 2 shifts, create 12 variables:
           - x_nurse1_day1_shift1, x_nurse1_day1_shift2, x_nurse1_day2_shift1, x_nurse1_day2_shift2
           - x_nurse2_day1_shift1, x_nurse2_day1_shift2, x_nurse2_day2_shift1, x_nurse2_day2_shift2
           - x_nurse3_day1_shift1, x_nurse3_day1_shift2, x_nurse3_day2_shift1, x_nurse3_day2_shift2

## RULE 3: CONSTRAINT CAPTURE
- Capture ALL constraints mentioned in the problem
- If problem says "max 10% per stock", create individual stock variables
- If problem says "max 30% per sector", create sector-level constraints
- Ensure constraints can be mathematically enforced

## RULE 4: VALIDATION CHECK
Before finalizing your model, ask:
- Are ALL variables actually decision variables in this problem?
- Do ALL constraints reflect the actual limitations described?
- Does the objective match the actual goal stated?
- Can the model enforce ALL stated constraints?

# KNOWLEDGE BASE CONTEXT
{kb_ctx}

# PROBLEM-TYPE GUIDANCE
{kb_guidance}

PROBLEM DESCRIPTION:
{problem_description}

CONTEXT:
- Intent: {intent_data.get('intent', 'unknown') if intent_data else 'unknown'}
- Industry: {industry}
- Optimization Type: {opt_type}
- Selected Solver: {selected_solver}
- Solver Capabilities: {', '.join(solver_capabilities)}

REQUIRED REASONING PROCESS:
You MUST show your work for each step. Do not skip any reasoning.

Step 1 - Decision Analysis:
What are the key decisions to be made in this problem? List each decision clearly.

Step 2 - Constraint Analysis:
What are the limitations and requirements? List each constraint clearly.

Step 3 - Objective Analysis:
What should be optimized? What is the goal?

         Step 4 - Variable Design:
         How do the decisions translate to mathematical variables? Define each variable with its meaning, type, and bounds.
         **CRITICAL**: For multi-dimensional problems, create variables for EACH combination. Do not use generic variables.
         **Example**: For "10 nurses × 7 days × 3 shifts", create 210 specific variables like x_nurse1_day1_shift1, x_nurse1_day1_shift2, etc.
         
         **VARIABLE EXPANSION REQUIREMENTS**:
         - Count the total number of combinations needed
         - Create a separate variable for each combination
         - Use descriptive names that include all dimensions
         - List ALL variables explicitly in the variables array
         - Do NOT use mathematical notation (Σ, etc.) in variable names
         - Do NOT create generic variables like x_n_d_s

Step 5 - Constraint Formulation:
How do the limitations translate to mathematical constraints? Write each constraint as a mathematical expression.

Step 6 - Objective Formulation:
How does the goal translate to an objective function? Write the mathematical expression.

Step 7 - Validation:
Verify that every variable is used in at least one constraint or the objective function.

# PROBLEM-SPECIFIC VARIABLE EXPANSION GUIDANCE

## PORTFOLIO OPTIMIZATION
If this is a portfolio problem:
- If individual stocks are mentioned (e.g., "AAPL, MSFT, GOOGL"), create individual stock variables
- If sector constraints exist (e.g., "max 30% per sector"), create sector-level constraints
- If stock constraints exist (e.g., "max 10% per stock"), create individual stock constraints
- Example: For 5 stocks in 4 sectors, you need 5 stock variables, not 4 sector variables

         ## SCHEDULING PROBLEMS
         If this is a scheduling problem:
         - For "N nurses × D days × S shifts", create N×D×S variables (x_nurse_day_shift)
         - For "T tasks × R resources", create T×R variables (x_task_resource)
         - For "P projects × D developers", create P×D variables (x_project_developer)
         - Example: 10 nurses × 7 days × 3 shifts = 210 variables, not 1 generic variable
         
         **NURSE SCHEDULING EXAMPLE**:
         For "3 nurses × 2 days × 2 shifts", create these 12 variables:
         - x_nurse1_day1_shift1, x_nurse1_day1_shift2, x_nurse1_day2_shift1, x_nurse1_day2_shift2
         - x_nurse2_day1_shift1, x_nurse2_day1_shift2, x_nurse2_day2_shift1, x_nurse2_day2_shift2  
         - x_nurse3_day1_shift1, x_nurse3_day1_shift2, x_nurse3_day2_shift1, x_nurse3_day2_shift2
         Each variable should be binary (0 or 1) with descriptive names including all dimensions.

## ROUTING PROBLEMS
If this is a routing problem:
- For "V vehicles × L locations", create V×L variables (x_vehicle_location)
- For "V vehicles × C customers", create V×C variables (x_vehicle_customer)
- For "V vehicles × L locations × T time periods", create V×L×T variables
- Example: 5 vehicles × 20 customers = 100 variables, not 2 generic variables

## PRODUCTION PLANNING
If this is a production planning problem:
- For "P products × T time periods", create P×T variables (x_product_time)
- For "P products × M machines × T time periods", create P×M×T variables
- For "P products × W warehouses", create P×W variables (x_product_warehouse)
- Example: 4 products × 6 months = 24 variables, not 2 generic variables

## RESOURCE ALLOCATION
If this is a resource allocation problem:
- For "R resources × T tasks", create R×T variables (x_resource_task)
- For "R resources × P projects × T time periods", create R×P×T variables
- For "A agents × J jobs", create A×J variables (x_agent_job)
- Example: 5 developers × 3 projects = 15 variables, not 2 generic variables

# OUTPUT FORMAT
Provide JSON with this EXACT structure:

{{
  "reasoning_steps": {{
    "step1_decision_analysis": "List of key decisions identified",
    "step2_constraint_analysis": "List of limitations and requirements", 
    "step3_objective_analysis": "Goal and optimization target",
    "step4_variable_design": "How decisions translate to variables",
    "step5_constraint_formulation": "How limitations translate to constraints",
    "step6_objective_formulation": "How goal translates to objective function",
    "step7_validation": "Verification that all variables are used"
  }},
  "model_type": "{opt_type}",
  "variables": [
    {{
      "name": "x1",
      "type": "continuous", 
      "bounds": "0 to 1",
      "description": "Allocation to stock 1 (fraction)"
    }}
  ],
  "objective": {{
    "type": "maximize",
    "expression": "0.12*x1 + 0.08*x2 + 0.10*x3 + 0.06*x4",
    "description": "Expected portfolio return"
  }},
  "constraints": [
    {{
      "expression": "x1 + x2 + x3 + x4 = 1",
      "description": "Total allocation must equal 100%"
    }}
  ],
  "model_complexity": "medium",
  "estimated_solve_time": 0.1,
  "mathematical_formulation": "Complete mathematical description based on reasoning steps",
  "validation_summary": {{
    "variables_defined": 4,
    "constraints_defined": 5,
    "objective_matches_problem": true,
    "model_is_feasible": true,
    "all_variables_used": true,
    "reasoning_completed": true
  }}
}}

Respond with valid JSON only:"""
                
                resp = await self.bedrock.invoke("anthropic.claude-3-haiku-20240307-v1:0", prompt, 6000)
                result = self._parse_json(resp)
                
                # Debug output
                logger.info(f"Model building attempt {attempt+1}:")
                logger.info(f"Raw response length: {len(resp) if resp else 0}")
                logger.info(f"Raw response preview: {resp[:200] if resp else 'None'}...")
                logger.info(f"Generated result keys: {list(result.keys()) if result else 'None'}")
                if result and 'raw_response' in result:
                    logger.info(f"Raw response in result: {result['raw_response'][:200]}...")
                if result and 'reasoning_steps' in result:
                    logger.info(f"Reasoning steps keys: {list(result['reasoning_steps'].keys()) if result['reasoning_steps'] else 'None'}")
                
                if self._validate_model_v2(result):
                    result.setdefault('model_type', opt_type)
                    
                    # Try to build MathOpt model if available
                    mathopt_result = None
                    if HAS_MATHOPT:
                        try:
                            mathopt_builder = MathOptModelBuilder()
                            mathopt_result = mathopt_builder.build_model_from_reasoning(result)
                            if mathopt_result.get('status') == 'success':
                                result['mathopt_model'] = mathopt_result
                                logger.info("MathOpt model built successfully")
                        except Exception as e:
                            logger.warning(f"MathOpt model building failed: {e}")
                    
                    return {
                        "status": "success",
                        "step": "model_building",
                        "timestamp": datetime.now().isoformat(),
                        "result": result,
                        "message": f"Model built with 7-step reasoning{' + MathOpt' if mathopt_result and mathopt_result.get('status') == 'success' else ''} (attempt {attempt+1})"
                    }
                else:
                    logger.warning(f"Model validation failed on attempt {attempt+1}")
                    if result:
                        logger.warning(f"Missing keys: {[k for k in ['variables', 'constraints', 'objective', 'reasoning_steps'] if k not in result]}")
                        if 'reasoning_steps' in result:
                            required_steps = ['step1_decision_analysis', 'step2_constraint_analysis', 'step3_objective_analysis', 'step4_variable_design', 'step5_constraint_formulation', 'step6_objective_formulation', 'step7_validation']
                            missing_steps = [s for s in required_steps if s not in result['reasoning_steps']]
                            if missing_steps:
                                logger.warning(f"Missing reasoning steps: {missing_steps}")
            
            return {"status": "error", "step": "model_building", "error": "Validation failed after retries"}
            
        except Exception as e:
            logger.error(f"Model error: {e}")
            return {"status": "error", "step": "model_building", "error": str(e)}
    
    def _validate_model(self, data: Dict) -> bool:
        if not data.get('variables') or not data.get('constraints') or not data.get('objective'):
            return False
        
        var_names = {v['name'] for v in data['variables'] if isinstance(v, dict)}
        all_text = ' '.join(c.get('expression', '') for c in data['constraints'] if isinstance(c, dict))
        all_text += ' ' + data['objective'].get('expression', '') if isinstance(data.get('objective'), dict) else ''
        
        return all(name in all_text for name in var_names)
    
    def _validate_model_v2(self, data: Dict) -> bool:
        """Enhanced validation for v2.0 models with 7-step reasoning."""
        # Basic structure validation
        if not data.get('variables') or not data.get('constraints') or not data.get('objective'):
            return False
        
        # Check for reasoning steps
        if not data.get('reasoning_steps'):
            return False
        
        # Check that all 7 steps are present
        required_steps = [
            'step1_decision_analysis', 'step2_constraint_analysis', 'step3_objective_analysis',
            'step4_variable_design', 'step5_constraint_formulation', 'step6_objective_formulation',
            'step7_validation'
        ]
        reasoning_steps = data.get('reasoning_steps', {})
        if not all(step in reasoning_steps for step in required_steps):
            return False
        
        # Variable usage validation
        var_names = {v['name'] for v in data['variables'] if isinstance(v, dict)}
        all_text = ' '.join(c.get('expression', '') for c in data['constraints'] if isinstance(c, dict))
        all_text += ' ' + data['objective'].get('expression', '') if isinstance(data.get('objective'), dict) else ''
        
        # All variables must be used
        if not all(name in all_text for name in var_names):
            return False
        
        # Check validation summary
        validation_summary = data.get('validation_summary', {})
        if not validation_summary.get('all_variables_used', False):
            return False
        
        return True
    
    async def solve_optimization(
        self,
        problem_description: str,
        intent_data: Optional[Dict] = None,
        data_analysis: Optional[Dict] = None,
        model_building: Optional[Dict] = None
    ) -> Dict[str, Any]:
        try:
            if not model_building or 'result' not in model_building:
                return {"status": "error", "error": "No model"}
            
            model_spec = ModelSpec.from_dict(model_building['result'])
            solver_result = solve_real_optimization(model_spec.to_dict())
            
            validation = self.validator.validate(solver_result, model_spec)
            
            if not validation['is_valid'] and solver_result.get('status') == 'optimal':
                logger.warning(f"Validation errors: {validation['errors']}")
            
            solver_result['validation'] = validation
            
            return {
                "status": "success" if validation['is_valid'] else "error",
                "step": "optimization_solution",
                "timestamp": datetime.now().isoformat(),
                "result": solver_result,
                "message": f"Solved: {solver_result.get('status')}"
            }
            
        except Exception as e:
            logger.error(f"Solve error: {e}")
            return {"status": "error", "step": "optimization_solution", "error": str(e)}
    
    async def select_solver(self, optimization_type: str, problem_size: Optional[Dict] = None, performance_requirement: str = "balanced") -> Dict[str, Any]:
        try:
            result = self.solver_selector.select_solver(optimization_type, problem_size or {}, performance_requirement)
            return {
                "status": "success",
                "step": "solver_selection",
                "timestamp": datetime.now().isoformat(),
                "result": result,
                "message": f"Selected: {result['selected_solver']}"
            }
        except Exception as e:
            return {"status": "error", "step": "solver_selection", "error": str(e)}
    
    async def explain_optimization(
        self,
        problem_description: str,
        intent_data: Optional[Dict] = None,
        data_analysis: Optional[Dict] = None,
        model_building: Optional[Dict] = None,
        optimization_solution: Optional[Dict] = None
    ) -> Dict[str, Any]:
        try:
            # Validate that we have actual optimization results to explain
            if not optimization_solution or optimization_solution.get('status') != 'success':
                return {
                    "status": "error",
                    "step": "explainability",
                    "error": "Cannot explain optimization results: No successful optimization found",
                    "message": "Optimization must be completed successfully before explanation can be generated"
                }
            
            # Validate that we have actual results
            result_data = optimization_solution.get('result', {})
            if not result_data or result_data.get('status') != 'optimal':
                return {
                    "status": "error", 
                    "step": "explainability",
                    "error": "Cannot explain optimization results: No optimal solution found",
                    "message": "Optimal solution required for business explanation"
                }
            
            status = result_data.get('status', 'unknown')
            objective_value = result_data.get('objective_value', 0)
            optimal_values = result_data.get('optimal_values', {})
            
            prompt = f"""Explain optimization result to business stakeholders.

PROBLEM: {problem_description}
OPTIMIZATION STATUS: {status}
OBJECTIVE VALUE: {objective_value}
OPTIMAL VALUES: {optimal_values}

IMPORTANT: Only provide explanations based on the actual optimization results above. Do not make up or estimate values.

JSON only:
{{
  "executive_summary": {{
    "problem_statement": "Clear statement of the original problem",
    "key_findings": ["actual finding 1", "actual finding 2"],
    "business_impact": "Actual quantified impact based on objective value"
  }},
  "implementation_guidance": {{
    "next_steps": ["step 1", "step 2"],
    "risk_considerations": ["risk 1"]
  }}
}}"""
            
            resp = await self.bedrock.invoke("anthropic.claude-3-haiku-20240307-v1:0", prompt, 3000)
            result = self._parse_json(resp)
            
            return {
                "status": "success",
                "step": "explainability",
                "timestamp": datetime.now().isoformat(),
                "result": result,
                "message": "Explanation generated based on actual optimization results"
            }
        except Exception as e:
            return {"status": "error", "step": "explainability", "error": str(e)}
    
    async def simulate_scenarios(
        self,
        problem_description: str,
        optimization_solution: Optional[Dict] = None,
        scenario_parameters: Optional[Dict] = None,
        simulation_type: str = "monte_carlo",
        num_trials: int = 10000
    ) -> Dict[str, Any]:
        try:
            # Validate that we have actual optimization results to simulate
            if not optimization_solution or optimization_solution.get('status') != 'success':
                return {
                    "status": "error",
                    "step": "simulation_analysis",
                    "error": "Cannot simulate scenarios: No successful optimization found",
                    "message": "Optimization must be completed successfully before simulation can be performed"
                }
            
            # Validate that we have actual results
            result_data = optimization_solution.get('result', {})
            if not result_data or result_data.get('status') != 'optimal':
                return {
                    "status": "error",
                    "step": "simulation_analysis", 
                    "error": "Cannot simulate scenarios: No optimal solution found",
                    "message": "Optimal solution required for scenario simulation"
                }
            
            if simulation_type != "monte_carlo" or not HAS_MONTE_CARLO:
                return {
                    "status": "error",
                    "error": f"Only Monte Carlo supported (NumPy required)",
                    "available_simulations": ["monte_carlo"],
                    "roadmap": ["discrete_event", "agent_based"]
                }
            
            obj_val = result_data.get('objective_value', 0)
            if obj_val == 0:
                return {
                    "status": "error",
                    "step": "simulation_analysis",
                    "error": "Cannot simulate scenarios: Zero objective value",
                    "message": "Valid objective value required for meaningful simulation"
                }
            
            np.random.seed(42)
            scenarios = np.random.normal(obj_val, obj_val * 0.1, num_trials)
            
            return {
                "status": "success",
                "step": "simulation_analysis",
                "timestamp": datetime.now().isoformat(),
                "result": {
                    "simulation_type": "monte_carlo",
                    "num_trials": num_trials,
                    "risk_metrics": {
                        "mean": float(np.mean(scenarios)),
                        "std_dev": float(np.std(scenarios)),
                        "percentile_5": float(np.percentile(scenarios, 5)),
                        "percentile_95": float(np.percentile(scenarios, 95))
                    }
                },
                "message": f"Monte Carlo completed ({num_trials} trials) based on actual optimization results"
            }
        except Exception as e:
            return {"status": "error", "step": "simulation_analysis", "error": str(e)}
    
    async def get_workflow_templates(self) -> Dict[str, Any]:
        try:
            return {
                "status": "success",
                "workflow_templates": self.workflow_manager.get_all_workflows(),
                "total_workflows": 21
            }
        except Exception as e:
            return {"status": "error", "error": str(e)}
    
    async def execute_workflow(self, industry: str, workflow_id: str, user_input: Optional[Dict] = None) -> Dict[str, Any]:
        try:
            problem_desc = f"{workflow_id} for {industry}"
            
            intent_result = await self.classify_intent(problem_desc, industry)
            data_result = await self.analyze_data(problem_desc, intent_result.get('result'))
            model_result = await self.build_model(problem_desc, intent_result.get('result'), data_result.get('result'))
            solve_result = await self.solve_optimization(problem_desc, intent_result.get('result'), data_result.get('result'), model_result)
            explain_result = await self.explain_optimization(problem_desc, intent_result.get('result'), data_result.get('result'), model_result, solve_result.get('result'))
            
            return {
                "status": "success",
                "workflow_type": workflow_id,
                "industry": industry,
                "steps_completed": 5,
                "results": {
                    "intent_classification": intent_result,
                    "data_analysis": data_result,
                    "model_building": model_result,
                    "optimization_solution": solve_result,
                    "explainability": explain_result
                }
            }
        except Exception as e:
            return {"status": "error", "error": str(e)}

# ============================================================================
# GLOBAL INSTANCE & WRAPPERS
# ============================================================================

_tools_instance = None

def get_tools() -> DcisionAITools:
    global _tools_instance
    if _tools_instance is None:
        _tools_instance = DcisionAITools()
    return _tools_instance

async def classify_intent(problem_description: str, context: Optional[str] = None) -> Dict[str, Any]:
    return await get_tools().classify_intent(problem_description, context)

async def analyze_data(problem_description: str, intent_data: Optional[Dict] = None) -> Dict[str, Any]:
    return await get_tools().analyze_data(problem_description, intent_data)

async def build_model(problem_description: str, intent_data: Optional[Dict] = None, data_analysis: Optional[Dict] = None, solver_selection: Optional[Dict] = None) -> Dict[str, Any]:
    return await get_tools().build_model(problem_description, intent_data, data_analysis, solver_selection)

async def solve_optimization(problem_description: str, intent_data: Optional[Dict] = None, data_analysis: Optional[Dict] = None, model_building: Optional[Dict] = None) -> Dict[str, Any]:
    return await get_tools().solve_optimization(problem_description, intent_data, data_analysis, model_building)

async def select_solver(optimization_type: str, problem_size: Optional[Dict] = None, performance_requirement: str = "balanced") -> Dict[str, Any]:
    return await get_tools().select_solver(optimization_type, problem_size, performance_requirement)

async def explain_optimization(problem_description: str, intent_data: Optional[Dict] = None, data_analysis: Optional[Dict] = None, model_building: Optional[Dict] = None, optimization_solution: Optional[Dict] = None) -> Dict[str, Any]:
    return await get_tools().explain_optimization(problem_description, intent_data, data_analysis, model_building, optimization_solution)

async def simulate_scenarios(problem_description: str, optimization_solution: Optional[Dict] = None, scenario_parameters: Optional[Dict] = None, simulation_type: str = "monte_carlo", num_trials: int = 10000) -> Dict[str, Any]:
    return await get_tools().simulate_scenarios(problem_description, optimization_solution, scenario_parameters, simulation_type, num_trials)

async def get_workflow_templates() -> Dict[str, Any]:
    return await get_tools().get_workflow_templates()

async def execute_workflow(industry: str, workflow_id: str, user_input: Optional[Dict] = None) -> Dict[str, Any]:
    return await get_tools().execute_workflow(industry, workflow_id, user_input)
