#!/usr/bin/env python3
"""
OrchestratorAI: Autonomous Multi-Layer Script Orchestration System
"""

import argparse
import pickle
import random
from typing import Dict, List, Any
from datetime import datetime

from .crispo_core import (
    OrchestrationContext,
    LayerParameters,
    GAOptimizer,
    RLAgent,
    AttentionRouter,
    CodeGenerator,
    MetaLearner,
    TaskMetadata,
    Verifier,
    SkiRentalContext,
    OneMaxSearchContext
)
from .solution_registry import query_registry
from .advanced_crispo import (
    FederatedOptimizer
)
from . import transfer_learning_pipeline as tlp
from . import nas_pipeline as nas

# ============================================================================
# CRISPO ENGINE
# ============================================================================

class Crispo:
    """Main orchestration engine for the Crispo system.

    This class integrates all the core components (GA, RL, Code Generation,
    etc.) and advanced features (Transfer Learning, NAS) to execute the
    end-to-end script generation and optimization pipeline. It manages the
    flow of control through the different phases of orchestration.
    """

    def __init__(self, context: OrchestrationContext, meta_learner: MetaLearner):
        """Initializes the Crispo engine.

        Args:
            context (OrchestrationContext): The global context for this
                orchestration run, containing the project name and objective.
            meta_learner (MetaLearner): The MetaLearner instance, which may be
                pre-loaded with historical knowledge.
        """
        self.context = context
        self.ga_optimizer = GAOptimizer()
        self.rl_agent = RLAgent()
        self.attention_router = AttentionRouter()
        self.code_generator = CodeGenerator()
        self.meta_learner = meta_learner
        self.verifier = Verifier()

        # Initialize advanced components
        self.federated_optimizer = FederatedOptimizer(num_clients=10)

    def orchestrate(
        self,
        project_type: str,
        domain: str,
        complexity: float,
        enable_transfer_learning: bool,
        enable_nas: bool,
        enable_federated_optimization: bool,
        trust_parameter: float
    ) -> List[str]:
        """Executes the full, multi-phase orchestration pipeline.

        This method guides the process from strategy selection through code
        generation, optimization, and verification.

        Args:
            project_type (str): The category of the project (e.g., 'data_pipeline').
            domain (str): The application domain (e.g., 'finance').
            complexity (float): The complexity of the task, from 0.0 to 1.0.
            enable_transfer_learning (bool): Flag to enable the Transfer
                Learning engine.
            enable_nas (bool): Flag to enable Neural Architecture Search.
            enable_federated_optimization (bool): Flag to enable Federated
                Optimization.
            trust_parameter (float): The trust parameter (lambda) for
                learning-augmented algorithms.

        Returns:
            List[str]: A list of strings, where each string is the generated
                Python code for a single layer of the pipeline.
        """
        print("🚀 " + "="*68)
        print("ORCHESTRATOR AI: AUTONOMOUS EXECUTION")
        print(f"Project: {self.context.project}")
        print(f"Objective: {self.context.objective}")
        print("="*70)

        # Get optimal strategy from meta-learner
        strategy = self.meta_learner.get_optimal_strategy(project_type, complexity)
        generations = strategy.get('ga_generations', 10)
        episodes = strategy.get('rl_episodes', 5)

        print(f"🧠 Meta-Learner Strategy: GA Gens={generations}, RL Eps={episodes}")

        # Advanced features
        if enable_nas:
            print("\n🤖 PHASE 1.2: Production Neural Architecture Search")
            # This pipeline was generated by OrchestratorAI itself.
            nas_search_space = {
                'num_layers': [2, 3, 4],
                'learning_rate': [0.01, 0.005, 0.001]
            }
            best_arch = nas.select_best_architecture(nas_search_space, num_samples=5)

        # Define a template for layer parameters
        template_params = LayerParameters(
            layer_id=0,
            weights={'complexity': complexity, 'execution': 1.0},
            biases={'logging': 0.0, 'error_handling': 0.0},
            temperature=1.0
        )

        # Phase 1: Evolve base parameters with GA
        print("\n📊 PHASE 1: Genetic Algorithm Evolution")
        evolved_params = self.ga_optimizer.execute(
            template_params,
            context={'desired_complexity': complexity},
            generations=generations
        )

        if enable_transfer_learning:
            print("\n🔧 PHASE 1.5: Production Transfer Learning Pipeline")
            # This pipeline was generated by OrchestratorAI itself.
            model_path = f"model_store/{project_type}_model.json"
            loaded_model = tlp.load_model(model_path)
            evolved_params = tlp.apply_model(evolved_params, loaded_model)
            tlp.log_to_registry(model_path, project_type)

        # Phase 2: Fine-tune parameters with RL
        print("\n🎯 PHASE 2: Reinforcement Learning Fine-Tuning")
        final_params = self.rl_agent.execute(
            evolved_params,
            context={'complexity': complexity},
            episodes=episodes
        )

        # Phase 3: Generate and Execute Pipeline
        print("\n💻 PHASE 3: Pipeline Generation and Execution")
        generated_scripts = []
        pipeline_context = {}  # Initialize the data pipeline context

        # If it's a learning-augmented algorithm, generate a complete solution package
        objective = self.context.objective.lower()
        if "ski rental" in objective or "one-max search" in objective:
            print("  - Generating a complete LAA Solution Package (Algorithm + Predictor)...")
            algorithm_script = self.code_generator.generate(final_params, 0, self.context.objective, trust_parameter)
            predictor_script = self.code_generator._generate_predictor_template()

            # Save the package to disk
            with open("generated_algorithm.py", "w") as f:
                f.write(algorithm_script)
            with open("generated_predictor.py", "w") as f:
                f.write(predictor_script)

            generated_scripts = [algorithm_script, predictor_script]
            print("  - Solution Package saved to generated_algorithm.py and generated_predictor.py")
        else:
            for i in range(3): # Generate 3 layers for a standard pipeline
                script = self.code_generator.generate(final_params, i, self.context.objective, trust_parameter)
                generated_scripts.append(script)
                print(f"  - Generated Layer {i} with Temp={final_params.temperature:.2f}")

        if enable_federated_optimization:
            client_data_sizes = [random.randint(100, 1000) for _ in range(self.federated_optimizer.num_clients)]
            final_params = self.federated_optimizer.optimize(final_params, client_data_sizes)

        # Phase 4: Verification and Feedback
        print("\n🔬 PHASE 4: Verification and Feedback")

        # If it's a learning-augmented algorithm, use the new evaluation method
        objective = self.context.objective.lower()
        problem_context = None
        if "ski rental" in objective:
            problem_context = SkiRentalContext()
        elif "one-max search" in objective:
            problem_context = OneMaxSearchContext()

        if problem_context:
            # The new verifier takes the paths to the generated scripts
            laa_metrics = self.verifier.evaluate_learning_augmented_algorithm(
                algorithm_script_path="generated_algorithm.py",
                predictor_script_path="generated_predictor.py",
                trust_parameter=trust_parameter,
                problem_context=problem_context
            )
            print(f"  - Co-Designed Solution Competitive Ratio: {laa_metrics.get('competitive_ratio', 0.0):.2f}")
            success_metrics = laa_metrics
        else:
            # Fallback to the original pipeline verification
            metrics = self.verifier.verify_pipeline(generated_scripts)
            final_quality = metrics['overall_quality']
            print(f"  - Final Aggregated Quality: {final_quality:.2f}")
            success_metrics = {'overall_quality': final_quality}


        # Record task for meta-learning
        task_metadata = TaskMetadata(
            task_id=f"{project_type}-{datetime.now().isoformat()}",
            project_type=project_type,
            complexity_level=complexity,
            domain=domain,
            success_metrics=success_metrics,
            optimal_config=strategy,
            timestamp=datetime.now().isoformat()
        )
        self.meta_learner.record_task(task_metadata)

        return generated_scripts

# ============================================================================
# COMMAND-LINE INTERFACE
# ============================================================================

def main():
    """Command-line entry point for the Crispo system.

    This function parses command-line arguments, initializes the
    OrchestrationContext and MetaLearner (loading from a file if specified),
    creates the main Crispo engine, runs the orchestration process,
    and saves the MetaLearner's state if requested.
    """
    parser = argparse.ArgumentParser(description="Crispo: Autonomous Co-Design of ML Predictors and Algorithms")
    parser.add_argument("--project", type=str, default="AutoCode_Genesis", help="Project name.")
    parser.add_argument("--objective", type=str, default="Generate a self-optimizing multi-layer data processing script", help="The main objective.")
    parser.add_argument("--project_type", type=str, default="data_pipeline", help="Type of the project (e.g., 'data_pipeline', 'web_scraper').")
    parser.add_argument("--domain", type=str, default="data_engineering", help="Domain of the project.")
    parser.add_argument("--complexity", type=float, default=0.8, help="Complexity of the project (0.0 to 1.0).")
    parser.add_argument("--load-metaknowledge", type=str, help="Path to load MetaLearner state from.")
    parser.add_argument("--save-metaknowledge", type=str, help="Path to save MetaLearner state to.")
    parser.add_argument("--enable-transfer-learning", action="store_true", help="Enable Transfer Learning.")
    parser.add_argument("--enable-nas", action="store_true", help="Enable Neural Architecture Search.")
    parser.add_argument("--enable-federated-optimization", action="store_true", help="Enable Federated Optimization.")
    parser.add_argument("--trust-parameter", type=float, default=0.8, help="Trust parameter (lambda) for the learning-augmented algorithm.")
    parser.add_argument("--query-registry", type=str, help="Query the solution registry. E.g., 'competitive_ratio:1.5'")

    args = parser.parse_args()

    # Handle registry query if requested and exit immediately.
    if args.query_registry:
        try:
            metric, threshold_str = args.query_registry.split(':')
            threshold = float(threshold_str)
            results = query_registry(metric, threshold)

            print("\n" + "="*70)
            print("SOLUTION REGISTRY QUERY RESULTS")
            print(f"Solutions with {metric} <= {threshold}:")
            print("="*70)
            if not results:
                print("No solutions found matching the criterion.")
            else:
                for res in results:
                    print(f"  - Problem: {res['problem_type']}, Version: {res['version']}, Metrics: {res['metrics']}")
            print("="*70)
            return  # Exit after querying
        except ValueError:
            print("Invalid query format. Please use 'metric:value', e.g., 'competitive_ratio:1.5'")
            return

    context = OrchestrationContext(
        project=args.project,
        objective=args.objective
    )

    # Initialize or load the MetaLearner
    if args.load_metaknowledge:
        try:
            with open(args.load_metaknowledge, 'rb') as f:
                meta_learner = pickle.load(f)
            print(f"🧠 Meta-knowledge loaded from {args.load_metaknowledge}")
        except FileNotFoundError:
            print(f"⚠️  Warning: Meta-knowledge file not found at '{args.load_metaknowledge}'. Starting fresh.")
            meta_learner = MetaLearner()
        except (pickle.UnpicklingError, EOFError):
            print(f"⚠️  Warning: Could not unpickle meta-knowledge from '{args.load_metaknowledge}'. File may be corrupted. Starting fresh.")
            meta_learner = MetaLearner()
    else:
        meta_learner = MetaLearner()

    crispo = Crispo(context, meta_learner)

    # Load the RLAgent's Q-table from the MetaLearner
    if hasattr(meta_learner, 'rl_q_table'):
        crispo.rl_agent.load_q_table(meta_learner.rl_q_table)

    final_scripts = crispo.orchestrate(
        project_type=args.project_type,
        domain=args.domain,
        complexity=args.complexity,
        enable_transfer_learning=args.enable_transfer_learning,
        enable_nas=args.enable_nas,
        enable_federated_optimization=args.enable_federated_optimization,
        trust_parameter=args.trust_parameter
    )

    # Save the MetaLearner's state if requested
    if args.save_metaknowledge:
        # First, update the meta_learner with the latest Q-table from the RLAgent
        meta_learner.rl_q_table = crispo.rl_agent.get_q_table()
        with open(args.save_metaknowledge, 'wb') as f:
            pickle.dump(meta_learner, f)
        print(f"🧠 Meta-knowledge saved to {args.save_metaknowledge}")

    print("\n" + "="*70)
    print("ORCHESTRATION COMPLETE")
    print("="*70)

    for i, script in enumerate(final_scripts):
        print(f"\n--- Layer {i} ---\n")
        print(script)

if __name__ == "__main__":
    main()
