import sys
import os
import json
import time
import psutil
import logging
from datetime import datetime

# Adjust the system path to allow importing from the app/ directory
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))

# Import the chunking function
from app.services.parser.smart_chunker import generate_parent_child_chunks

# Setup basic logging for error capturing
logger = logging.getLogger("benchmark_chunking")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

# Output Directories
BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))
INPUT_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-parsing/Output-Parsing')
OUTPUT_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-chunking/Output-Chunking')
MATRIX_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-chunking/Perform-Matrix')
ERROR_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-chunking/Log-Error')

MATRIX_FILE = os.path.join(MATRIX_DIR, 'performance_matrix_chunking.md')
ERROR_FILE = os.path.join(ERROR_DIR, 'chunking_errors.log')

# Setup File Handler for Errors
file_handler = logging.FileHandler(ERROR_FILE, mode='a')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

def get_process_metrics(process):
    try:
        mem_info = process.memory_info()
        rss_mb = mem_info.rss / (1024 * 1024)
        cpu_percent = process.cpu_percent(interval=None)
        return rss_mb, cpu_percent
    except psutil.NoSuchProcess:
        return 0.0, 0.0

def sanitize_filename(filename: str) -> str:
    # Just grab the base name before extension for simplicity
    return os.path.splitext(filename)[0]

def log_performance_matrix(metrics: dict):
    file_exists = os.path.isfile(MATRIX_FILE)
    
    with open(MATRIX_FILE, 'a', encoding='utf-8') as f:
        if not file_exists:
            f.write("| TIMESTAMP | DOCUMENT | TOTAL PARENT CHUNK | TOTAL CHILD CHUNK | P-CHUNK AVERAGE SIZE (CHAR) | C-CHUNK AVERAGE SIZE (CHAR) | P-HEADER RATE | RAM | CPU | TIME (s) | STATUS | NOTE |\n")
            f.write("| :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :--- | :--- |\n")
        
        row = (
            f"| {metrics['timestamp']} "
            f"| {metrics['document'][:30]}... "
            f"| {metrics.get('p_chunks', 'N/A')} "
            f"| {metrics.get('c_chunks', 'N/A')} "
            f"| {metrics.get('p_avg_size', 'N/A')} "
            f"| {metrics.get('c_avg_size', 'N/A')} "
            f"| {metrics.get('p_header_rate', 'N/A')} "
            f"| {metrics.get('peak_ram', 'N/A')} MB "
            f"| {metrics.get('peak_cpu', 'N/A')} % "
            f"| {metrics.get('total_time', 'N/A')} "
            f"| {metrics['status']} "
            f"| {metrics.get('note', '-')} |\n"
        )
        f.write(row)

def calculate_metrics(p_chunks, c_chunks):
    p_count = len(p_chunks)
    c_count = len(c_chunks)
    
    p_total_chars = sum(len(str(chunk.get('content', ''))) for chunk in p_chunks)
    c_total_chars = sum(len(str(chunk.get('content', ''))) for chunk in c_chunks)
    
    p_avg_size = round(p_total_chars / p_count) if p_count > 0 else 0
    c_avg_size = round(c_total_chars / c_count) if c_count > 0 else 0
    
    p_with_header = sum(1 for chunk in p_chunks if chunk.get('metadata_header') and str(chunk.get('metadata_header')).strip() != "")
    p_header_rate = round((p_with_header / p_count) * 100, 1) if p_count > 0 else 0.0
    
    return {
        'p_chunks': p_count,
        'c_chunks': c_count,
        'p_avg_size': p_avg_size,
        'c_avg_size': c_avg_size,
        'p_header_rate': f"{p_header_rate}%"
    }

def process_file(file_path):
    filename = os.path.basename(file_path)
    base_name = sanitize_filename(filename)
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    metrics = {
        'timestamp': timestamp,
        'document': filename,
        'status': 'Failed',
        'note': '-'
    }
    
    process = psutil.Process(os.getpid())
    process.cpu_percent(interval=None) # Initialize CPU counter
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            markdown_content = f.read()
            
        start_time = time.time()
        start_ram, _ = get_process_metrics(process)
        
        # Target function execution
        parents, children = generate_parent_child_chunks(
            markdown_text=markdown_content,
            file_hash="benchmark_hash",
            title=base_name
        )
        
        end_time = time.time()
        end_ram, peak_cpu = get_process_metrics(process)
        
        peak_ram = max(start_ram, end_ram)
        total_time = round(end_time - start_time, 2)
        
        # Calculate semantic metrics
        chunk_metrics = calculate_metrics(parents, children)
        metrics.update(chunk_metrics)
        
        # Save output
        output_data = {
            "metadata": {
                "source_file": filename,
                "timestamp": timestamp,
                "metrics": chunk_metrics
            },
            "parent_chunks": parents,
            "child_chunks": children
        }
        
        out_file = os.path.join(OUTPUT_DIR, f"{base_name}_chunked.json")
        with open(out_file, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, ensure_ascii=False, indent=2)
            
        metrics.update({
            'peak_ram': round(peak_ram, 1),
            'peak_cpu': round(peak_cpu, 1),
            'total_time': total_time,
            'status': 'Success'
        })
        
    except Exception as e:
        error_msg = str(e)
        logger.error(f"Failed processing {filename}: {error_msg}", exc_info=True)
        metrics['note'] = "Exception logged"
    
    log_performance_matrix(metrics)
    print(f"Processed: {filename} -> Status: {metrics['status']}")

if __name__ == "__main__":
    print("Starting Chunking Benchmark...")
    
    # Process only specific files to control execution as per plan
    target_files = [
        "69467_284308_1_PB_parsed.md",
        "Reviving_Life_on_the_Edge_Joint_Score_Based_Graph_Generation_of_Rich_Edge_Attributes_parsed.md",
        "SemanticGraph2Vec_Semantic_graph_embedding_for_text_representation_parsed.md",
        "Sistem_Budidaya_Ikan_Lele_Clarias_sp_dalam_Kolam_Terpal_dengan_parsed.md",
        "The_Hybrid_Multimodal_Graph_Index_HMGI_A_ComprehensiveFramework_for_Integrated_Relational_and_Vector_Search_parsed.md",
        "Retrieval_Augmented_Generation_for_Knowledge_Intensive_NLP_Tasks_parsed.md",
        "PaperQA_Retrieval_Augmented_Generative_Agent_for_Scientific_Research_parsed.md",
        "Proses_Pembuatan_Pupuk_Kompos_dan_Penggunaan_Pupuk_Kompos_di_TPS_3R_parsed.md"

    ]
    
    for filename in target_files:
        file_path = os.path.join(INPUT_DIR, filename)
        if os.path.exists(file_path):
            process_file(file_path)
        else:
            print(f"File not found: {file_path}")
            
    print("Benchmark completed. Check Performance Matrix for results.")
