#!/usr/bin/env python3

import time
import logging
import json
from config import *
from api_client import call_llama_api, call_llama_api_with_reprompt
from prompts import (create_mathematical_verification_prompt, 
                    create_detailed_verification_diagnosis_prompt,
                    create_mathematical_solution_prompt,
                    create_mathematical_solution_prompt_with_lesson,
                    create_incremental_problem_description_prompt,
                    create_problem_description_prompt)

def extract_mathematical_formulation_from_alternating(final_or_analysis, final_implementation):
    """Extract the intended mathematical formulation from converged alternating optimization iterations"""
    
    math_mapping = final_or_analysis.get("current_optimization_to_schema_mapping", {})
    data_mapping = final_implementation.get("data_mapping", {})
    optimization_formulation = final_or_analysis.get("optimization_formulation", {})
    
    return {
        "objective_function": {
            "goal": optimization_formulation.get("objective", "").split()[0] if optimization_formulation.get("objective") else "",
            "description": optimization_formulation.get("objective", ""),
            "coefficient_sources": math_mapping.get("objective_coefficients", {})
        },
        "decision_variables": {
            "definition": optimization_formulation.get("decision_variables", ""),
            "source_mapping": math_mapping.get("decision_variables", {})
        },
        "constraints": [
            {
                "type": "resource_constraint",
                "description": optimization_formulation.get("constraints", ""),
                "coefficient_sources": math_mapping.get("constraint_bounds", {}),
                "mathematical_form": "Standard constraint form based on business requirements"
            }
        ],
        "sample_coefficients": data_mapping.get("sample_data_rows", {})
    }

def verify_mathematical_consistency(client, model, problem_description, mathematical_solution, expected_alternating_formulation, db_id):
    """Verify consistency between generated mathematical solution and expected alternating formulation"""
    
    logging.info(f"    Verifying mathematical consistency...")
    
    verification_prompt = create_mathematical_verification_prompt(
        problem_description, mathematical_solution, expected_alternating_formulation, db_id
    )
    
    try:
        verification_response = call_llama_api_with_reprompt(
            client, model, verification_prompt,
            f"Mathematical Verification - {db_id}",
            "Mathematical Verification", db_id, "verification"
        )
        
        verification_result = verification_response
        
        consistency_score = verification_result.get("verification_results", {}).get("consistency_score", 0.0)
        passes_verification = consistency_score >= VERIFICATION_THRESHOLD
        
        logging.info(f"    Verification score: {consistency_score:.2f}, Passes: {passes_verification}")
        
        return verification_result, consistency_score, passes_verification
        
    except Exception as e:
        logging.error(f"    ERROR: Mathematical verification failed: {e}")
        raise e

def verify_with_targeted_incremental_regeneration(client, model, original_seed_problem, problem_description, mathematical_solution, expected_formulation, db_id):
    """Enhanced intelligent verification with incremental targeted regeneration based on detailed diagnosis"""
    
    current_problem_description = problem_description
    current_mathematical_solution = mathematical_solution
    
    for attempt in range(MAX_VERIFICATION_ATTEMPTS + 1):
        logging.info(f"    Incremental verification attempt {attempt + 1}...")
        
        # FIXED: If this is the first attempt and we don't have a mathematical solution, generate one
        if attempt == 0 and not current_mathematical_solution:
            logging.info(f"    Generating initial mathematical solution...")
            math_solution_prompt = create_mathematical_solution_prompt(current_problem_description, db_id)
            current_mathematical_solution = call_llama_api(
                client, model, math_solution_prompt,
                f"Initial Mathematical Solution Generation - {db_id}"
            )
            time.sleep(1)
        
        # If this is not the first attempt and we don't have a mathematical solution, generate one
        if attempt > 0 and not current_mathematical_solution:
            math_solution_prompt = create_mathematical_solution_prompt(current_problem_description, db_id)
            current_mathematical_solution = call_llama_api(
                client, model, math_solution_prompt,
                f"Mathematical Solution Generation Attempt {attempt} - {db_id}"
            )
            time.sleep(1)
        
        # Verify current version
        verification_result, consistency_score, passes_verification = verify_mathematical_consistency(
            client, model, current_problem_description, current_mathematical_solution, expected_formulation, db_id
        )
        
        if passes_verification:
            logging.info(f"    Incremental verification passed on attempt {attempt + 1} (score: {consistency_score:.2f})")
            return current_problem_description, current_mathematical_solution, verification_result, consistency_score, True
            
        if attempt >= MAX_VERIFICATION_ATTEMPTS:
            logging.info(f"    Maximum incremental verification attempts reached. Using best result (score: {consistency_score:.2f})")
            break
            
        # Enhanced diagnosis with specific issue identification
        logging.info(f"    Detailed diagnosis of verification failure (attempt {attempt + 1})...")
        diagnosis_prompt = create_detailed_verification_diagnosis_prompt(
            original_seed_problem, current_problem_description, current_mathematical_solution, 
            verification_result, expected_formulation, db_id
        )
        
        try:
            diagnosis = call_llama_api_with_reprompt(
                client, model, diagnosis_prompt,
                f"Detailed Verification Diagnosis Attempt {attempt + 1} - {db_id}",
                "Detailed Verification Diagnosis", db_id, f"diagnosis_{attempt + 1}"
            )
            
            # Extract diagnosis results
            primary_issue = diagnosis.get("root_cause_analysis", {}).get("primary_issue_source", "mathematical_solution")
            lesson = diagnosis.get("targeted_lesson", {})
            modification_strategy = lesson.get("modification_strategy", "incremental")
            
            logging.info(f"    Diagnosis: Primary issue source = {primary_issue}")
            logging.info(f"    Modification strategy = {modification_strategy}")
            logging.info(f"    Key insight: {lesson.get('key_insight', 'No insight provided')}")
            
            time.sleep(1)
            
            # Apply targeted incremental regeneration based on diagnosis
            if primary_issue in ["problem_description", "both"]:
                if modification_strategy == "incremental":
                    logging.info(f"    Making incremental improvements to problem description...")
                    current_problem_description = make_incremental_improvements(
                        current_problem_description, lesson, original_seed_problem, expected_formulation, db_id, client, model
                    )
                else:
                    logging.info(f"    Complete regeneration of problem description...")
                    current_problem_description = regenerate_problem_description_with_lesson(
                        original_seed_problem, lesson, expected_formulation, db_id, client, model
                    )
                
                time.sleep(1)
                
                # If we regenerated problem description, we need to regenerate mathematical solution too
                current_mathematical_solution = ""
                
            if primary_issue in ["mathematical_solution", "both"] or not current_mathematical_solution:
                logging.info(f"    Regenerating mathematical solution based on diagnosis...")
                math_solution_prompt = create_mathematical_solution_prompt_with_lesson(
                    current_problem_description, lesson, db_id
                )
                
                current_mathematical_solution = call_llama_api(
                    client, model, math_solution_prompt,
                    f"Mathematical Solution Regeneration Attempt {attempt + 1} - {db_id}"
                )
                time.sleep(1)
        
        except Exception as e:
            logging.error(f"    ERROR in incremental regeneration attempt {attempt + 1}: {e}")
            # Fall back to simple regeneration
            logging.info(f"    Falling back to simple mathematical solution regeneration...")
            math_solution_prompt = create_mathematical_solution_prompt(current_problem_description, db_id)
            current_mathematical_solution = call_llama_api(
                client, model, math_solution_prompt,
                f"Fallback Mathematical Solution Attempt {attempt + 1} - {db_id}"
            )
            time.sleep(1)
    
    return current_problem_description, current_mathematical_solution, verification_result, consistency_score, False

def make_incremental_improvements(current_description, lesson, original_seed_problem, expected_formulation, db_id, client, model):
    """Make incremental improvements to problem description"""
    
    incremental_prompt = create_incremental_problem_description_prompt(
        current_description,
        lesson,
        original_seed_problem.get("final_or_analysis", {}),
        original_seed_problem.get("final_implementation", {}),
        original_seed_problem.get("schema_sql", ""),
        original_seed_problem.get("data_sql", ""),
        original_seed_problem.get("data_dictionary", {}),
        original_seed_problem.get("business_configuration", {}),
        expected_formulation,
        db_id
    )
    
    improved_description = call_llama_api(
        client, model, incremental_prompt,
        f"Incremental Problem Description Improvement - {db_id}"
    )
    
    return improved_description

def regenerate_problem_description_with_lesson(original_seed_problem, lesson, expected_formulation, db_id, client, model):
    """Regenerate problem description with lesson from diagnosis"""
    
    lesson_section = ""
    if lesson:
        lesson_section = f"""
BASED ON VERIFICATION FAILURE DIAGNOSIS:
- Key Insight: {lesson.get('key_insight', '')}
- Specific Guidance: {lesson.get('specific_guidance', '')}
- Focus Areas: {', '.join(lesson.get('focus_areas', []))}

Please pay special attention to the above guidance when regenerating the problem description.
"""
    
    from utils import extract_business_configuration_requirements, format_business_requirements_for_prompt
    
    business_configuration = original_seed_problem.get("business_configuration", {})
    scalar_parameters, business_formulas = extract_business_configuration_requirements(business_configuration)
    business_requirements = format_business_requirements_for_prompt(scalar_parameters, business_formulas)
    
    data_dict_section = ""
    data_dictionary = original_seed_problem.get("data_dictionary", {})
    if data_dictionary:
        data_dict_section = f"""
DATA DICTIONARY:
{json.dumps(data_dictionary, indent=2)}
"""

    business_config_section = ""
    if business_configuration:
        business_config_section = f"""
BUSINESS CONFIGURATION:
{json.dumps(business_configuration, indent=2)}

Business Configuration Design: 
Our system separates business logic design from value determination:
- Configuration Logic (business_configuration_logic.json): Templates designed by data engineers with sample_value for scalars and actual formulas for business logic
- Configuration Values (business_configuration.json): Realistic values determined by domain experts for scalar parameters only
- Design Rationale: Ensures business logic consistency while allowing flexible parameter tuning
"""

    expected_objective = expected_formulation.get("objective_function", {}).get("description", "")
    expected_decision_vars = expected_formulation.get("decision_variables", {}).get("definition", "")
    expected_constraints = expected_formulation.get("constraints", [])
    
    auto_context_requirements = []
    auto_context_requirements.append(f"- Business decisions match expected decision variables: {expected_decision_vars}")
    auto_context_requirements.append(f"- Operational parameters align with expected linear objective: {expected_objective}")
    auto_context_requirements.extend([f"- {req}" for req in business_requirements])
    auto_context_requirements.extend([
        "- Use natural language to precisely describe linear mathematical relationships",
        "- NO mathematical formulas, equations, or symbolic notation",
        "- Present data as current operational information",
        "- Focus on precise operational decision-making that leads to linear formulations",
        "- Resource limitations match expected linear constraints",
        "- Avoid scenarios requiring variable products, divisions, or other nonlinear relationships",
        "- Include specific operational parameters that map to expected coefficient sources",
        "- Reference business configuration parameters where appropriate"
    ])

    final_or_analysis = original_seed_problem.get("final_or_analysis", {})
    schema_sql = original_seed_problem.get("schema_sql", "")
    data_sql = original_seed_problem.get("data_sql", "")

    prompt = f"""You are a business analyst creating structured optimization problem documentation.

{lesson_section}

DATA SOURCES EXPLANATION:
- FINAL OR ANALYSIS: Final converged optimization problem from alternating process (iteration {final_or_analysis.get('iteration', 'unknown')}), contains business context and schema mapping evaluation
- DATABASE SCHEMA: Current database structure after iterative adjustments  
- DATA DICTIONARY: Business meanings and optimization roles of tables and columns
- CURRENT STORED VALUES: Realistic business data generated by triple expert (business + data + optimization)
- BUSINESS CONFIGURATION: Scalar parameters and business logic formulas separated from table data

CRITICAL REQUIREMENTS: 
- Ensure problem description naturally leads to LINEAR or MIXED-INTEGER optimization formulation
- Make business context consistent with the intended decision variables and objectives
- Align constraint descriptions with expected mathematical constraints
- Ensure data descriptions map clearly to expected coefficient sources
- Maintain business authenticity while fixing mathematical consistency issues
- Avoid business scenarios that would naturally require nonlinear relationships (variable products, divisions, etc.)

AUTO-EXTRACTED CONTEXT REQUIREMENTS:
{chr(10).join(auto_context_requirements)}

FINAL OR ANALYSIS:
{json.dumps(final_or_analysis, indent=2)}

FINAL DATABASE SCHEMA:
```sql
{schema_sql}
```

CURRENT STORED VALUES:
```sql
{data_sql}
```
{data_dict_section}
{business_config_section}

TASK: Create structured markdown documentation for SECTIONS 1-3 ONLY (Problem Description), paying special attention to the diagnostic guidance above.

EXACT MARKDOWN STRUCTURE TO FOLLOW:

# Complete Optimization Problem and Solution: {db_id}

## 1. Problem Context and Goals

### Context  
[Regenerate business context that naturally aligns with LINEAR optimization formulation. Ensure:]
{chr(10).join(auto_context_requirements)}
- CRITICAL: Include ALL business configuration information (scalar parameters AND business logic formulas) in natural business language

### Goals  
[Regenerate goals that clearly lead to LINEAR mathematical objective:]
- Optimization goal: {expected_formulation.get("objective_function", {}).get("goal", "")}
- Metric to optimize: {expected_objective}
- Success measurement aligned with expected coefficient sources
- Use natural language to precisely describe linear optimization goal
- NO mathematical formulas, equations, or symbolic notation
- CRITICAL: Express all relevant business configuration logic in business terms

## 2. Constraints    

[Regenerate constraints that directly match expected LINEAR mathematical constraints:]
{chr(10).join([f"- Expected constraint: {c.get('description', '')} (Form: {c.get('mathematical_form', '')})" for c in expected_constraints])}

[Each constraint should be described in business terms that naturally lead to LINEAR mathematical forms (no variable products or divisions)]
- CRITICAL: Include ALL business configuration constraints and business logic formulas expressed in natural business language

## 3. Available Data  

### Database Schema  
```sql
{schema_sql}
```

### Data Dictionary  
[Create comprehensive business-oriented data dictionary mapping tables and columns to their business purposes and optimization roles - NOT technical database terms. Base this on the data_dictionary provided but rewrite in business language that clearly connects to the expected linear mathematical formulation]


### Current Stored Values  
```sql
{data_sql}
```

STOP HERE - Do not include section 4 (Mathematical Optimization Formulation).
"""
    
    regenerated_description = call_llama_api(
        client, model, prompt,
        f"Complete Problem Description Regeneration with Lesson - {db_id}"
    )
    
    return regenerated_description