#!/usr/bin/env python3

import sys
import os
import time
import argparse
from pathlib import Path
from openai import OpenAI

import config
from config import *
from file_manager import setup_logging, load_tables_from_dir, save_alternating_dataset_instance
from core_generation import generate_alternating_optimization_problem_with_verification
from utils import should_skip_database
from prompts import create_initial_or_expert_prompt
from solver_execution import run_code_in_virtual_env

def main():
    """Main function with argument parsing and execution"""
    parser = argparse.ArgumentParser(
        description="Enhanced Alternating Optimization Problem Generator with Incremental Intelligent Verification Diagnosis (Singleton Version)"
    )
    parser.add_argument(
        "--database", 
        required=True,
        help="Database name to process"
    )
    parser.add_argument(
        "--spider-dir", 
        default=SPIDER_DIR,
        help=f"Path to Spider dataset directory (default: {SPIDER_DIR})"
    )
    parser.add_argument(
        "--output-dir", 
        default=OUTPUT_BASE_DIR,  
        help=f"Output directory (default: {OUTPUT_BASE_DIR})"
    )

    args = parser.parse_args()
    
    # Update config module values based on arguments
    if args.output_dir != config.OUTPUT_BASE_DIR:
        config.OUTPUT_BASE_DIR = args.output_dir
    if args.spider_dir != config.SPIDER_DIR:
        config.SPIDER_DIR = args.spider_dir

    # Setup logging for the specific database
    logger = setup_logging(args.database, args.output_dir)
    
    logger.info("FIXED: Enhanced Alternating Optimization Problem Generator with Incremental Intelligent Verification Diagnosis (Singleton Version)")
    logger.info("=" * 200)
    logger.info(f"Processing single database: {args.database}")
    logger.info(f"Model: {MODEL_NAME}")
    logger.info(f"Endpoint: {BASE_URL}")
    logger.info(f"Alternating Optimization Algorithm: OR Expert <-> Data Engineer (max {MAX_ALTERNATING_ITERATIONS} iterations)")
    logger.info(f"Professional Boundaries: OR Expert (optimization modeling only) | Data Engineer (all database decisions)")
    logger.info(f"Business Configuration: Scalar parameters with sample_value + Business logic formulas with actual expressions")
    logger.info(f"Mathematical Constraints: Linear Programming (LP) and Mixed-Integer Programming (MIP) ONLY")
    logger.info(f"Optimization Environment: {OPTIM_VENV_PATH}")
    logger.info(f"Constraint Count Guidance: {MIN_CONSTRAINTS} to {MAX_CONSTRAINTS} constraints for optimization feasibility")
    logger.info(f"Sample Data Guidance: {MIN_SAMPLE_ROWS} to {MAX_SAMPLE_ROWS} rows per table for realistic scenarios")
    logger.info(f"Row Count Rule: Tables must generate at least {MIN_SAMPLE_ROWS} meaningful rows, otherwise moved to business_configuration_logic.json")
    logger.info(f"Triple Expert: Business + Data + Optimization expert for realistic data including business configuration")
    logger.info(f"FIXED INCREMENTAL INTELLIGENT VERIFICATION: Enhanced root cause diagnosis with targeted minimal modifications (threshold {VERIFICATION_THRESHOLD}, max {MAX_VERIFICATION_ATTEMPTS} attempts)")
    logger.info(f"Template-Guided Solver Execution: Gurobipy 12.0.2 + DOCplex 2.29.245 + Pyomo 6.9.2 with {SOLVER_TIMEOUT}s timeout")
    logger.info(f"Cross-Solver Analysis: OR Expert analyzes all solver results and provides final business recommendations")
    logger.info(f"Retry Logic: Automatic retry up to {MAX_JSON_REPROMPT_ATTEMPTS} attempts when JSON parsing fails, up to {MAX_SOLVER_RETRY_ATTEMPTS} attempts when solver values are inconsistent")
    logger.info(f"Solver Consistency Tolerance: {SOLVER_CONSISTENCY_TOLERANCE}")
    logger.info(f"FIXED FEATURE: Initial mathematical solution generation before verification ensures Section 4 is always included")
    logger.info(f"NO FALLBACK: All failures result in intelligent incremental retries with diagnostic guidance, no fallback responses generated")

    # Setup API client
    os.environ['RITS_API_KEY'] = RITS_API_KEY
    api_key = os.environ.get("RITS_API_KEY")
    if not api_key:
        logger.error("ERROR: Please set RITS_API_KEY environment variable")
        sys.exit(1)

    client = OpenAI(
        api_key="dummy",
        base_url=BASE_URL,
        default_headers={"RITS_API_KEY": api_key},
        timeout=300
    )

    # Test API connection
    logger.info("\nTesting Llama-3.3-70B API connection...")
    try:
        test_response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "user", "content": "Hello, please respond with 'FIXED Enhanced Alternating Optimization API connection successful with incremental intelligent verification (Singleton)'"}],
            max_tokens=50,
            temperature=0.1
        )
        if test_response.choices:
            logger.info(f"SUCCESS: {test_response.choices[0].message.content}")
        else:
            logger.warning("WARNING: API responded but no content received")
    except Exception as e:
        logger.error(f"ERROR: API connection failed: {e}")
        logger.info("Continuing anyway...")

    # Test optimization environment
    logger.info(f"\nTesting optimization environment: {OPTIM_VENV_PATH}")
    if os.path.exists(OPTIM_VENV_PATH):
        logger.info(f"SUCCESS: Optimization environment found at {OPTIM_VENV_PATH}")
        
        # Test basic packages
        try:
            test_code = "import gurobipy; import docplex; import pyomo; print('All solver packages available for FIXED enhanced linear optimization with incremental intelligent verification (Singleton)')"
            stdout, stderr = run_code_in_virtual_env(test_code)
            if stderr is None:
                logger.info(f"SUCCESS: {stdout.strip()}")
            else:
                logger.warning(f"WARNING: Package test failed: {stderr}")
        except Exception as e:
            logger.warning(f"WARNING: Could not test optimization packages: {e}")
    else:
        logger.warning(f"WARNING: Optimization environment not found at {OPTIM_VENV_PATH}")
        logger.warning("Solver execution may fail. Please set OPTIM_VENV_PATH correctly.")

    # Test template loading
    logger.info("\nTesting solver template loading...")
    from solver_execution import load_solver_templates
    templates = load_solver_templates()
    
    template_status = []
    for solver_type in ["gurobipy", "docplex", "pyomo"]:
        if templates.get(solver_type):
            template_status.append(f"{solver_type}: {len(templates[solver_type])} chars")
        else:
            template_status.append(f"{solver_type}: missing")
    
    logger.info(f"Template status: {', '.join(template_status)}")

    # Load Spider data
    tables_data = load_tables_from_dir(args.spider_dir)
    
    # Find the specific database
    target_db = None
    for db in tables_data:
        if db['db_id'] == args.database:
            target_db = db
            break
    
    if not target_db:
        logger.error(f"ERROR: Database '{args.database}' not found in Spider dataset")
        sys.exit(1)
    
    # Check if database should be skipped
    if should_skip_database(target_db):
        table_count = len(target_db['table_names_original'])
        logger.error(f"ERROR: Database '{args.database}' has {table_count} tables (> {MAX_TABLES} limit)")
        sys.exit(1)
    
    original_table_count = len(target_db['table_names_original'])
    logger.info(f"\nProcessing: {args.database} ({original_table_count} tables)")
    
    # Create output directory
    os.makedirs(config.OUTPUT_BASE_DIR, exist_ok=True)
    
    # Process the specific database
    start_time = time.time()
    
    try:
        # Enhanced alternating optimization problem generation with incremental intelligent verification
        logger.info(f"  FIXED Enhanced alternating optimization generation with incremental intelligent verification diagnosis...")
        (complete_documentation, problem_description, mathematical_solution,
         final_or_analysis, final_implementation, 
         iteration_history, debug_history, verification_history,
         expected_alternating_formulation, final_verification_result, 
         consistency_score, solver_results, solver_codes, 
         triple_expert_result, solver_analysis_result) = generate_alternating_optimization_problem_with_verification(
            client, MODEL_NAME, target_db, None)
        
        time.sleep(1)
        
        # Get table info for saving
        _, table_info = create_initial_or_expert_prompt(target_db)
        
        save_alternating_dataset_instance(
            config.OUTPUT_BASE_DIR, args.database, table_info,
            complete_documentation, problem_description, mathematical_solution,
            final_or_analysis, final_implementation, 
            iteration_history, debug_history, verification_history,
            expected_alternating_formulation, final_verification_result, consistency_score,
            solver_results, solver_codes, triple_expert_result, solver_analysis_result)
        
        execution_time = time.time() - start_time
        
        # Log final statistics
        iterations_completed = len([h for h in iteration_history if h.get("type") == "or_expert"])
        final_table_count = len(final_implementation.get("schema_adjustment_decisions", {}).get("tables_to_create", []))
        business_config_params = len(final_implementation.get("business_configuration_logic_updates", {}).get("configuration_parameters", {}))
        
        # Extract solver status information
        gurobipy_status = solver_results.get("gurobipy", {}).get("status", "unknown")
        docplex_status = solver_results.get("docplex", {}).get("status", "unknown")
        pyomo_status = solver_results.get("pyomo", {}).get("status", "unknown")
        gurobipy_value = solver_results.get("gurobipy", {}).get("optimal_value")
        docplex_value = solver_results.get("docplex", {}).get("optimal_value")
        pyomo_value = solver_results.get("pyomo", {}).get("optimal_value")
        gurobipy_time = solver_results.get("gurobipy", {}).get("execution_time", 0)
        docplex_time = solver_results.get("docplex", {}).get("execution_time", 0)
        pyomo_time = solver_results.get("pyomo", {}).get("execution_time", 0)
        gurobipy_retry = solver_results.get("gurobipy", {}).get("retry_attempt", "N/A")
        docplex_retry = solver_results.get("docplex", {}).get("retry_attempt", "N/A")
        pyomo_retry = solver_results.get("pyomo", {}).get("retry_attempt", "N/A")
        
        # Extract cross-solver analysis information
        consistency_evaluation = solver_analysis_result.get("consistency_evaluation", {})
        final_recommendation = solver_analysis_result.get("final_recommendation", {})
        values_consistent = consistency_evaluation.get("values_consistent", False)
        final_confidence = final_recommendation.get("confidence", "unknown")
        recommended_value = final_recommendation.get("recommended_optimal_value")
        
        logger.info(f"\nSingle Database Processing Complete")
        logger.info(f"Database: {args.database}")
        logger.info(f"Total execution time: {execution_time:.1f} seconds")
        logger.info(f"Alternating iterations completed: {iterations_completed}")
        logger.info(f"Business configuration parameters: {business_config_params}")
        logger.info(f"Mathematical verification score: {consistency_score:.3f}")
        logger.info(f"Tables: {original_table_count} -> {final_table_count}")
        logger.info(f"")
        logger.info(f"Solver Results:")
        logger.info(f"  Gurobipy: {gurobipy_status} (value: {gurobipy_value}, time: {gurobipy_time:.2f}s, retry: {gurobipy_retry})")
        logger.info(f"  DOCplex: {docplex_status} (value: {docplex_value}, time: {docplex_time:.2f}s, retry: {docplex_retry})")
        logger.info(f"  Pyomo: {pyomo_status} (value: {pyomo_value}, time: {pyomo_time:.2f}s, retry: {pyomo_retry})")
        logger.info(f"")
        logger.info(f"Cross-Solver Analysis:")
        logger.info(f"  Values consistent: {values_consistent}")
        logger.info(f"  Final recommended value: {recommended_value}")
        logger.info(f"  Confidence: {final_confidence}")
        logger.info(f"")
        logger.info(f"Business context: {final_or_analysis.get('business_context', 'Unknown')[:50]}...")
        logger.info(f"Optimization goal: {final_or_analysis.get('optimization_formulation', {}).get('objective', 'Unknown')[:50]}...")
        logger.info(f"Triple expert data realistic: {triple_expert_result.get('validation', {}).get('business_realistic', False)}")
        
        # Final success message
        verification_status = "VERIFIED" if consistency_score >= VERIFICATION_THRESHOLD else "UNVERIFIED"
        convergence_status = "CONVERGED" if iterations_completed < MAX_ALTERNATING_ITERATIONS else "MAX_ITER"
        data_status = "REALISTIC" if triple_expert_result.get('validation', {}).get('business_realistic', False) else "BASIC"
        
        logger.info(f"\nFinal Status: SUCCESS [{verification_status}] [{convergence_status}] [{data_status}] [LINEAR] [INCREMENTAL_INTELLIGENT]")
        logger.info(f"Output saved to: {args.output_dir}/{args.database}/")
        
        # Log file location
        log_file = Path(args.output_dir) / args.database / "logs" / f"{args.database}_generation.log"
        logger.info(f"Detailed log saved to: {log_file}")
        
    except Exception as e:
        execution_time = time.time() - start_time
        logger.error(f"ERROR: Failed to process {args.database}: {e}")
        logger.error(f"Execution time before failure: {execution_time:.1f} seconds")
        
        # Log file location even on failure
        log_file = Path(args.output_dir) / args.database / "logs" / f"{args.database}_generation.log"
        logger.info(f"Error log saved to: {log_file}")
        
        sys.exit(1)

if __name__ == "__main__":
    main()