import os
import shutil
import importlib.util
from models.training import Training
from models.TrainingProject import TrainingProject

def load_base_config(selected_model):
    """Load base configuration for a specific YOLOX model"""
    model_name = selected_model.lower().replace('-', '_').replace('.pth', '')
    base_config_path = os.path.join(os.path.dirname(__file__), '..', 'data', f'{model_name}.py')
    
    if not os.path.exists(base_config_path):
        raise Exception(f'Base configuration not found for model: {model_name} at {base_config_path}')
    
    # Load the module dynamically
    spec = importlib.util.spec_from_file_location(f"base_config_{model_name}", base_config_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    
    # Extract all attributes from BaseExp class
    base_exp = module.BaseExp()
    base_config = {}
    for attr in dir(base_exp):
        if not attr.startswith('_'):
            base_config[attr] = getattr(base_exp, attr)
    
    return base_config

def generate_yolox_exp(training_id):
    """Generate YOLOX exp.py file"""
    # Fetch training row from DB
    training = Training.query.get(training_id)
    if not training:
        training = Training.query.filter_by(project_details_id=training_id).first()
    
    if not training:
        raise Exception(f'Training not found for trainingId or project_details_id: {training_id}')
    
    # If transfer_learning is 'coco', generate exp using base config + custom settings
    if training.transfer_learning == 'coco':
        exp_content = generate_yolox_inference_exp(training_id, use_base_config=True)
        return {'type': 'custom', 'expContent': exp_content}
    
    # If transfer_learning is 'sketch', generate custom exp.py
    if training.transfer_learning == 'sketch':
        exp_content = generate_yolox_inference_exp(training_id, use_base_config=False)
        return {'type': 'custom', 'expContent': exp_content}
    
    raise Exception(f'Unknown transfer_learning type: {training.transfer_learning}')

def save_yolox_exp(training_id, out_path):
    """Save YOLOX exp.py to specified path"""
    exp_result = generate_yolox_exp(training_id)
    
    if exp_result['type'] == 'custom' and 'expContent' in exp_result:
        with open(out_path, 'w') as f:
            f.write(exp_result['expContent'])
        return out_path
    elif exp_result['type'] == 'default' and 'expPath' in exp_result:
        # Optionally copy the file if outPath is different
        if exp_result['expPath'] != out_path:
            shutil.copyfile(exp_result['expPath'], out_path)
        return out_path
    else:
        raise Exception('Unknown expResult type or missing content')

def generate_yolox_inference_exp(training_id, options=None, use_base_config=False):
    """Generate inference exp.py using DB values
    
    Args:
        training_id: The training/project_details ID
        options: Optional overrides for data paths
        use_base_config: If True, load base config and only override with user-defined values
    """
    if options is None:
        options = {}
    
    training = Training.query.get(training_id)
    if not training:
        training = Training.query.filter_by(project_details_id=training_id).first()
    
    if not training:
        raise Exception(f'Training not found for trainingId or project_details_id: {training_id}')
    
    # Always use the training_id (project_details_id) for annotation file names
    project_details_id = training.project_details_id
    
    data_dir = options.get('data_dir', '/home/kitraining/To_Annotate/')
    train_ann = options.get('train_ann', f'coco_project_{training_id}_train.json')
    val_ann = options.get('val_ann', f'coco_project_{training_id}_valid.json')
    test_ann = options.get('test_ann', f'coco_project_{training_id}_test.json')
    
    # Get num_classes from TrainingProject.classes JSON
    num_classes = 80
    try:
        training_project = TrainingProject.query.get(project_details_id)
        if training_project and training_project.classes:
            classes_arr = training_project.classes
            if isinstance(classes_arr, str):
                import json
                classes_arr = json.loads(classes_arr)
            
            if isinstance(classes_arr, list):
                num_classes = len([c for c in classes_arr if c not in [None, '']])
            elif isinstance(classes_arr, dict):
                num_classes = len([k for k, v in classes_arr.items() if v not in [None, '']])
    except Exception as e:
        print(f'Could not determine num_classes from TrainingProject.classes: {e}')
    
    # Initialize config dictionary
    config = {}
    
    # If using base config (transfer learning from COCO), load protected parameters first
    if use_base_config and training.selected_model:
        try:
            base_config = load_base_config(training.selected_model)
            config.update(base_config)
            print(f'Loaded base config for {training.selected_model}: {list(base_config.keys())}')
        except Exception as e:
            print(f'Warning: Could not load base config for {training.selected_model}: {e}')
            print('Falling back to custom settings only')
    
    # Override with user-defined values from training table (only if they exist and are not None)
    user_overrides = {
        'depth': training.depth,
        'width': training.width,
        'input_size': training.input_size,
        'mosaic_scale': training.mosaic_scale,
        'test_size': training.test_size,
        'enable_mixup': training.enable_mixup,
        'max_epoch': training.max_epoch,
        'warmup_epochs': training.warmup_epochs,
        'warmup_lr': training.warmup_lr,
        'basic_lr_per_img': training.basic_lr_per_img,
        'scheduler': training.scheduler,
        'no_aug_epochs': training.no_aug_epochs,
        'min_lr_ratio': training.min_lr_ratio,
        'ema': training.ema,
        'weight_decay': training.weight_decay,
        'momentum': training.momentum,
        'print_interval': training.print_interval,
        'eval_interval': training.eval_interval,
        'test_conf': training.test_conf,
        'nms_thre': training.nms_thre,
        'mosaic_prob': training.mosaic_prob,
        'mixup_prob': training.mixup_prob,
        'hsv_prob': training.hsv_prob,
        'flip_prob': training.flip_prob,
        'degrees': training.degrees,
        'translate': training.translate,
        'shear': training.shear,
        'mixup_scale': training.mixup_scale,
        'activation': training.activation,
    }
    
    # Only override if value is explicitly set (not None)
    for key, value in user_overrides.items():
        if value is not None:
            config[key] = value
    
    # Apply any additional options overrides
    config.update(options)
    
    # Set defaults for any missing required parameters
    config.setdefault('depth', 1.00)
    config.setdefault('width', 1.00)
    config.setdefault('input_size', [640, 640])
    config.setdefault('mosaic_scale', [0.1, 2])
    config.setdefault('random_size', [10, 20])
    config.setdefault('test_size', [640, 640])
    config.setdefault('enable_mixup', False)
    config.setdefault('exp_name', 'inference_exp')
    
    # Build exp content
    exp_content = f'''#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import os

from yolox.exp import Exp as MyExp


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.data_dir = "{data_dir}"
        self.train_ann = "{train_ann}"
        self.val_ann = "{val_ann}"
        self.test_ann = "{test_ann}"
        self.num_classes = {num_classes}
'''
    
    # Set pretrained_ckpt if transfer_learning is 'coco'
    if training.transfer_learning and isinstance(training.transfer_learning, str) and training.transfer_learning.lower() == 'coco':
        yolox_base_dir = '/home/kitraining/Yolox/YOLOX-main'
        selected_model = training.selected_model.replace('.pth', '') if training.selected_model else ''
        if selected_model:
            exp_content += f"        self.pretrained_ckpt = r'{yolox_base_dir}/pretrained/{selected_model}.pth'\n"
    
    # Format arrays
    def format_value(val):
        if isinstance(val, (list, tuple)):
            return '(' + ', '.join(map(str, val)) + ')'
        elif isinstance(val, bool):
            return str(val)
        elif isinstance(val, str):
            return f'"{val}"'
        else:
            return str(val)
    
    # Add all config parameters to exp
    for key, value in config.items():
        if key not in ['exp_name']:  # exp_name is handled separately
            exp_content += f"        self.{key} = {format_value(value)}\n"
    
    # Add exp_name at the end (uses dynamic path)
    exp_content += f'''        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
'''
    
    return exp_content

def save_yolox_inference_exp(training_id, out_path, options=None):
    """Save inference exp.py to custom path"""
    exp_content = generate_yolox_inference_exp(training_id, options, use_base_config=False)
    with open(out_path, 'w') as f:
        f.write(exp_content)
    return out_path
