import sys
import os
import json
import time
import logging
from datetime import datetime
import threading
import tiktoken
from dotenv import load_dotenv
from openai import OpenAI

# Load environment variables
load_dotenv(override=True)

# Adjust system path to import app modules
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))

from app.services.parser.llm_extractor import extract_pillars_with_llm

# Setup Logging
logger = logging.getLogger("benchmark_extraction")
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)

# Directories
BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))
INPUT_CHUNKS_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-chunking/Output-Chunking')
INPUT_RETRIEVAL_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-retrieval/Output-JSON-C') # folder hasil retrieval (Output-JSON-C untuk Candidate C)

OUTPUT_JSON_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-extraction/Output-JSON')
OUTPUT_MD_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-extraction/Output-Human-Verification')
MATRIX_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-extraction/Perform-Matrix')
ERROR_DIR = os.path.join(BASE_DIR, 'archive/Testing-Pipeline/benchmark-extraction/Log-Error')

MATRIX_FILE = os.path.join(MATRIX_DIR, 'performance_matrix_extraction.md')
ERROR_FILE = os.path.join(ERROR_DIR, 'extraction_errors.log')

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

# ==========================================
# AGNOSTIC PROVIDER & PRICING MATRIX CONFIG
# ==========================================

# Tombol Sakelar Lokal Skrip (Ubah ini untuk ganti tes ke OpenAI atau DashScope)
TESTING_PROVIDER = "dashscope" # Options: "dashscope", "openai"
# Opsional: Paksa override nama model jika ingin ngetes model selain di .env
TESTING_MODEL_OVERRIDE = None # Contoh: "gpt-4o-mini"

# Matriks Harga per 1.000.000 Token (USD)
PRICING_MATRIX = {
    "qwen3-30b-a3b-thinking-2507": {"input": 0.10, "output": 0.30},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "gpt-4o": {"input": 5.00, "output": 15.00},
    "default": {"input": 0.10, "output": 0.30}
}

api_semaphore = threading.BoundedSemaphore(3)
_tiktoken_encoder = tiktoken.get_encoding("cl100k_base")

def get_llm_client_and_model():
    """Smart auto-detect and construct OpenAI client based on TESTING_PROVIDER switch."""
    if TESTING_PROVIDER == "openai":
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            logger.error("OPENAI_API_KEY is missing in .env")
            sys.exit(1)
            
        client = OpenAI(api_key=api_key, timeout=60.0)
        # Use override if exists, else fallback to a standard testing model
        model_name = TESTING_MODEL_OVERRIDE or "gpt-4o-mini"
        logger.info(f"Initialized OpenAI Client. Model: {model_name}")
        return client, model_name
        
    elif TESTING_PROVIDER == "dashscope":
        api_key = os.getenv("DASHSCOPE_API_KEY")
        if not api_key:
            logger.error("DASHSCOPE_API_KEY is missing in .env")
            sys.exit(1)
            
        client = OpenAI(
            base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            api_key=api_key,
            timeout=60.0
        )
        model_name = TESTING_MODEL_OVERRIDE or os.getenv("DASHSCOPE_MODEL", "qwen3-30b-a3b-thinking-2507")
        logger.info(f"Initialized DashScope Client. Model: {model_name}")
        return client, model_name
        
    else:
        logger.error(f"Unknown TESTING_PROVIDER: {TESTING_PROVIDER}")
        sys.exit(1)

def calculate_cost(model_name: str, token_in: int, token_out: int) -> float:
    """Kalkulasi biaya dinamis berdasarkan PRICING_MATRIX."""
    pricing = PRICING_MATRIX.get(model_name, PRICING_MATRIX["default"])
    cost_in = (token_in / 1_000_000) * pricing["input"]
    cost_out = (token_out / 1_000_000) * pricing["output"]
    return round(cost_in + cost_out, 6)

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 | MODEL | TIKTOKEN INPUT | LLM OUT TOKEN | LATENCY (s) | COST ($) | STATUS | NOTE |\n")
            f.write("| :--- | :--- | :--- | :---: | :---: | :---: | :---: | :--- | :--- |\n")
        
        doc_name = metrics['document']
        short_name = (doc_name[:25] + '...') if len(doc_name) > 28 else doc_name
        model_name = metrics.get('model', 'Unknown')
        short_model = (model_name[:15] + '...') if len(model_name) > 18 else model_name
        
        cost_str = f"${metrics.get('cost', 0.0):.6f}" if metrics.get('cost') is not None else "N/A"
        
        row = (
            f"| {metrics['timestamp']} "
            f"| {short_name} "
            f"| {short_model} "
            f"| {metrics.get('token_in', 'N/A')} "
            f"| {metrics.get('token_out', 'N/A')} "
            f"| {metrics.get('latency', 'N/A')} "
            f"| {cost_str} "
            f"| {metrics['status']} "
            f"| {metrics.get('note', '-')} |\n"
        )
        f.write(row)

def create_human_verification_md(doc_name: str, extraction_data: dict) -> str:
    metadata = extraction_data.get("metadata", {})
    pillars = extraction_data.get("pillars", {})
    
    lines = [
        f"# Extraction Verification: {doc_name}",
        "",
        "## METADATA",
        f"**Title:** {metadata.get('title', '')}",
        f"**Keywords:** {', '.join(metadata.get('keywords', []))}",
        "",
        "---",
        "",
        "## PILLARS",
        ""
    ]
    
    pillar_names = ["background", "methodology", "gap", "objective", "futurework"]
    for p in pillar_names:
        lines.append(f"### {p.upper()}")
        content = pillars.get(p, "*TIDAK TERSEDIA*")
        # Ensure it's not totally empty visually if the value is empty string
        if not str(content).strip():
            content = "*KOSONG (LLM mengembalikan string kosong)*"
        lines.append(f"{content}\n")
        
    return "\n".join(lines)

def assemble_context(chunk_data: dict, retrieval_data: dict) -> str:
    """Replicate logic from ingestion.py to build the stuffed context."""
    parent_chunks = chunk_data.get('parent_chunks', [])
    
    # 1. Deduplicate Top Parent IDs
    unique_parent_ids = set()
    for result in retrieval_data.get('retrieval_results', []):
        top_parents = result.get('top_parents', [])
        unique_parent_ids.update(top_parents)
        
    # 2. Extract texts
    selected_parent_texts = []
    for p in parent_chunks:
        if p.get("parent_id") in unique_parent_ids:
            selected_parent_texts.append(p.get("content", ""))
            
    stuffed_context = "\n\n---\n\n".join(selected_parent_texts)
    
    # 3. Metadata Injection (first 1500 chars)
    if parent_chunks:
        first_chunk_text = parent_chunks[0].get("content", "")
        global_context = first_chunk_text[:1500]
        stuffed_context = f"--- BAGIAN RELEVAN: METADATA & JUDUL ARTIKEL ---\n{global_context}\n\n" + stuffed_context
    else:
        logger.warning("No parent chunks found to extract metadata.")
        
    return stuffed_context

def process_file(chunk_file_path: str, retrieval_file_path: str, client: OpenAI, model_name: str):
    filename = os.path.basename(chunk_file_path)
    base_name = filename.replace('_chunked.json', '')
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    metrics = {
        'timestamp': timestamp,
        'document': base_name,
        'model': model_name,
        'status': 'Failed',
        'note': '-'
    }
    
    logger.info(f"Processing {base_name}...")
    
    try:
        # Load Data
        with open(chunk_file_path, 'r', encoding='utf-8') as f:
            chunk_data = json.load(f)
            
        with open(retrieval_file_path, 'r', encoding='utf-8') as f:
            retrieval_data = json.load(f)
            
        # Context Assembler
        raw_text = assemble_context(chunk_data, retrieval_data)
        
        if not raw_text.strip():
            raise ValueError("Assembled context is empty.")
            
        # Calculate Tiktoken Input (cl100k_base)
        truncated_text = raw_text[:25000]
        token_count = len(_tiktoken_encoder.encode(truncated_text))
        metrics['token_in'] = token_count
        
        # LLM Invocation
        t_start = time.time()
        extraction_result, usage_stats = extract_pillars_with_llm(
            text=truncated_text,
            client=client,
            semaphore=api_semaphore,
            model_name=model_name
        )
        t_end = time.time()
        
        latency = round(t_end - t_start, 2)
        metrics['latency'] = latency
        metrics['token_out'] = usage_stats.get('completion_tokens', 0)
        
        # Calculate Cost
        metrics['cost'] = calculate_cost(model_name, metrics['token_in'], metrics['token_out'])
        
        # Validate Success
        if not extraction_result.get("pillars"):
            metrics['status'] = 'Failed'
            metrics['note'] = 'JSON pillars empty/failed'
        else:
            metrics['status'] = 'Success'
            
        # Generate Artifacts
        json_output_path = os.path.join(OUTPUT_JSON_DIR, f"{base_name}_extraction_results.json")
        with open(json_output_path, 'w', encoding='utf-8') as f:
            json.dump(extraction_result, f, ensure_ascii=False, indent=2)
            
        md_content = create_human_verification_md(base_name, extraction_result)
        md_output_path = os.path.join(OUTPUT_MD_DIR, f"{base_name}_extraction_human.md")
        with open(md_output_path, 'w', encoding='utf-8') as f:
            f.write(md_content)

    except Exception as e:
        error_msg = f"{type(e).__name__}: {str(e)}"
        logger.error(f"Failed processing {base_name}: {error_msg}", exc_info=True)
        metrics['note'] = "Exception logged"
        
    log_performance_matrix(metrics)
    logger.info(f"Finished {base_name} -> Status: {metrics['status']}, Cost: ${metrics.get('cost', 0):.4f}")

if __name__ == "__main__":
    logger.info("Starting Extraction Benchmark...")
    
    # Init client using the agnostic setup
    client, model_name = get_llm_client_and_model()
    
    # target_bases = [
    #     "69467_284308_1_PB_parsed_retrieval_results.json",
    #     "PaperQA_Retrieval_Augmented_Generative_Agent_for_Scientific_Research_parsed_retrieval_results.json",
    #     "Proses_Pembuatan_Pupuk_Kompos_dan_Penggunaan_Pupuk_Kompos_di_TPS_3R_parsed_retrieval_results.json",
    #     "RAGAS_Automated_Evaluation_of_Retrieval_Augmented_Generation_parsed_retrieval_results.json",
    #     "Retrieval_Augmented_Generation_for_Knowledge_Intensive_NLP_Tasks_parsed_retrieval_results.json",
    #     "Reviving_Life_on_the_Edge_Joint_Score_Based_Graph_Generation_of_Rich_Edge_Attributes_parsed_retrieval_results.json",
    #     "SciLit_Joint_Literature_Discovery_Summarization_and_Citation_Generation_parsed_retrieval_results.json",
    #     "SemanticGraph2Vec_Semantic_graph_embedding_for_text_representation_parsed_retrieval_results.json",
    #     "Sistem_Budidaya_Ikan_Lele_Clarias_sp_dalam_Kolam_Terpal_dengan_parsed_retrieval_results.json",
    #     "The_Hybrid_Multimodal_Graph_Index_HMGI_A_ComprehensiveFramework_for_Integrated_Relational_and_Vector_Search_parsed_retrieval_results.json"
    # ]
    target_bases = [
        "69467_284308_1_PB_parsed",
        "PaperQA_Retrieval_Augmented_Generative_Agent_for_Scientific_Research_parsed",
        "Proses_Pembuatan_Pupuk_Kompos_dan_Penggunaan_Pupuk_Kompos_di_TPS_3R_parsed",
        "RAGAS_Automated_Evaluation_of_Retrieval_Augmented_Generation_parsed",
        "Retrieval_Augmented_Generation_for_Knowledge_Intensive_NLP_Tasks_parsed",
        "Reviving_Life_on_the_Edge_Joint_Score_Based_Graph_Generation_of_Rich_Edge_Attributes_parsed",
        "SciLit_Joint_Literature_Discovery_Summarization_and_Citation_Generation_parsed",
        "SemanticGraph2Vec_Semantic_graph_embedding_for_text_representation_parsed",
        "Sistem_Budidaya_Ikan_Lele_Clarias_sp_dalam_Kolam_Terpal_dengan_parsed",
        "The_Hybrid_Multimodal_Graph_Index_HMGI_A_ComprehensiveFramework_for_Integrated_Relational_and_Vector_Search_parsed"
]
    
    for base in target_bases:
        chunk_file = os.path.join(INPUT_CHUNKS_DIR, f"{base}_chunked.json")
        retrieval_file = os.path.join(INPUT_RETRIEVAL_DIR, f"{base}_retrieval_results.json")
        
        if os.path.exists(chunk_file) and os.path.exists(retrieval_file):
            process_file(chunk_file, retrieval_file, client, model_name)
        else:
            logger.warning(f"Missing input files for {base}. Skipping.")
            
    logger.info("Extraction Benchmark Completed.")
