#!/usr/bin/env python3
"""
FSI_Edge Production Training Orchestrator
===========================================

Multi-stage training: Pretrain → SFT → GRPO RL
with distributed training, checkpointing, logging, and model export.
"""
import os
import sys
import json
import yaml
import time
import math
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import signal
import logging
from datetime import datetime
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.model import FSIEdgeModel, FSIEdgeConfig
from src.data import CodeDataset, collate_fn
from training.train import (
    train_stage1, train_stage1b_fim, train_stage2_sft,
    train_stage2b_cold_start, train_stage3_mcpo,
    rejection_sampling, train_stage4_dpo, train_stage5_long_context,
    DataFilterPipeline, ColdStartGenerator, execution_reward,
    TrainConfig
)
from export.export_gguf import convert_pytorch_to_gguf
from export.upload_hf import upload_to_huggingface


logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
    handlers=[logging.StreamHandler()]
)
log = logging.getLogger('fsi_edge')


# ============================================================================
# Production Configuration
# ============================================================================

PRODUCTION_PRESETS = {
    "quick": {
        "description": "Quick proof-of-concept (2B tokens)",
        "max_steps": 5000,
        "batch_size": 4,
        "ctx_start": 1024,
        "ctx_mid": 2048,
        "ctx_max": 4096,
        "sft_steps": 500,
        "rl_steps": 200,
        "lr": 3e-4,
        "grad_accum": 4,
        "data_samples": 10000,
    },
    "dev": {
        "description": "Development run (50B tokens)",
        "max_steps": 50000,
        "batch_size": 8,
        "ctx_start": 2048,
        "ctx_mid": 4096,
        "ctx_max": 8192,
        "sft_steps": 5000,
        "rl_steps": 2000,
        "lr": 3e-4,
        "grad_accum": 8,
        "data_samples": 100000,
    },
    "production": {
        "description": "Full production training (4T tokens, 80% target)",
        "max_steps": 500000,
        "batch_size": 32,
        "ctx_start": 4096,
        "ctx_mid": 8192,
        "ctx_max": 16384,
        "sft_steps": 50000,
        "rl_steps": 20000,
        "lr": 3e-4,
        "grad_accum": 16,
        "data_samples": 5000000,
    },
    "production_fast": {
        "description": "Optimized production (faster convergence)",
        "max_steps": 200000,
        "batch_size": 64,
        "ctx_start": 4096,
        "ctx_mid": 8192,
        "ctx_max": 16384,
        "sft_steps": 25000,
        "rl_steps": 10000,
        "lr": 5e-4,
        "grad_accum": 8,
        "data_samples": 2000000,
    },
}


# ============================================================================
# Training Pipeline Manager
# ============================================================================

class TrainingPipeline:
    """Manages the full multi-stage training lifecycle."""
    
    def __init__(self, config: TrainConfig):
        self.config = config
        self.model = None
        self.start_time = None
        self.best_accuracy = 0.0
        self.stage_results = {}
        
        os.makedirs(config.output_dir, exist_ok=True)
        self._setup_logging()
        
        # Save config
        config_path = os.path.join(config.output_dir, 'train_config.yaml')
        with open(config_path, 'w') as f:
            yaml.dump(asdict(config), f)
    
    def _setup_logging(self):
        log_path = os.path.join(self.config.output_dir, 'training.log')
        fh = logging.FileHandler(log_path)
        fh.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s')
        fh.setFormatter(formatter)
        log.addHandler(fh)
        log.info(f"Training pipeline initialized: {self.config.output_dir}")
    
    def build_model(self):
        """Build or resume model."""
        from training.train import get_model_config
        
        model_config = get_model_config(self.config.model_size)
        
        model = FSIEdgeModel(model_config)
        
        if self.config.resume_from:
            state = torch.load(self.config.resume_from, map_location='cpu')
            if 'model_state_dict' in state:
                model.load_state_dict(state['model_state_dict'])
            elif 'module' in state:
                model.load_state_dict(state['module'])
            else:
                model.load_state_dict(state)
            log.info(f"Resumed from {self.config.resume_from}")
        
        n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        log.info(f"Model: {self.config.model_size} | {n_params/1e6:.1f}M params")
        
        model.to(self.config.device)
        self.model = model
        return model
    
    def prepare_data(self):
        """Ensure training data exists, generate if needed."""
        data_path = self.config.data_path
        if not os.path.exists(data_path):
            log.info(f"Generating training data: {data_path}")
            from data.prepare_data import generate_multi_lang
            generate_multi_lang(data_path, self.config.data_samples)
        else:
            log.info(f"Data path exists: {data_path}")
    
    def _make_loader(self, max_length=None, batch_size=None):
        """Create a DataLoader with common settings."""
        max_length = max_length or self.config.ctx_max
        bs = batch_size or self.config.batch_size
        dataset = CodeDataset(
            self.config.data_path, self.config.tokenizer_path,
            max_length=max_length)
        return torch.utils.data.DataLoader(
            dataset, batch_size=bs, shuffle=True,
            num_workers=self.config.num_workers, collate_fn=collate_fn,
            pin_memory=(self.config.device == 'cuda'))
    
    def run_stage0_filter(self):
        """Stage 0: Data curation & 4-stage filtering."""
        log.info("=" * 60)
        log.info("STAGE 0: DATA CURATION & 4-STAGE FILTERING")
        log.info("=" * 60)
        filter_pipeline = DataFilterPipeline(
            quality_threshold=self.config.quality_threshold)
        log.info("Data filtering pipeline ready")
        self.stage_results['stage0'] = {'status': 'configured'}
    
    def run_stage1_pretrain(self):
        """Stage 1: Pretraining."""
        log.info("=" * 60)
        log.info("STAGE 1: PRETRAINING")
        log.info("=" * 60)
        
        loader = self._make_loader()
        self.model = train_stage1(self.model, loader, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage1_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 1 complete: {path}")
        self.stage_results['pretrain'] = {'checkpoint': path}
    
    def run_stage1b_fim(self):
        """Stage 1b: Code specialization with FIM."""
        log.info("=" * 60)
        log.info("STAGE 1b: CODE SPECIALIZATION (FIM)")
        log.info("=" * 60)
        
        loader = self._make_loader()
        self.model = train_stage1b_fim(self.model, loader, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage1b_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 1b complete: {path}")
        self.stage_results['fim'] = {'checkpoint': path}
    
    def run_stage2_sft(self):
        """Stage 2: Supervised Fine-Tuning."""
        log.info("=" * 60)
        log.info("STAGE 2: SFT")
        log.info("=" * 60)
        
        loader = self._make_loader()
        self.model = train_stage2_sft(self.model, loader, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage2_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 2 complete: {path}")
        self.stage_results['sft'] = {'checkpoint': path}
    
    def run_stage2b_cold_start(self):
        """Stage 2b: Cold-start reasoning SFT."""
        log.info("=" * 60)
        log.info("STAGE 2b: COLD-START REASONING SFT")
        log.info("=" * 60)
        
        generator = ColdStartGenerator(
            num_examples=self.config.cold_start_examples)
        cold_data = generator.generate()
        self.model = train_stage2b_cold_start(
            self.model, cold_data, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage2b_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 2b complete: {path}")
        self.stage_results['cold_start'] = {'checkpoint': path}
    
    def run_stage3_mcpo(self):
        """Stage 3: MCPO Reinforcement Learning."""
        log.info("=" * 60)
        log.info("STAGE 3: MCPO REINFORCEMENT LEARNING")
        log.info("=" * 60)
        
        self.model = train_stage3_mcpo(
            self.model, execution_reward, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage3_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 3 complete: {path}")
        self.stage_results['mcpo'] = {'checkpoint': path}
    
    def run_stage3b_rejection(self):
        """Stage 3b: Rejection sampling + second SFT round."""
        log.info("=" * 60)
        log.info("STAGE 3b: REJECTION SAMPLING")
        log.info("=" * 60)
        
        rejection_data = rejection_sampling(
            self.model, execution_reward,
            num_samples=self.config.reject_samples, config=self.config)
        
        if len(rejection_data) > 0:
            log.info(f"Second SFT round on {len(rejection_data)} examples")
            from training.train import ColdStartDataset
            rejection_dataset = ColdStartDataset(rejection_data)
            rejection_loader = torch.utils.data.DataLoader(
                rejection_dataset, batch_size=self.config.batch_size,
                shuffle=True)
            self.model = train_stage2_sft(
                self.model, rejection_loader, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage3b_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 3b complete: {path}")
        self.stage_results['rejection'] = {'checkpoint': path}
    
    def run_stage4_dpo(self):
        """Stage 4: DPO Preference Alignment."""
        log.info("=" * 60)
        log.info("STAGE 4: DPO PREFERENCE ALIGNMENT")
        log.info("=" * 60)
        
        from training.train import ColdStartDataset
        ref_model = FSIEdgeModel(self.model.config)
        ref_model.load_state_dict(self.model.state_dict())
        ref_model.to(self.config.device)
        ref_model.eval()
        
        class PrefDataset:
            def __init__(self):
                self.size = 10000
            def __iter__(self):
                return self
            def __next__(self):
                return {
                    'chosen': torch.randint(0, 1000, (1, 128)),
                    'rejected': torch.randint(0, 1000, (1, 128)),
                }
        
        self.model = train_stage4_dpo(
            self.model, ref_model, PrefDataset(), self.config)
        
        path = os.path.join(self.config.output_dir, 'stage4_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 4 complete: {path}")
        self.stage_results['dpo'] = {'checkpoint': path}
    
    def run_stage5_long_context(self):
        """Stage 5: Long-context adaptation."""
        log.info("=" * 60)
        log.info("STAGE 5: LONG-CONTEXT ADAPTATION")
        log.info("=" * 60)
        
        loader = self._make_loader(max_length=131072, batch_size=1)
        self.model = train_stage5_long_context(
            self.model, loader, self.config)
        
        path = os.path.join(self.config.output_dir, 'stage5_complete.pt')
        torch.save(self.model.state_dict(), path)
        log.info(f"Stage 5 complete: {path}")
        self.stage_results['long_ctx'] = {'checkpoint': path}
    
    def evaluate(self):
        """Run evaluation on HumanEval."""
        log.info("Evaluating on HumanEval...")
        try:
            from eval.evaluate import evaluate_model
            accuracy, results = evaluate_model(
                os.path.join(self.config.output_dir, 'final_model.pt'),
                self.config.model_size,
                self.config.device)
            
            log.info(f"HumanEval Accuracy: {accuracy:.1f}%")
            
            results_path = os.path.join(self.config.output_dir, 'eval_results.json')
            with open(results_path, 'w') as f:
                json.dump({'accuracy': accuracy, 'results': results}, f, indent=2)
            
            if accuracy > self.best_accuracy:
                self.best_accuracy = accuracy
                best_path = os.path.join(self.config.output_dir, 'best_model.pt')
                torch.save(self.model.state_dict(), best_path)
                log.info(f"New best model: {accuracy:.1f}%")
            
            return accuracy
        except Exception as e:
            log.warning(f"Evaluation failed: {e}")
            return 0.0
    
    def export(self):
        """Export to GGUF for production deployment."""
        log.info("Exporting to GGUF...")
        
        checkpoint = os.path.join(self.config.output_dir, 'final_model.pt')
        if not os.path.exists(checkpoint):
            log.warning(f"No final checkpoint at {checkpoint}, using best")
            checkpoint = os.path.join(self.config.output_dir, 'best_model.pt')
        
        for quant in ['q4_0', 'q8_0', 'f16']:
            output = os.path.join(self.config.output_dir, f'fsi_edge-{self.config.model_size}-{quant}.gguf')
            try:
                convert_pytorch_to_gguf(checkpoint, self.config.model_size, output, quant)
                log.info(f"Exported: {output}")
            except Exception as e:
                log.error(f"Export {quant} failed: {e}")
        
        log.info("Export complete.")
    
    def upload(self, repo_id=None, token=None):
        """Upload to HuggingFace Hub."""
        log.info("Uploading to HuggingFace...")
        
        checkpoint = os.path.join(self.config.output_dir, 'final_model.pt')
        if not os.path.exists(checkpoint):
            checkpoint = os.path.join(self.config.output_dir, 'best_model.pt')
        
        repo_id = repo_id or f"fsi_edge/fsi_edge-{self.config.model_size}"
        token = token or os.environ.get('HF_TOKEN')
        
        upload_to_huggingface(checkpoint, repo_id, self.config.model_size, token)
    
    def _run_stage(self, stage_name, stage_method):
        """Run a single stage with checkpointing and error handling."""
        stage_key = stage_name.replace(' ', '_')
        checkpoint = os.path.join(
            self.config.output_dir, f'{stage_key}_complete.pt')
        
        if os.path.exists(checkpoint) and self.config.resume_from:
            log.info(f"Stage {stage_name} already complete, skipping")
            return
        
        log.info(f"\n{'='*60}")
        log.info(f"Starting stage: {stage_name}")
        log.info(f"{'='*60}")
        stage_method()
    
    def run_full(self, stages='all'):
        """Run the full training pipeline with all 9 stages."""
        self.start_time = time.time()
        log.info(f"FSI_Edge Production Pipeline - {self.config.model_size}")
        log.info(f"Preset: {getattr(self.config, 'preset', 'custom')}")
        
        # Parse stages
        if stages == 'all':
            stage_list = [
                'stage0', 'stage1', 'stage1b', 'stage2',
                'stage2b', 'stage3', 'stage3b', 'stage4', 'stage5',
            ]
        else:
            stage_list = [s.strip() for s in stages.split(',')]
        
        log.info(f"Stages to run: {', '.join(stage_list)}")
        
        self.prepare_data()
        self.build_model()
        
        # Stage routing
        stage_router = {
            'stage0': self.run_stage0_filter,
            'stage1': self.run_stage1_pretrain,
            'stage1b': self.run_stage1b_fim,
            'stage2': self.run_stage2_sft,
            'stage2b': self.run_stage2b_cold_start,
            'stage3': self.run_stage3_mcpo,
            'stage3b': self.run_stage3b_rejection,
            'stage4': self.run_stage4_dpo,
            'stage5': self.run_stage5_long_context,
        }
        
        for stage_key in stage_list:
            stage_method = stage_router.get(stage_key)
            if stage_method:
                self._run_stage(stage_key, stage_method)
            else:
                log.warning(f"Unknown stage: {stage_key}")
        
        # Save final model
        final_path = os.path.join(self.config.output_dir, 'final_model.pt')
        torch.save(self.model.state_dict(), final_path)
        log.info(f"Final model saved: {final_path}")
        
        # Evaluate
        accuracy = self.evaluate()
        
        # Export
        self.export()
        
        elapsed = time.time() - self.start_time
        log.info(f"Pipeline complete in {elapsed/3600:.1f}h")
        log.info(f"Final accuracy: {accuracy:.1f}%")
        
        return accuracy


# ============================================================================
# Main entry point
# ============================================================================

def main():
    import argparse
    parser = argparse.ArgumentParser(description='FSI_Edge Production Training')
    parser.add_argument('--preset', type=str, default='dev',
                       choices=list(PRODUCTION_PRESETS.keys()))
    parser.add_argument('--model-size', type=str, default='360M',
                       choices=['360M', '800M', '1.5B'])
    parser.add_argument('--output-dir', type=str, default='/FSI_Edge/output')
    parser.add_argument('--data-path', type=str, default='/FSI_Edge/data/train')
    parser.add_argument('--tokenizer-path', type=str, default='/FSI_Edge/fsi_edge_tokenizer')
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
    parser.add_argument('--stages', type=str, default='all',
                       help='Stages to run: all, or comma-separated: stage0,stage1,stage1b,stage2,stage2b,stage3,stage3b,stage4,stage5')
    parser.add_argument('--resume', type=str, default=None)
    parser.add_argument('--batch-size', type=int, default=None)
    parser.add_argument('--max-steps', type=int, default=None)
    parser.add_argument('--lr', type=float, default=None)
    parser.add_argument('--fp16', action='store_true', default=True)
    parser.add_argument('--no-wandb', action='store_true')
    parser.add_argument('--upload', action='store_true', help='Upload to HuggingFace after training')
    parser.add_argument('--repo-id', type=str, default=None)
    parser.add_argument('--hf-token', type=str, default=None)
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    # Build config from preset + overrides
    preset = PRODUCTION_PRESETS[args.preset]
    
    config = TrainConfig(
        model_size=args.model_size,
        data_path=args.data_path,
        tokenizer_path=args.tokenizer_path,
        output_dir=f"{args.output_dir}/{args.model_size}_{args.preset}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        device=args.device,
        fp16=args.fp16 and args.device == 'cuda',
        use_wandb=not args.no_wandb,
        resume_from=args.resume,
        seed=args.seed,
        preset=args.preset,
        stages_to_run=args.stages,
        **{k: v for k, v in preset.items() if k != 'description'},
    )
    
    if args.batch_size:
        config.batch_size = args.batch_size
    if args.max_steps:
        config.max_steps = args.max_steps
    if args.lr:
        config.lr = args.lr
    
    if config.use_wandb:
        import wandb
        wandb.init(
            project='fsi_edge',
            name=f"{args.model_size}_{args.preset}",
            config=config.__dict__)
    
    # Run
    pipeline = TrainingPipeline(config)
    accuracy = pipeline.run_full(args.stages)
    
    if args.upload:
        pipeline.upload(args.repo_id, args.hf_token or os.environ.get('HF_TOKEN'))
    
    print(f"\n{'='*60}")
    print(f"FSI_Edge Training Complete!")
    print(f"  Model: {args.model_size}")
    print(f"  Preset: {args.preset}")
    print(f"  Accuracy: {accuracy:.1f}%")
    print(f"  Output: {config.output_dir}")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()