import sys
import os
import json
import time
import psutil
import logging
from datetime import datetime
import numpy as np

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

from app.services.parser.semantic_retriever import SemanticRetriever

# Setup basic logging for error capturing
logger = logging.getLogger("benchmark_retrieval")
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-chunking/Output-Chunking')
OUTPUT_JSON_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-retrieval/Output-JSON-C')
OUTPUT_MD_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-retrieval/Output-Human-Verification-C')
MATRIX_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-retrieval/Perform-Matrix')
ERROR_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-retrieval/Log-Error')

os.makedirs(OUTPUT_JSON_DIR, exist_ok=True)
os.makedirs(OUTPUT_MD_DIR, exist_ok=True)
os.makedirs(MATRIX_DIR, exist_ok=True)
os.makedirs(ERROR_DIR, exist_ok=True)

MATRIX_FILE = os.path.join(MATRIX_DIR, 'performance_matrix_retrieval_c.md')
ERROR_FILE = os.path.join(ERROR_DIR, 'retrieval_errors_c.log')

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

# The 5 Queries (Candidate C: The Anchor-Density Strategy)
QUERY_TEXTS = [
    "latar belakang masalah, urgensi penelitian, motivasi mendasar, dilatar belakangi oleh, fenomena krusial, konteks situasi saat ini, dan alasan utama studi ini dilakukan.",
    "rincian teknis metodologi, arsitektur sistem, desain eksperimen, teknik pengumpulan dataset, algoritma yang diusulkan, sampel, dan prosedur pengujian.",
    "keterbatasan penelitian terdahulu, kelemahan metode sebelumnya, kekosongan literatur, masalah yang belum terselesaikan, namun demikian, akan tetapi, dan hasil yang belum jelas.",
    "tujuan utama penelitian, sasaran spesifik, kontribusi inti makalah, hasil yang ingin dicapai, mengeksplorasi, dan solusi yang diusulkan untuk menyelesaikan masalah.",
    "saran untuk penelitian selanjutnya, rekomendasi pekerjaan di masa depan, arah pengembangan masa depan, dan peluang studi berikutnya untuk mengatasi batasan saat ini."
]
PILLAR_NAMES = ["Background", "Methodology", "Gap", "Objective", "Futurework"]

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 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 | CHUNKS EMBEDDED | EMBED TIME (s) | SEARCH TIME (s) | AVG SCORE (BG/MET/GAP/OBJ/FUT) | RAM | CPU | STATUS | NOTE |\n")
            f.write("| :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :--- | :--- |\n")
        
        doc_name = metrics['document']
        short_name = (doc_name[:25] + '...') if len(doc_name) > 28 else doc_name
        
        row = (
            f"| {metrics['timestamp']} "
            f"| {short_name} "
            f"| {metrics.get('chunks_embedded', 'N/A')} "
            f"| {metrics.get('embed_time', 'N/A')} "
            f"| {metrics.get('search_time', 'N/A')} "
            f"| {metrics.get('avg_score_str', 'N/A')} "
            f"| {metrics.get('peak_ram', 'N/A')} MB "
            f"| {metrics.get('peak_cpu', 'N/A')} % "
            f"| {metrics['status']} "
            f"| {metrics.get('note', '-')} |\n"
        )
        f.write(row)

class BenchmarkRetriever(SemanticRetriever):
    """Wrapper to expose cosine similarity scores for testing."""
    def find_top_k_parents_with_scores(self, queries, child_chunks, top_k=3):
        if not child_chunks or not queries:
            return {}, {}, 0.0, 0.0

        t0 = time.time()
        child_texts = [c["content"] for c in child_chunks]
        child_embeddings = self._embed_batch(child_texts)
        child_vecs = np.array(child_embeddings)
        
        query_embeddings = self._embed_batch(queries)
        query_vecs = np.array(query_embeddings)
        embed_time = time.time() - t0
        
        t1 = time.time()
        results_ids = {}
        results_scores = {}
        
        for q_idx, query in enumerate(queries):
            q_vec = query_vecs[q_idx]
            
            if np.linalg.norm(q_vec) == 0:
                 results_ids[query] = []
                 results_scores[query] = []
                 continue
                 
            scores = []
            for c_idx, c_vec in enumerate(child_vecs):
                score = self._cosine_similarity(q_vec, c_vec)
                scores.append((score, child_chunks[c_idx]["parent_id"]))
            
            scores.sort(key=lambda x: x[0], reverse=True)
            
            top_parents = []
            top_scores = []
            seen_parents = set()
            for score, p_id in scores:
                if p_id not in seen_parents:
                    top_parents.append(p_id)
                    top_scores.append(score)
                    seen_parents.add(p_id)
                if len(top_parents) >= top_k:
                    break
                    
            results_ids[query] = top_parents
            results_scores[query] = top_scores
            
        search_time = time.time() - t1
        return results_ids, results_scores, embed_time, search_time

def create_human_verification_md(doc_name, results_ids, results_scores, parent_chunks_map):
    lines = [f"# Retrieval Verification: {doc_name}", ""]
    
    for i, query in enumerate(QUERY_TEXTS):
        pillar = PILLAR_NAMES[i]
        lines.append(f"## Pilar: {pillar}")
        lines.append(f"**Query:** _{query}_")
        lines.append("")
        
        top_ids = results_ids.get(query, [])
        top_scores = results_scores.get(query, [])
        
        if not top_ids:
            lines.append("> Tidak ada hasil ditemukan.\n")
            continue
            
        for rank, p_id in enumerate(top_ids):
            score = top_scores[rank] if rank < len(top_scores) else 0.0
            parent_data = parent_chunks_map.get(p_id, {})
            
            header = parent_data.get('metadata_header', 'No Header')
            content = parent_data.get('content', 'No Content')
            
            lines.append(f"### Rank {rank+1} (Score: {score:.4f})")
            lines.append(f"**Header:** {header}")
            lines.append(f"**Content:**")
            lines.append(f"```text\n{content}\n```")
            lines.append("")
            
    return "\n".join(lines)

def process_file(file_path, retriever):
    filename = os.path.basename(file_path)
    base_name = filename.replace('_chunked.json', '')
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    metrics = {
        'timestamp': timestamp,
        'document': base_name,
        '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:
            data = json.load(f)
            
        parent_chunks = data.get('parent_chunks', [])
        child_chunks = data.get('child_chunks', [])
        
        # --- Data Purification Layer (Aggressive V2) ---
        import re
        
        EXTENDED_IRRELEVANT_HEADERS = [
            "reference", "references", "bibliography", "daftar pustaka", "referensi",
            "acknowledgement", "acknowledgements", "acknowledgment", "acknowledgments",
            "ucapan terima kasih", "funding", "pendanaan",
            "appendix", "appendices", "lampiran",
            "table of contents", "daftar isi",
            "author contributions", "kontribusi penulis"
        ]
        
        def is_header_blacklisted(header):
            if not header: return False
            normalized = header.strip().lower()
            return any(token in normalized for token in EXTENDED_IRRELEVANT_HEADERS)
            
        def clean_text_noise(text):
            if not text: return ""
            # Hapus blok Teks Gambar Hasil OCR secara ekstrem
            text = re.sub(r'> \*\*\[Teks Gambar Hasil OCR - .*?\]\*\*(?::)?.*?(?=\n\n|\Z)', '', text, flags=re.DOTALL | re.IGNORECASE)
            text = re.sub(r'\*+- Start of picture text -\*+.*?\*+- End of picture text -\*+', '', text, flags=re.DOTALL | re.IGNORECASE)
            # Hapus tabel markdown (baris dengan banyak pipe)
            text = re.sub(r'(?im)^.*\|.*\|.*$', '', text)
            # Hapus baris yang mengandung email, korespondensi, dll
            text = re.sub(r'(?im)^.*(?:e-?mail|@|korespondensi).*$', '', text)
            return text.strip()

        def is_content_garbage(text):
            if not text: return True
            lowered = text.lower()
            
            # 1. Deteksi kalimat pengakuan/afiliasi (Mencegah header lolos)
            if any(kw in lowered for kw in ["acknowledgment", "acknowledgement", "ucapan terima kasih", "funded by", "didanai oleh", "thank"]):
                return True
                
            # 2. Deteksi referensi murni
            ref_brackets = len(re.findall(r'\[\d{1,3}\]', text))
            ref_years = len(re.findall(r'\(\d{4}\)', text))
            urls = len(re.findall(r'http[s]?://|doi\s*:', text, flags=re.IGNORECASE))
            
            if (ref_brackets + ref_years + urls) > 4 or urls > 1:
                return True
                
            return False

        purified_parents = []
        for p in parent_chunks:
            if is_header_blacklisted(p.get("metadata_header")):
                continue
            
            cleaned_content = clean_text_noise(p.get("content", ""))
            
            if is_content_garbage(cleaned_content) or len(cleaned_content) < 100:
                continue
                
            p["content"] = cleaned_content
            purified_parents.append(p)
                
        valid_parent_ids = {p["parent_id"] for p in purified_parents}
        purified_children = []
        for c in child_chunks:
            if c.get("parent_id") not in valid_parent_ids:
                continue
            cleaned_content = clean_text_noise(c.get("content", ""))
            if len(cleaned_content) < 30 or is_content_garbage(cleaned_content):
                continue
            c["content"] = cleaned_content
            purified_children.append(c)

        parent_chunks = purified_parents
        child_chunks = purified_children
        # -------------------------------
        
        if not child_chunks:
            raise ValueError("No child chunks found in JSON after purification")
            
        metrics['chunks_embedded'] = len(child_chunks)
        
        parent_chunks_map = {p['parent_id']: p for p in parent_chunks}
        
        start_ram, _ = get_process_metrics(process)
        
        # Target function execution
        results_ids, results_scores, embed_time, search_time = retriever.find_top_k_parents_with_scores(
            queries=QUERY_TEXTS,
            child_chunks=child_chunks,
            top_k=3
        )
        
        end_ram, peak_cpu = get_process_metrics(process)
        peak_ram = max(start_ram, end_ram)
        
        # Calculate Avg Scores per pillar
        avg_scores = []
        for q in QUERY_TEXTS:
            scores = results_scores.get(q, [])
            avg_val = sum(scores)/len(scores) if scores else 0.0
            avg_scores.append(f"{avg_val:.2f}")
            
        avg_score_str = " / ".join(avg_scores)
        
        # Generate Outputs
        md_content = create_human_verification_md(base_name, results_ids, results_scores, parent_chunks_map)
        md_path = os.path.join(OUTPUT_MD_DIR, f"{base_name}_human_verification.md")
        with open(md_path, 'w', encoding='utf-8') as f:
            f.write(md_content)
            
        json_output = {
            "document": base_name,
            "timestamp": timestamp,
            "retrieval_results": []
        }
        for i, q in enumerate(QUERY_TEXTS):
            json_output["retrieval_results"].append({
                "pillar": PILLAR_NAMES[i],
                "query": q,
                "top_parents": results_ids.get(q, []),
                "scores": results_scores.get(q, [])
            })
            
        json_path = os.path.join(OUTPUT_JSON_DIR, f"{base_name}_retrieval_results.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(json_output, f, ensure_ascii=False, indent=2)
            
        metrics.update({
            'embed_time': round(embed_time, 2),
            'search_time': round(search_time, 2),
            'avg_score_str': avg_score_str,
            'peak_ram': round(peak_ram, 1),
            'peak_cpu': round(peak_cpu, 1),
            '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 Retrieval Benchmark...")
    print("Initializing SemanticRetriever (Loading Model)...")
    try:
        retriever = BenchmarkRetriever()
        print("Model loaded successfully.")
    except Exception as e:
        print(f"Failed to load model: {e}")
        sys.exit(1)
    
    # Process files
    target_files = [
        "Reviving_Life_on_the_Edge_Joint_Score_Based_Graph_Generation_of_Rich_Edge_Attributes_parsed_chunked.json",
        "69467_284308_1_PB_parsed_chunked.json",
        "SemanticGraph2Vec_Semantic_graph_embedding_for_text_representation_parsed_chunked.json",
        "Sistem_Budidaya_Ikan_Lele_Clarias_sp_dalam_Kolam_Terpal_dengan_parsed_chunked.json",
        "The_Hybrid_Multimodal_Graph_Index_HMGI_A_ComprehensiveFramework_for_Integrated_Relational_and_Vector_Search_parsed_chunked.json",
        "Retrieval_Augmented_Generation_for_Knowledge_Intensive_NLP_Tasks_parsed_chunked.json",
        "PaperQA_Retrieval_Augmented_Generative_Agent_for_Scientific_Research_parsed_chunked.json",
        "Proses_Pembuatan_Pupuk_Kompos_dan_Penggunaan_Pupuk_Kompos_di_TPS_3R_parsed_chunked.json"
    ]
    
    for filename in target_files:
        file_path = os.path.join(INPUT_DIR, filename)
        if os.path.exists(file_path):
            process_file(file_path, retriever)
        else:
            print(f"File not found: {file_path}")
            
    print("Benchmark completed. Check Output directories for results.")