#!/usr/bin/env python3
import os
# below line clears the console
os.system('clear')
# below line calls gpustat if the gpu id is 0
if os.environ.get('RANK') == '0':
    os.system('gpustat')
import re
import torch
import logging
from typing import List, Dict, Any, Optional
from datasets import load_dataset, Dataset
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GRPOConfig, GRPOTrainer
from accelerate import Accelerator
import torch.distributed as dist
from premsql.executors import SQLiteExecutor
from premsql.datasets import BirdDataset
from premsql.evaluator import Text2SQLEvaluator
import pandas as pd
import json
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Constants
SYSTEM_PROMPT = """You are an expert SQL analyst. For the given question and database schema, generate a SQL query.
First analyze the schema and requirements, then write the query.

The database schema and tables are:
{schema}

Respond in the following format:

<reasoning>
[Step by step analysis of the requirements and schema]
</reasoning>
<sql>
[Your SQL query]
</sql>"""

class PremSQLv2Trainer:
    def __init__(
        self,
        model_name: str,
        output_dir: str,
        dataset_folder: str = "./data",
        hf_token: Optional[str] = None
    ):
        self.model_name = model_name
        self.output_dir = output_dir
        self.dataset_folder = Path(dataset_folder)
        self.hf_token = hf_token
        self.device_count = torch.cuda.device_count()
        self.executor = SQLiteExecutor()
        self.evaluator = Text2SQLEvaluator(
            executor=self.executor,
            experiment_path=output_dir
        )
        
        # Create dataset directory if it doesn't exist
        self.dataset_folder.mkdir(parents=True, exist_ok=True)
        
        self.setup_environment()

    def setup_environment(self):
        """Setup training environment and distributed training"""
        os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
        
        if self.device_count > 1:
            if os.environ.get('RANK') is not None and os.environ.get('WORLD_SIZE') is not None:
                if not dist.is_initialized():
                    dist.init_process_group(backend='nccl')
            else:
                logger.info("Multiple GPUs detected but distributed training environment variables not set. Running on single GPU.")

    @staticmethod
    def extract_reasoning(text: str) -> str:
        """Extract reasoning from model output"""
        try:
            reasoning = text.split("<reasoning>")[-1]
            reasoning = reasoning.split("</reasoning>")[0]
            return reasoning.strip()
        except IndexError:
            return ""

    @staticmethod
    def extract_sql_query(text: str) -> str:
        """Extract SQL query from model output"""
        try:
            sql = text.split("<sql>")[-1]
            sql = sql.split("</sql>")[0]
            return sql.strip()
        except IndexError:
            return ""

    def reasoning_quality_reward(self, text: str) -> float:
        """Evaluate the quality of reasoning"""
        reasoning = self.extract_reasoning(text)
        if not reasoning:
            return 0.0

        reward = 0.0
        
        # Check for key reasoning components
        if "schema" in reasoning.lower(): reward += 0.2
        if "table" in reasoning.lower(): reward += 0.2
        if "join" in reasoning.lower(): reward += 0.2
        if "where" in reasoning.lower(): reward += 0.2
        if "condition" in reasoning.lower(): reward += 0.2
        
        # Check for step-by-step analysis
        steps = [s for s in reasoning.split('\n') if s.strip()]
        if len(steps) >= 3:
            reward += 0.3
        
        # Check for schema understanding
        if any(word in reasoning.lower() for word in ["primary key", "foreign key", "relationship"]):
            reward += 0.2
            
        return min(reward, 1.5)  # Cap at 1.5

    def show_gpu_usage(self):
        """Show GPU usage if on main process"""
        if os.environ.get('RANK') == '0':
            os.system('gpustat')

    def get_training_dataset(self) -> Dataset:
        """Load and prepare BirdBench dataset"""
        if os.environ.get('RANK') == '0':
            logger.info("Loading BirdBench dataset...")
        
        # Initialize BirdDataset without forcing download
        bird_dataset = BirdDataset(
            split="train",
            dataset_folder=self.dataset_folder,
            hf_token=self.hf_token,
            force_download=False
        )
        
        # Format the examples with schema information
        formatted_examples = []
        for example in bird_dataset.dataset:
            # Get database schema
            db_path = f"{self.dataset_folder}/train_databases/{example['db_id']}/{example['db_id']}.sqlite"
            schema = self.get_database_schema(db_path)
            
            # Format prompt with schema
            prompt = [
                {'role': 'system', 'content': SYSTEM_PROMPT.format(schema=schema)},
                {'role': 'user', 'content': f"Database: {example['db_id']}\nQuestion: {example['question']}"}
            ]
            formatted_examples.append({
                'prompt': prompt,
                'ground_truth': example['SQL'],
                'db_id': example['db_id']
            })
        
        # Convert to Dataset
        dataset = Dataset.from_list(formatted_examples)
        
        # Add tokenized inputs
        tokenizer = AutoTokenizer.from_pretrained(
            self.model_name,
            padding_side="right",
            token=self.hf_token
        )
        
        def tokenize_function(examples):
            # Apply chat template
            texts = [tokenizer.apply_chat_template(p, tokenize=False) for p in examples['prompt']]
            tokenized = tokenizer(texts, padding=True, truncation=True)
            
            # Keep original fields
            return {
                'input_ids': tokenized['input_ids'],
                'attention_mask': tokenized['attention_mask'],
                'prompt': examples['prompt'],
                'ground_truth': examples['ground_truth'],
                'db_id': examples['db_id']
            }
        
        # Tokenize while keeping original fields
        dataset = dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=None
        )
        
        if os.environ.get('RANK') == '0':
            logger.info(f"Dataset loaded with {len(dataset)} examples")
            self.show_gpu_usage()
        
        return dataset

    def get_database_schema(self, db_path: str) -> str:
        """Extract database schema from SQLite database"""
        try:
            import sqlite3
            conn = sqlite3.connect(db_path)
            cursor = conn.cursor()
            
            # Get all tables
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
            tables = cursor.fetchall()
            
            schema_info = []
            for table in tables:
                table_name = table[0]
                # Get columns for each table
                cursor.execute(f"PRAGMA table_info({table_name});")
                columns = cursor.fetchall()
                
                # Format table schema
                column_info = [f"{col[1]} ({col[2]})" for col in columns]
                schema_info.append(f"Table: {table_name}\nColumns: {', '.join(column_info)}\n")
            
            conn.close()
            return "\n".join(schema_info)
        except Exception as e:
            if os.environ.get('RANK') == '0':
                logger.warning(f"Failed to get schema for {db_path}: {str(e)}")
            return "Schema information unavailable"

    def execution_accuracy_reward(
        self,
        prompts: List,
        completions: List,
        **kwargs
    ) -> List[float]:
        """Reward function based on SQL execution accuracy"""
        responses = [completion[0]['content'] for completion in completions]
        rewards = []

        # Get ground truth from dataset
        dataset = kwargs.get('train_dataset')
        batch_size = len(prompts)
        current_idx = kwargs.get('current_step', 0) * batch_size

        for i, response in enumerate(responses):
            query = self.extract_sql_query(response)
            reward = 0.0
            
            try:
                if query:
                    # Get corresponding ground truth
                    example = dataset[current_idx + i]
                    db_path = f"{self.dataset_folder}/train_databases/{example['db_id']}/{example['db_id']}.sqlite"
                    
                    # Debug information
                    if os.environ.get('RANK') == '0':
                        debug_info = {
                            'query': query,
                            'ground_truth': example['ground_truth'],
                            'db_path': db_path
                        }
                        print("\n=== Debug Info ===")
                        print(f"Full Response:\n{response}")
                        print(f"\nExtracted Query:\n{query}")
                        print(f"\nDB Path: {db_path}")
                        print(f"DB Path exists: {os.path.exists(db_path)}")
                        print(f"\nExample Data:")
                        print(json.dumps(example, indent=2))
                        print("\nEvaluator Input:")
                        print(json.dumps(debug_info, indent=2))
                        print("================\n")
                    
                    result = self.evaluator.execute(
                        metric_name="accuracy",
                        model_responses=[{
                            'query': query,
                            'ground_truth': example['ground_truth'],
                            'db_path': db_path
                        }],
                        filter_by="db_id",
                        meta_time_out=10
                    )
                    
                    # Debug evaluator result
                    if os.environ.get('RANK') == '0':
                        print("Evaluator Result:")
                        print(json.dumps(result, indent=2))
                    
                    # Reward based on accuracy
                    reward = result['accuracy'] * 2.0
                    
                    # Add small rewards for SQL components
                    reward += 0.1 * query.upper().count("SELECT")
                    reward += 0.1 * query.upper().count("FROM")
                    reward += 0.2 * query.upper().count("JOIN")
                    
            except Exception as e:
                if os.environ.get('RANK') == '0':
                    logger.warning(f"Query execution error: {str(e)}")
                    logger.warning(f"Error occurred with query: {query}")
                    logger.warning(f"Full response was: {response}")
                    
                if all(k in query.upper() for k in ["SELECT", "FROM"]):
                    reward = 0.1
                    
            rewards.append(reward)
            
        return rewards

    def setup_training_args(self) -> GRPOConfig:
        """Configure training arguments"""
        return GRPOConfig(
            output_dir=self.output_dir,
            run_name=f"{self.model_name.split('/')[-1]}-GRPO-v2",
            learning_rate=5e-6,
            adam_beta1=0.9,
            adam_beta2=0.99,
            weight_decay=0.1,
            warmup_ratio=0.1,
            lr_scheduler_type='cosine',
            logging_steps=1,
            bf16=True,
            per_device_train_batch_size=2,
            gradient_accumulation_steps=2,
            num_generations=16,
            max_prompt_length=512,
            max_completion_length=256,
            num_train_epochs=1,
            save_steps=50,
            max_grad_norm=0.1,
            log_on_each_node=True,
            use_vllm=False,
            report_to="wandb",  # Using wandb for tracking
            ddp_find_unused_parameters=False,
            push_to_hub=False,
            hub_model_id=None,
            hub_token=None,
        )

    def train(self):
        """Main training function"""
        try:
            # Load dataset
            dataset = self.get_training_dataset()
            if os.environ.get('RANK') == '0':
                logger.info(f"Loaded dataset with {len(dataset)} examples")

            # Initialize model and tokenizer
            model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.bfloat16,
            ).cuda()
            
            tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            tokenizer.pad_token = tokenizer.eos_token
            
            # Setup training arguments
            training_args = self.setup_training_args()
            
            class GPUMonitorCallback(transformers.TrainerCallback):
                def __init__(self, trainer):
                    self.trainer = trainer
                    
                def on_save(self, args, state, control, **kwargs):
                    if os.environ.get('RANK') == '0' and state.global_step % 50 == 0:
                        self.trainer.show_gpu_usage()
            
            # Initialize trainer
            trainer = GRPOTrainer(
                model=model,
                processing_class=tokenizer,
                reward_funcs=[self.execution_accuracy_reward],
                args=training_args,
                train_dataset=dataset,
                callbacks=[GPUMonitorCallback(self)]
            )
            
            # Start training
            trainer.train()
            
            # Save model
            if dist.get_rank() == 0:
                logger.info(f"Saving model to {self.output_dir}")
                trainer.save_model()
                self.show_gpu_usage()  # Show final GPU usage
            
            # Cleanup distributed training
            if dist.is_initialized():
                dist.destroy_process_group()
            
        except Exception as e:
            if os.environ.get('RANK') == '0':
                logger.error(f"Training failed with error: {str(e)}")
            if dist.is_initialized():
                dist.destroy_process_group()
            raise

def main():
    """Main entry point"""
    model_name = "premai-io/prem-1B-SQL"
    output_dir = "outputs/Prem-1B-SQL-GRPO-v2"
    dataset_folder = "./data/birdbench"
    
    hf_token = "hf_"
    
    trainer = PremSQLv2Trainer(
        model_name=model_name,
        output_dir=output_dir,
        dataset_folder=dataset_folder,
        hf_token=hf_token
    )
    trainer.train()

if __name__ == "__main__":
    main()