# the code is structured to automate the process of extracting text from PDF documents, 
# generating question-answer pairs from that text, 
# fine-tuning a language model based on the generated data, and evaluating the model's performance.
# The use of caching, error handling, and logging improves the robustness and maintainability of the script, making it suitable for practical applications in natural language processing and machine learning.

import json
import re
import PyPDF2
import time
from openai import OpenAI
import fitz  # PyMuPDF
import pytesseract
from PIL import Image
import io
import os
from pathlib import Path
# We have to check which Torch version for Xformers (2.3 -> 0.0.27)
from torch import __version__; from packaging.version import Version as V
from unsloth import FastLanguageModel
import torch
import json
from datasets import load_dataset
from unsloth.chat_templates import get_chat_template
from trl import SFTTrainer
from transformers import TrainingArguments, TrainerCallback
from unsloth import is_bfloat16_supported
import gc
import logging
import sys
import subprocess
import requests
# Set the GLOO_SOCKET_IFNAME environment variable
# os.environ["GLOO_SOCKET_IFNAME"] = "lo"

# Configure logging
# logging.basicConfig(level=logging.INFO)
# input_data = json.loads(sys.argv[1])  

# # Process the data (example: just print it here)
# response = f"Received data: {input_data}"
def wait_for_server(max_attempts=60):
    """Wait for the vLLM server to become available."""
    url = "http://localhost:8000/health"
    for attempt in range(max_attempts):
        try:
            response = requests.get(url)
            if response.status_code == 200:
                logging.info("vLLM server is ready!")
                return True
        except requests.exceptions.RequestException as e:
            logging.info(f"Server not ready yet: {e}. Retrying in {2**attempt} seconds...")
            time.sleep(2**attempt)


def log_output(pipe, log_func):
    """Helper function to log output from a subprocess pipe."""
    for line in iter(pipe.readline, ''):
        log_func(line.strip())

def start_vllm_server(model_name):
    cmd = [
        "vllm",
        "serve",
        "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
        "--gpu_memory_utilization=0.99",
        "--max_model_len=8192",
        "--enable-chunked-prefill=False",
        "--num_scheduler_steps=2"
    ]

    logging.info(f"Starting vLLM server with command: {' '.join(cmd)}")

    # Start the server subprocess
    server_process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1
    )

    # # Use threads to handle stdout and stderr in real-time
    # threading.Thread(target=log_output, args=(server_process.stdout, logging.info), daemon=True).start()
    # threading.Thread(target=log_output, args=(server_process.stderr, logging.error), daemon=True).start()

    # Wait for the server to become ready
    if not wait_for_server():
        server_process.terminate()
        raise Exception("Server failed to start in time.")
    
    return server_process


try:
    # Validate that we have an argument
    if len(sys.argv) < 2:
        raise ValueError("No input JSON provided")
        
    # Parse the JSON input from command line argument
    input_data = json.loads(sys.argv[1])

    # Validate required fields
    required_fields = ['pdf_file', 'system_prompt', 'model_name', 'max_step', 'learning_rate', 'epochs']
    missing_fields = [field for field in required_fields if field not in input_data]
    if missing_fields:
        raise ValueError(f"Missing required fields: {', '.join(missing_fields)}")
        
    # Your existing pipeline code here
    # Access the fields from input_data dictionary:
    pdf_file = input_data['pdf_file']
    system_prompt = input_data['system_prompt']
    model_name = input_data['model_name']
    max_step = input_data['max_step']
    learning_rate = input_data['learning_rate']
    epochs = input_data['epochs']

    # Rest of your pipeline implementation...

except json.JSONDecodeError as e:
    logging.error(f"Invalid JSON input: {str(e)}")
    sys.exit(1)
except ValueError as e:
    logging.error(str(e))
    sys.exit(1)
except Exception as e:
    logging.error(f"Pipeline error: {str(e)}")
    sys.exit(1)
# Initialize the OpenAI Client with your RunPod API Key and Endpoint URL

def get_cache_filename(pdf_path):
    pdf_stat = os.stat(pdf_path)
    pdf_modified_time = pdf_stat.st_mtime
    base_name = Path(pdf_path).stem
    cache_filename = f"{base_name}_{pdf_modified_time}.txt"
    cache_dir = "pdf_cache"
    return os.path.join(cache_dir, cache_filename)

def save_text_cache(cache_path, text_pages):
    os.makedirs(os.path.dirname(cache_path), exist_ok=True)
    with open(cache_path, 'w', encoding='utf-8') as f:
        json.dump(text_pages, f, ensure_ascii=False, indent=2)

def load_text_cache(cache_path):
    try:
        with open(cache_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return None

def extract_pdf_text_by_page(pdf_path):
    cache_path = get_cache_filename(pdf_path)
    cached_text = load_text_cache(cache_path)

    if cached_text is not None:
        logging.info(f"Loading text from cache: {cache_path}")
        return cached_text

    logging.info("Cache not found, extracting text from PDF.")
    pdf_text_pages = []
    
    try:
        with open(pdf_path, "rb") as pdf_file:
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            total_pages = len(pdf_reader.pages)

            for page_num in range(total_pages):
                page = pdf_reader.pages[page_num]
                page_text = page.extract_text() or ""

                if not page_text.strip():
                    pdf_doc = fitz.open(pdf_path)
                    pdf_page = pdf_doc.load_page(page_num)
                    image_list = pdf_page.get_images(full=True)

                    if image_list:
                        for img in image_list:
                            xref = img[0]
                            base_image = pdf_doc.extract_image(xref)
                            image_bytes = base_image["image"]
                            image = Image.open(io.BytesIO(image_bytes))
                            ocr_text = pytesseract.image_to_string(image)
                            page_text += ocr_text.strip()

                    pdf_doc.close()
                pdf_text_pages.append(page_text.strip() if page_text else "")
        save_text_cache(cache_path, pdf_text_pages)
    except Exception as e:
        print(f"Error extracting PDF text: {str(e)}")

    return pdf_text_pages or []


def clean_json_response(response):
    """
    Attempt to extract a valid JSON array from a response, even if it is incomplete or contains unterminated strings.
    """
    response_str = "\n".join(response) if isinstance(response, list) else response
    response_str = response_str.replace("```json", "").replace("```", "").strip()

    # Search for a JSON array pattern using regex, even if it’s incomplete
    json_array_match = re.search(r'(\[.*?\])', response_str, re.DOTALL)
    
    # Attempt to parse as JSON and handle unterminated strings by retrying with a trimmed response
    if json_array_match:
        json_str = json_array_match.group(1)

        try:
            json_data = json.loads(json_str)
            return json_str  # Return valid JSON if successful

        except json.JSONDecodeError as e:
            # If error is due to unterminated strings or missing brackets, try trimming and re-validating
            print(f"JSON extraction failed: {e}")
            # Try a progressive trim of the response
            for i in range(len(json_str), 0, -10):
                try:
                    trimmed_json_str = json_str[:i] + "]"  # Ensure it ends with a closing bracket
                    json_data = json.loads(trimmed_json_str)
                    return trimmed_json_str
                except json.JSONDecodeError:
                    continue  # Keep trimming until valid or no more retries

    print("JSON extraction failed: Incomplete or malformed response.")
    return ""  # Return empty string if parsing fails

def generate_qa_from_chunk(text_chunk, retries=2, max_tokens=2500):
    """
    Generate QA pairs based on a text chunk, with JSON validation, chunk splitting, and retry logic.
    """
    if not text_chunk.strip():
        return ""  # Return empty if the chunk is blank

    if len(text_chunk) > max_tokens:
        # Recursively split chunk into smaller sections if it exceeds max token limit
        half = len(text_chunk) // 2
        return generate_qa_from_chunk(text_chunk[:half], retries) + \
               generate_qa_from_chunk(text_chunk[half:], retries)
    prompt = f"""You are an AI assistant tasked with generating informative question-answer pairs from text-based documents.

INPUT CONTEXT:
{text_chunk}

TASK:
Generate relevant question-answer pairs from the provided text. Each pair must:
1. Be directly based on the information in the text
2. Include a clear, specific question
3. Provide an accurate, complete response
4. Follow the exact JSON format specified below

OUTPUT FORMAT REQUIREMENTS:
1. Respond ONLY with a JSON array
2. Each object must contain exactly two fields:
- "prompt": the question
- "response": the complete answer
3. Include no text outside the JSON array
4. Follow this exact structure:

[
    {{
        "prompt": "What is the daily allowance for a Subedar on domestic travel?",
        "response": "A Subedar is entitled to a daily allowance of Rupees 600 for domestic travel. This allowance covers meals and minor incidental expenses."
    }},
    {{
        "prompt": "How much reimbursement can be claimed for travel by train for a Lieutenant Colonel?",
        "response": "A Lieutenant Colonel is entitled to AC 1st class travel by train. The full fare for AC 1st class is reimbursed, provided the journey is undertaken for official purposes and valid tickets are submitted."
    }},
    {{
        "prompt": "What is the limit for claiming hotel accommodation reimbursement for a Havildar?",
        "response": "A Havildar can claim up to Rupees 2,500 per night for hotel accommodations during official travel, subject to submission of valid receipts and adherence to the approved lodging limits."
    }}
]

Generate the QA pairs now, following the exact format shown above."""

    attempt = 0
    while attempt < retries:
        try:
            """Query the vLLM server with retries."""
            url = "http://localhost:8000/v1/chat/completions"
            headers = {"Content-Type": "application/json"}
            data = {
                "model": f"hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt}
                ]
            }
            response = requests.post(url, headers=headers, json=data, timeout=300)
            if response and response.choices:
                response_text = response.choices[0].message.content.strip()
                
                # Use improved JSON validation with trimming for incomplete responses
                json_str = clean_json_response(response_text)
                if json_str:
                    return json_str
                else:
                    print("JSON response incomplete, retrying with split chunks.")
                    half = len(text_chunk) // 2
                    return generate_qa_from_chunk(text_chunk[:half], retries) + \
                           generate_qa_from_chunk(text_chunk[half:], retries)

        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {e}")
            attempt += 1
            time.sleep(5)  # Delay before retrying

    print("All attempts failed for this chunk.")
    return ""


def process_pdf_file_by_pages(pdf_path, output_json_file, pages_per_chunk=4, overlap=1):
    pdf_text_pages = extract_pdf_text_by_page(pdf_path)
    chunks = [' '.join(pdf_text_pages[i:i + pages_per_chunk]) 
              for i in range(0, len(pdf_text_pages) - pages_per_chunk + 1, pages_per_chunk - overlap)]

    print(f"Total chunks to process: {len(chunks)}")
    all_qa_pairs = []

    for idx, chunk in enumerate(chunks):
        print(f"Processing chunk {idx + 1}...")
        try:
            qa_pairs = generate_qa_from_chunk(chunk)

            if not qa_pairs:
                print(f"No QA pairs generated for chunk {idx + 1}")
                continue

            qa_pairs_cleaned = clean_json_response(qa_pairs)
            if not qa_pairs_cleaned:
                print(f"Failed to clean JSON for chunk {idx + 1}")
                continue

            try:
                qa_pairs_json = json.loads(qa_pairs_cleaned)
                all_qa_pairs.extend(qa_pairs_json)
                print(f"Chunk {idx + 1} processed successfully.")
            except json.JSONDecodeError as e:
                print(f"JSON decoding error for chunk {idx + 1}: {e}")
                print(f"Raw cleaned response: {qa_pairs_cleaned}")

        except Exception as e:
            print(f"Error processing chunk {idx + 1}: {e}")

    if all_qa_pairs:
        with open(output_json_file, 'w', encoding='utf-8') as json_file:
            json.dump(all_qa_pairs, json_file, ensure_ascii=False, indent=4)
        print(f"QA pairs saved to {output_json_file}")
    else:
        print("No QA pairs were successfully processed.")


# pdf_path = "/home/ubuntu/Diksha/finetuning_1/Finetuning_Complete/Dockerization/TravelEnglish 1.pdf"
pdf_file=input_data["pdf_file"]


output_json_file = 'output_json.json'
process_pdf_file_by_pages(pdf_file, output_json_file)
print(f"QA pairs saved to {output_json_file}")

#Here Starts the Finetuning Process

max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.


model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)


def load_data_from_json(file_path):
    return load_dataset('json', data_files=file_path)['train']


tokenizer = get_chat_template(
    tokenizer,
    chat_template="chatml",  # Supports zephyr, chatml, mistral, llama, alpaca, vicuna, vicuna_old, unsloth
    mapping={"role": "from", "content": "value", "user": "human", "assistant": "gpt"},  # ShareGPT style
    map_eos_token=True,  # Maps <|im_end|> to </s> instead
)

system_prompt=input_data["system_prompt"]

def formatting_prompts_func(examples):
    prompts = examples['prompt']
    responses = examples['response']

    text = []

    for prompt, response in zip(prompts, responses):
        prompt = str(prompt).strip()
        response = str(response).strip()

        # Ensure punctuation for the prompt
        if not prompt.endswith(('?', '.', '!', ':')):
            prompt += '.'

        # Capitalize the response and ensure punctuation
        response = response.capitalize()
        if not response.endswith(('.', '?', '!')):
            response += '.'

        # Create conversation in dictionary format, including the system prompt
        convo = [
            {'from': 'system', 'value': system_prompt},
            {'from': 'human', 'value': prompt},
            {'from': 'gpt', 'value': response}
        ]

        # Apply tokenizer's chat template to format each conversation
        text.append(tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False))

    return {"text": text}  # Return the formatted text

# Load your JSON file

dataset = load_data_from_json(output_json_file)

# Apply the formatting function to the dataset to create the 'text' column
dataset = dataset.map(formatting_prompts_func, batched=True) # This line is crucial to add the 'text' column

# Print a sample to verify the formatting
print(dataset["text"][0])



class TrackBestModelCallback(TrainerCallback):

    def __init__(self, output_dir):
        super().__init__()
        self.best_loss = float('inf')
        self.best_step = 0
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)



    def on_log(self, args, state, control, logs=None, **kwargs):
        train_loss = logs.get("loss")
        if train_loss is not None and train_loss < self.best_loss:
            self.best_loss = train_loss
            self.best_step = state.global_step

            # Save the model to disk
            model_path = os.path.join(self.output_dir, f"best_model_step_{self.best_step}.pt")
            torch.save(kwargs['model'].state_dict(), model_path)
            print(f"New best model saved at {model_path} with loss: {self.best_loss}")

            # Remove the previous best model if it exists

            for file in os.listdir(self.output_dir):
                if file.startswith("best_model_step_") and file != f"best_model_step_{self.best_step}.pt":
                    os.remove(os.path.join(self.output_dir, file))

    def on_train_end(self, args, state, control, **kwargs):
        print("Training ended. Best model is saved on disk.")

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=4,  # Reduced to balance between speed and CPU usage
    packing=False,  # Re-enabled packing for efficiency
    args=TrainingArguments(
        per_device_train_batch_size=4,  # Reduced to improve speed
        gradient_accumulation_steps=2,  # Reduced to increase update frequency
        warmup_steps=100,  # Reduced warmup steps
        num_train_epochs=epochs,
        max_steps=max_step,
        learning_rate=learning_rate,  # Slightly reduced for stability with smaller batch size
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=50,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="cosine",
        seed=3407,
        output_dir="outputs",
        max_grad_norm=1.0,
        dataloader_num_workers=4,  # Adjusted to match dataset_num_proc
        gradient_checkpointing=True,  # Re-enabled to save memory
        adam_beta1=0.9,
        adam_beta2=0.999,
        adam_epsilon=1e-8,
        ddp_find_unused_parameters=False,
        report_to="none",  # Disable wandb logging if you're not using it
    ),

)


#@title Show current memory stats

gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")


# In your training script:

output_dir = "outputs"
best_model_tracker = TrackBestModelCallback(output_dir)
trainer.add_callback(best_model_tracker)



# Start training

trainer_stats = trainer.train()
print("Training completed.")
print(f"Best loss: {best_model_tracker.best_loss} at step {best_model_tracker.best_step}")
print("Final training loss:", trainer_stats.training_loss)


#@title Show final memory and time stats

used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory         /max_memory*100, 3)
lora_percentage = round(used_memory_for_lora/max_memory*100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")

# Saving the model for VLLM
# Function to save and push model in different formats
def save_and_push_model(format, save_method, model_name):
    try:
        model.save_pretrained_merged("model", tokenizer, save_method=save_method)
        model.push_to_hub_merged(f"PharynxAI/{model_name}", tokenizer, save_method=save_method, private=True, token=os.getenv('HF_Token'))
        print(f"Successfully saved and pushed model in {format} format.")
    except Exception as e:
        print(f"Error while saving or pushing model in {format} format: {e}")
if __name__ == "__main__":

    server_process = None
    
    try:
        # # Start vLLM server
        server_process = start_vllm_server("hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4")

    except Exception as e:
        logging.error(f"An error occurred: {e}")
        sys.exit(1)
        
    finally:
        # Cleanup: terminate the server process if it exists
        if server_process:
            logging.info("Shutting down vLLM server...")
            server_process.terminate()
            try:
                server_process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                logging.warning("Server didn't terminate gracefully, forcing kill...")
                server_process.kill()
                server_process.wait()
            logging.info("Server shutdown complete")

    # Assuming input_data is defined and contains 'model_name' 
    save_and_push_model("16bit", "merged_16bit", input_data['model_name'])
    save_and_push_model("4bit", "merged_4bit", input_data['model_name'])
    save_and_push_model("LoRA adapters", "lora", input_data['model_name'])