import sys
import os
import time
import logging
import json
import tiktoken
from dotenv import load_dotenv

PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

load_dotenv(override=True)

from neo4j import GraphDatabase
from openai import OpenAI

import app.services.ingestion
def monkey_patch_is_valid_embedding(vec):
    if not vec or len(vec) != 1536:
        return False
    import math
    if not all(math.isfinite(v) for v in vec):
        return False
    l2 = math.sqrt(sum(v * v for v in vec))
    return l2 > 0.0
app.services.ingestion.is_valid_embedding = monkey_patch_is_valid_embedding

from app.services.ingestion import (
    get_md5_hash,
    embed_text,
    create_literature_node,
    find_similar_nodes,
    PILLAR_THRESHOLD_MATRIX,
    VECTOR_TOP_K
)

ARCHIVE_DIR = os.path.join(PROJECT_ROOT, "archive", "benchmark-graph")
MATRIX_DIR = os.path.join(ARCHIVE_DIR, "Perform-Matrix")
ERROR_DIR = os.path.join(ARCHIVE_DIR, "Log-Error")
ERROR_FILE = os.path.join(ERROR_DIR, "graph_errors.log")

MATRIX_NODE_FILE = os.path.join(MATRIX_DIR, "performance_matrix_node.md")
MATRIX_EDGE_FILE = os.path.join(MATRIX_DIR, "performance_matrix_edge.md")

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

logger = logging.getLogger("benchmark_graph")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)

file_handler = logging.FileHandler(ERROR_FILE, mode='a')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

PRICE_EMBED_PER_M = 0.020  
_tiktoken_encoder = tiktoken.get_encoding("cl100k_base")

node_metrics = []

def record_node_metric(document: str, embed_model: str, pillars_count: int, token_in: int, cost: float, embed_lat: float, db_lat: float, total_time: float, status: str):
    now = time.strftime("%Y-%m-%d %H:%M:%S")
    node_metrics.append({
        "timestamp": now,
        "document": document,
        "embed_model": embed_model,
        "pillars_count": pillars_count,
        "token_in": token_in,
        "embed_cost": cost,
        "embed_latency": embed_lat,
        "db_insert_latency": db_lat,
        "total_time": total_time,
        "status": status
    })

def write_node_matrix():
    file_exists = os.path.isfile(MATRIX_NODE_FILE)
    with open(MATRIX_NODE_FILE, 'a', encoding='utf-8') as f:
        if not file_exists:
            f.write("| Timestamp | Document | Embed Model | Pillars Count | Token Input | Embed Cost ($) | Embed Latency | DB Insert Latency | Total Time | Status |\n")
            f.write("| :--- | :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :--- |\n")
        
        for m in node_metrics:
            cost_str = f"${m['embed_cost']:.6f}"
            row = (
                f"| {m['timestamp']} "
                f"| {m['document'][:20]}... "
                f"| {m['embed_model']} "
                f"| {m['pillars_count']} "
                f"| {m['token_in']} "
                f"| {cost_str} "
                f"| {m['embed_latency']:.2f}s "
                f"| {m['db_insert_latency']:.2f}s "
                f"| {m['total_time']:.2f}s "
                f"| {m['status']} |\n"
            )
            f.write(row)
            
    print("\n--- PERFORMANCE MATRIX: NODE GENERATION ---")
    print("| Timestamp | Document | Embed Model | Pillars Count | Token Input | Embed Cost ($) | Embed Latency | DB Insert Latency | Total Time | Status |")
    print("| :--- | :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :--- |")
    for m in node_metrics:
        cost_str = f"${m['embed_cost']:.6f}"
        print(f"| {m['timestamp']} | {m['document'][:20]}... | {m['embed_model']} | {m['pillars_count']} | {m['token_in']} | {cost_str} | {m['embed_latency']:.2f}s | {m['db_insert_latency']:.2f}s | {m['total_time']:.2f}s | {m['status']} |")

LITERATURE_DIR = os.path.join(PROJECT_ROOT, "Literature")
JSON_OUTPUT_DIR = os.path.join(PROJECT_ROOT, "archive", "Testing-Pipeline", "benchmark-extraction", "Output-JSON")

FILE_MAPPINGS = [
    ("69467_284308_1_PB_parsed_extraction_results.json", "69467-284308-1-PB.pdf"),
    ("PaperQA_Retrieval_Augmented_Generative_Agent_for_Scientific_Research_parsed_extraction_results.json", "PaperQA Retrieval-Augmented Generative Agent for Scientific Research.pdf"),
    ("Proses_Pembuatan_Pupuk_Kompos_dan_Penggunaan_Pupuk_Kompos_di_TPS_3R_parsed_extraction_results.json", "Proses Pembuatan Pupuk Kompos dan Penggunaan Pupuk Kompos  di TPS 3R.pdf"),
    ("RAGAS_Automated_Evaluation_of_Retrieval_Augmented_Generation_parsed_extraction_results.json", "RAGAS Automated Evaluation of Retrieval Augmented Generation.pdf"),
    ("Retrieval_Augmented_Generation_for_Knowledge_Intensive_NLP_Tasks_parsed_extraction_results.json", "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.pdf"),
    ("Reviving_Life_on_the_Edge_Joint_Score_Based_Graph_Generation_of_Rich_Edge_Attributes_parsed_extraction_results.json", "Reviving Life on the Edge Joint Score-Based Graph Generation of Rich Edge Attributes.pdf"),
    ("SciLit_Joint_Literature_Discovery_Summarization_and_Citation_Generation_parsed_extraction_results.json", "SciLit Joint Literature Discovery, Summarization and Citation Generation.pdf"),
    ("SemanticGraph2Vec_Semantic_graph_embedding_for_text_representation_parsed_extraction_results.json", "SemanticGraph2Vec Semantic graph embedding for text representation.pdf"),
    ("Sistem_Budidaya_Ikan_Lele_Clarias_sp_dalam_Kolam_Terpal_dengan_parsed_extraction_results.json", "Sistem Budidaya Ikan Lele (Clarias sp.) dalam Kolam Terpal dengan.pdf"),
    ("The_Hybrid_Multimodal_Graph_Index_HMGI_A_ComprehensiveFramework_for_Integrated_Relational_and_Vector_Search_parsed_extraction_results.json", "The Hybrid Multimodal Graph Index (HMGI): A ComprehensiveFramework for Integrated Relational and Vector Search.pdf")
]

MOCK_DATA = []
for json_file, pdf_file in FILE_MAPPINGS:
    json_path = os.path.join(JSON_OUTPUT_DIR, json_file)
    pdf_path = os.path.join(LITERATURE_DIR, pdf_file)
    
    if os.path.exists(json_path) and os.path.exists(pdf_path):
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            MOCK_DATA.append({
                "filename": pdf_file,
                "pdf_path": pdf_path,
                "metadata": data.get("metadata", {}),
                "pillars": data.get("pillars", {})
            })
    else:
        logger.warning(f"File missing for {pdf_file}. Checked JSON: {os.path.exists(json_path)}, PDF: {os.path.exists(pdf_path)}")

def setup_vector_indexes(driver):
    logger.info("Setting up Neo4j vector indexes...")
    t_start = time.time()
    created_count = 0
    pillars = ["background", "methodology", "gap", "objective", "futurework"]
    
    try:
        with driver.session() as session:
            for pillar in pillars:
                index_name = f"{pillar}_emb"
                prop_name = f"embedding_{pillar}"
                session.run(f"""
                CREATE VECTOR INDEX {index_name} IF NOT EXISTS
                FOR (n:Literature)
                ON (n.{prop_name})
                OPTIONS {{indexConfig: {{
                    `vector.dimensions`: 1536,
                    `vector.similarity_function`: 'cosine'
                }}}}
                """)
                created_count += 1
        
        latency = time.time() - t_start
        logger.info(f"Successfully ensured {created_count} vector indexes. Latency: {latency:.2f}s")
        return True
    except Exception as e:
        logger.error(f"Failed to setup vector indexes: {e}")
        return False

def main():
    logger.info("Starting Graph & Vector Database Benchmark Phase 1 (Setup)...")
    
    neo4j_uri = os.getenv("NEO4J_URI", "")
    neo4j_user = os.getenv("NEO4J_USERNAME", "neo4j")
    neo4j_pass = os.getenv("NEO4J_PASSWORD", "")
    openai_key = os.getenv("OPENAI_API_KEY", "")
    
    if not all([neo4j_uri, neo4j_user, neo4j_pass, openai_key]):
        logger.error("Missing Neo4j or OpenAI credentials in environment.")
        sys.exit(1)
        
    driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_pass))
    openai_client = OpenAI(api_key=openai_key, timeout=30.0)
    
    if not setup_vector_indexes(driver):
        logger.error("Phase 1 initialization failed.")
        sys.exit(1)
        
    logger.info("Phase 1 setup complete.")
    
    logger.info("Starting Phase 3: Node Injection & Embedding...")
    
    node_ids = []
    
    for item in MOCK_DATA:
        filename = item["filename"]
        pdf_path = item["pdf_path"]
        
        logger.info(f"Processing: {filename}")
        t_start_total = time.time()
        
        if not os.path.exists(pdf_path):
            logger.error(f"PDF not found: {pdf_path}")
            record_node_metric(filename, "text-embedding-3-small", 0, 0, 0.0, 0.0, 0.0, 0.0, "Failed - No PDF")
            continue
            
        file_hash = get_md5_hash(pdf_path)
        logger.info(f"Generated MD5 Hash: {file_hash}")
        
        t_start_embed = time.time()
        vectors = {}
        total_embed_in = 0
        pillars_embedded = 0
        
        for key in ["background", "methodology", "gap", "objective", "futurework"]:
            pillar_text = item["pillars"].get(key, "").strip()
            if not pillar_text:
                vectors[key] = None
                continue
                
            try:
                vec = embed_text(pillar_text, openai_client=openai_client)
                vectors[key] = vec
                if vec:
                    total_embed_in += len(_tiktoken_encoder.encode(pillar_text))
                    pillars_embedded += 1
            except Exception as e:
                logger.error(f"Embedding failed for {key}: {e}")
                vectors[key] = None
                
        latency_embed = time.time() - t_start_embed
        
        t_start_db = time.time()
        try:
            node_id = create_literature_node(
                driver=driver,
                file_hash=file_hash,
                filename=filename,
                metadata=item["metadata"],
                pillars=item["pillars"],
                vectors=vectors
            )
            node_ids.append((node_id, item))
            latency_db = time.time() - t_start_db
            
            latency_total = time.time() - t_start_total
            cost_embed = (total_embed_in / 1_000_000) * PRICE_EMBED_PER_M
            
            record_node_metric(
                document=filename,
                embed_model="text-embedding-3-small",
                pillars_count=pillars_embedded,
                token_in=total_embed_in,
                cost=cost_embed,
                embed_lat=latency_embed,
                db_lat=latency_db,
                total_time=latency_total,
                status="Success"
            )
            logger.info(f"Successfully injected node. ElementId: {node_id}")
            
        except Exception as e:
            latency_db = time.time() - t_start_db
            latency_total = time.time() - t_start_total
            record_node_metric(filename, "text-embedding-3-small", pillars_embedded, total_embed_in, 0.0, latency_embed, latency_db, latency_total, f"Failed DB: {e}"[:30])
            logger.error(f"Failed to create node: {e}")

    write_node_matrix()

    # --- Phase 4: Semantic Retrieval ---
    logger.info("Starting Phase 4: Semantic Retrieval...")
    
    # Structure for edge metrics
    edge_metrics = []
    def record_edge_metric(source_doc, target_doc, edge_type, token_in, token_out, cost, total_time, status, gate_retr, gate_coh, gate_dom, gate_reas):
        now = time.strftime("%Y-%m-%d %H:%M:%S")
        edge_metrics.append({
            "timestamp": now,
            "source_doc": source_doc,
            "target_doc": target_doc,
            "edge_type": edge_type,
            "token_in": token_in if token_in > 0 else "-",
            "token_out": token_out if token_out > 0 else "-",
            "cost": cost,
            "total_time": total_time,
            "status": status,
            "gate_retr": gate_retr,
            "gate_coh": gate_coh,
            "gate_dom": gate_dom,
            "gate_reas": gate_reas
        })

    def write_edge_matrix():
        file_exists = os.path.isfile(MATRIX_EDGE_FILE)
        with open(MATRIX_EDGE_FILE, 'a', encoding='utf-8') as f:
            if not file_exists:
                f.write("| Timestamp | Source Doc | Target Doc | Edge Type | Token In | Token Out | Cost ($) | Total Time | Status | Gating: Retrieval | Gating: Coherence | Gating: Domain | Gating: Reasoning |\n")
                f.write("| :--- | :--- | :--- | :--- | :---: | :---: | :---: | :---: | :--- | :--- | :--- | :--- | :--- |\n")
            
            for m in edge_metrics:
                cost_str = f"${m['cost']:.6f}" if m['cost'] > 0 else "-"
                row = (
                    f"| {m['timestamp']} "
                    f"| {m['source_doc'][:15]}... "
                    f"| {m['target_doc'][:15]}... "
                    f"| {m['edge_type']} "
                    f"| {m['token_in']} "
                    f"| {m['token_out']} "
                    f"| {cost_str} "
                    f"| {m['total_time']:.2f}s "
                    f"| {m['status']} "
                    f"| {m['gate_retr']} "
                    f"| {m['gate_coh']} "
                    f"| {m['gate_dom']} "
                    f"| {m['gate_reas']} |\n"
                )
                f.write(row)
        
        print("\n--- PERFORMANCE MATRIX: EDGE GENERATION ---")
        print("| Timestamp | Source Doc | Target Doc | Edge Type | Token In | Token Out | Cost ($) | Total Time | Status | Gating: Retrieval | Gating: Coherence | Gating: Domain | Gating: Reasoning |")
        print("| :--- | :--- | :--- | :--- | :---: | :---: | :---: | :---: | :--- | :--- | :--- | :--- | :--- |")
        for m in edge_metrics:
            cost_str = f"${m['cost']:.6f}" if m['cost'] > 0 else "-"
            print(f"| {m['timestamp']} | {m['source_doc'][:15]}... | {m['target_doc'][:15]}... | {m['edge_type']} | {m['token_in']} | {m['token_out']} | {cost_str} | {m['total_time']:.2f}s | {m['status']} | {m['gate_retr']} | {m['gate_coh']} | {m['gate_dom']} | {m['gate_reas']} |")

    # Assuming node_ids contains tuples of (node_id, item_dict)
    # We will use the first node as source and test against the rest
    source_node_id, source_item = node_ids[0]
    source_title = source_item["filename"]
    
    t_start_retrieval = time.time()
    audit_events = []
    
    try:
        similar_targets = find_similar_nodes(
            driver=driver,
            node_id=source_node_id,
            threshold_matrix=PILLAR_THRESHOLD_MATRIX,
            top_k=VECTOR_TOP_K,
            audit_events=audit_events
        )
        latency_retrieval = time.time() - t_start_retrieval
        
        # We need to process the audit_events to populate the Edge Matrix
        # since audit_events contains the pass/fail for each pillar retrieval.
        
        # Group audit events by target_id and pillar
        # To simulate realistic processing, we must understand what Neo4j returned.
        # Even if a pillar failed the threshold, we want to log it in the edge matrix.
        
        for ev in audit_events:
            target_title = ev.get('target_title', 'Unknown')
            pillar = ev.get('pillar', 'UNKNOWN').upper()
            score = ev.get('score', 0.0)
            status_retrieval = f"Pass ({score:.2f})" if ev.get('status') == 'PASS' else f"Fail ({score:.2f})"
            
            # If it failed retrieval, it stops here. We log it.
            if ev.get('status') != 'PASS':
                record_edge_metric(
                    source_doc=source_title,
                    target_doc=target_title,
                    edge_type=pillar,
                    token_in=0,
                    token_out=0,
                    cost=0.0,
                    total_time=latency_retrieval,
                    status="Reject",
                    gate_retr=status_retrieval,
                    gate_coh="-",
                    gate_dom="-",
                    gate_reas="-"
                )
                
        if similar_targets:
            t_start_gate = time.time()
            bundles = bundle_candidates_by_target(similar_targets)
            
            for target_id, bundle in bundles.items():
                target_title = bundle["title"]
                
                # Phase 5: Coherence Gate
                coherent, reason_coh = validate_semantic_coherence(bundle)
                status_coh = f"Pass ({len(bundle['pillar_scores'])} Match)" if coherent else f"Fail ({reason_coh})"
                
                if not coherent:
                    # Log rejection for each passed pillar in this bundle
                    for pillar_name in bundle["pillar_scores"].keys():
                        score = bundle["pillar_scores"][pillar_name]
                        record_edge_metric(
                            source_doc=source_title,
                            target_doc=target_title,
                            edge_type=pillar_name.upper(),
                            token_in=0,
                            token_out=0,
                            cost=0.0,
                            total_time=time.time() - t_start_retrieval,
                            status="Reject",
                            gate_retr=f"Pass ({score:.2f})",
                            gate_coh=status_coh,
                            gate_dom="-",
                            gate_reas="-"
                        )
                    continue

                # Phase 5: Domain Gate
                kw_source = source_item["metadata"].get("keywords", [])
                kw_target = bundle.get("keywords", [])
                
                prompt_domain = str(kw_source) + str(kw_target)
                tok_in_dom = len(_tiktoken_encoder.encode(prompt_domain)) + 150
                
                domain_res = evaluate_domain_compatibility_with_llm(ali_client, kw_source, kw_target)
                tok_out_dom = 15
                
                status_dom = "Pass" if domain_res["is_compatible"] else f"Fail ({domain_res['reason']})"
                
                if not domain_res["is_compatible"]:
                    for pillar_name in bundle["pillar_scores"].keys():
                        score = bundle["pillar_scores"][pillar_name]
                        record_edge_metric(
                            source_doc=source_title,
                            target_doc=target_title,
                            edge_type=pillar_name.upper(),
                            token_in=tok_in_dom,
                            token_out=tok_out_dom,
                            cost=(tok_in_dom * PRICE_LLM_IN_PER_M + tok_out_dom * PRICE_LLM_OUT_PER_M) / 1_000_000,
                            total_time=time.time() - t_start_retrieval,
                            status="Reject",
                            gate_retr=f"Pass ({score:.2f})",
                            gate_coh=status_coh,
                            gate_dom=status_dom,
                            gate_reas="-"
                        )
                    continue
                    
                # Phase 6: Edge Reasoning
                t_start_reason = time.time()
                
                pillars_to_reason = {}
                with driver.session() as session:
                    source_rec = session.run("MATCH (n:Literature) WHERE elementId(n) = $id RETURN n", id=source_node_id).single()["n"]
                
                for pillar_name, score in bundle["pillar_scores"].items():
                    source_excerpt = source_rec.get(f"text_{pillar_name}", "")
                    target_excerpt = bundle["pillar_texts"].get(pillar_name, "")
                    pillars_to_reason[pillar_name] = {
                        "source": source_excerpt,
                        "target": target_excerpt,
                        "score": score
                    }
                
                reason_prompt_len = sum(len(_tiktoken_encoder.encode(p["source"] + p["target"])) for p in pillars_to_reason.values()) + 300
                
                try:
                    reasonings = generate_edge_reasoning(
                        ali_client,
                        source_title,
                        bundle["title"],
                        pillars_to_reason
                    )
                    
                    tok_in_reas = reason_prompt_len
                    tok_out_reas = sum(len(_tiktoken_encoder.encode(v)) for v in reasonings.values())
                    
                    edges_to_create = []
                    matched_pillars = sorted(bundle["pillar_scores"].keys())
                    
                    for pillar_name, data in pillars_to_reason.items():
                        reasoning = reasonings.get(pillar_name, "")
                        quality_ok, quality_reason = validate_edge_reasoning_quality(reasoning)
                        
                        status_reas = "Pass" if quality_ok else f"Fail ({quality_reason})"
                        
                        # Calculate proportional cost for this pillar
                        prop_tok_in = (tok_in_dom + tok_in_reas) / len(pillars_to_reason)
                        prop_tok_out = (tok_out_dom + tok_out_reas) / len(pillars_to_reason)
                        prop_cost = (prop_tok_in * PRICE_LLM_IN_PER_M + prop_tok_out * PRICE_LLM_OUT_PER_M) / 1_000_000
                        
                        if quality_ok:
                            edges_to_create.append({
                                "source_id": source_node_id,
                                "target_id": target_id,
                                "score": data["score"],
                                "reasoning": reasoning,
                                "category": pillar_name,
                                "matched_pillars": matched_pillars,
                            })
                            record_edge_metric(
                                source_doc=source_title,
                                target_doc=target_title,
                                edge_type=pillar_name.upper(),
                                token_in=int(prop_tok_in),
                                token_out=int(prop_tok_out),
                                cost=prop_cost,
                                total_time=time.time() - t_start_retrieval,
                                status="Success",
                                gate_retr=f"Pass ({data['score']:.2f})",
                                gate_coh=status_coh,
                                gate_dom=status_dom,
                                gate_reas=status_reas
                            )
                        else:
                            record_edge_metric(
                                source_doc=source_title,
                                target_doc=target_title,
                                edge_type=pillar_name.upper(),
                                token_in=int(prop_tok_in),
                                token_out=int(prop_tok_out),
                                cost=prop_cost,
                                total_time=time.time() - t_start_retrieval,
                                status="Reject",
                                gate_retr=f"Pass ({data['score']:.2f})",
                                gate_coh=status_coh,
                                gate_dom=status_dom,
                                gate_reas=status_reas
                            )

                    # Phase 7: Edge Batch Insert
                    if edges_to_create:
                        rel_ids = create_lexicographical_edges_batch(driver, edges_to_create)
                        logger.info(f"Created {len(rel_ids)} edges for target {target_title}")

                except Exception as e:
                    logger.error(f"Edge generation failed: {e}")
                    for pillar_name, data in pillars_to_reason.items():
                        record_edge_metric(
                            source_doc=source_title,
                            target_doc=target_title,
                            edge_type=pillar_name.upper(),
                            token_in=tok_in_dom,
                            token_out=tok_out_dom,
                            cost=(tok_in_dom * PRICE_LLM_IN_PER_M + tok_out_dom * PRICE_LLM_OUT_PER_M) / 1_000_000,
                            total_time=time.time() - t_start_retrieval,
                            status="Reject",
                            gate_retr=f"Pass ({data['score']:.2f})",
                            gate_coh=status_coh,
                            gate_dom=status_dom,
                            gate_reas=f"Fail (Exception)"
                        )

        logger.info(f"Retrieval complete. Found {len(similar_targets)} passing pillars.")
        
    except Exception as e:
        logger.error(f"Retrieval failed: {e}")
        
    write_edge_matrix()

    driver.close()
    logger.info("Pipeline Execution Complete.")

if __name__ == "__main__":
    main()
