import os
import sys
import pandas as pd
import json
import re
from pathlib import Path
from typing import List

# Add get_knowledgeBase_tool to path
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
TOOL_DIR = os.path.join(SCRIPT_DIR, "get_knowledgeBase_tool")
if TOOL_DIR not in sys.path:
    sys.path.insert(0, TOOL_DIR)

# Import from tool directory
sys.path.insert(0, TOOL_DIR)
import ace_lib
from fetch_all_operators import fetch_operators, prompt_credentials
from fetch_all_documentation import (
    fetch_tutorials,
    fetch_tutorial_pages,
    fetch_page,
    _extract_page_id,
)
# Dataset fetching currently disabled per request
# from fetch_all_datasets import (
#     fetch_all_combinations,
#     fetch_datasets_for_combo,
#     merge_and_deduplicate,
# )


def ensure_knowledge_dir():
    """Ensure knowledge directory exists"""
    knowledge_dir = os.path.join(SCRIPT_DIR, "knowledge")
    os.makedirs(knowledge_dir, exist_ok=True)
    return knowledge_dir


def to_jsonable(value):
    """Convert values to JSON-serializable, handling NaN and nested structures."""
    try:
        if isinstance(value, float) and pd.isna(value):
            return None
    except TypeError:
        pass

    if isinstance(value, list):
        return [to_jsonable(v) for v in value if not (isinstance(v, float) and pd.isna(v))]
    if isinstance(value, dict):
        return {k: to_jsonable(v) for k, v in value.items()}
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    return str(value)


def safe_filename(name: str, suffix: str = "") -> str:
    base = re.sub(r"[^A-Za-z0-9._-]+", "_", str(name)).strip("_") or "doc"
    base = base[:80]
    return f"{base}{suffix}"


def process_operators(session: ace_lib.SingleSession, knowledge_dir: str):
    """
    Process operators and save as JSON files
    
    Args:
        session: Authenticated BRAIN session
        knowledge_dir: Directory to save JSON files
    """
    print("\n=== Processing Operators ===")
    
    # Fetch operators data
    print("Fetching operators...")
    operators_df = fetch_operators(session)
    
    if operators_df.empty:
        print("No operators found!")
        return
    
    print(f"Found {len(operators_df)} operator entries")
    
    # Get unique categories
    categories = sorted(operators_df['category'].dropna().unique())
    
    for category in categories:
        category_data = operators_df[operators_df['category'] == category].copy()
        
        # Create JSON file for this category
        filename = f"{category.replace(' ', '_').lower()}_operators.json"
        filepath = os.path.join(knowledge_dir, filename)
        
        print(f"Processing category: {category}")
        
        # Convert to list of dicts
        category_list = []
        for idx, row in category_data.iterrows():
            operator_dict = {}
            for col in row.index:
                value = row[col]
                operator_dict[col] = to_jsonable(value)
            category_list.append(operator_dict)
        
        # Save category JSON
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(category_list, f, ensure_ascii=False, indent=2)
        
        print(f"✓ Created {filename} with {len(category_list)} operators")


    # Dataset fetching intentionally disabled; keep for potential re-enable.
    # def process_datasets(session: ace_lib.SingleSession, dataset_dir: str):
    #     """Fetch datasets and save one JSON per region."""
    #     print("=== Processing Datasets ===")
    #
    #     print("Fetching valid instrument/region/delay/universe combinations...")
    #     options_df = fetch_all_combinations(session)
    #     if options_df is None or options_df.empty:
    #         print("No simulation options fetched; aborting dataset fetch.")
    #         return
    #
    #     all_datasets: list[pd.DataFrame] = []
    #     combo_idx = 0
    #
    #     for _, row in options_df.iterrows():
    #         instrument_type = row.get("InstrumentType")
    #         region = row.get("Region")
    #         delay = row.get("Delay")
    #         universes = row.get("Universe") or []
    #
    #         for universe in universes:
    #             combo_idx += 1
    #             print(f"[{combo_idx}] {instrument_type} / {region} / D{delay} / {universe}")
    #             try:
    #                 df = fetch_datasets_for_combo(session, instrument_type, region, delay, universe)
    #                 print(f"    -> {len(df)} rows")
    #                 all_datasets.append(df)
    #             except Exception as exc:
    #                 print(f"    -> Failed: {exc}")
    #
    #     if not all_datasets:
    #         print("No datasets fetched; nothing to save.")
    #         return
    #
    #     combined_df = pd.concat([df for df in all_datasets if not df.empty], ignore_index=True)
    #     if combined_df.empty:
    #         print("No datasets fetched; nothing to save.")
    #         return
    #
    #     regions = sorted(combined_df["param_region"].dropna().unique())
    #     print(f"Found regions: {', '.join(regions)}")
    #
    #     for region in regions:
    #         region_df = combined_df[combined_df["param_region"] == region]
    #         region_unique = merge_and_deduplicate([region_df])
    #
    #         region_list = []
    #         for _, row in region_unique.iterrows():
    #             record = {col: to_jsonable(row[col]) for col in row.index}
    #             region_list.append(record)
    #
    #         filename = f"{region.replace(' ', '_').lower()}_datasets.json"
    #         filepath = os.path.join(dataset_dir, filename)
    #         with open(filepath, "w", encoding="utf-8") as f:
    #             json.dump(region_list, f, ensure_ascii=False, indent=2)
    #
    #         print(f"✓ Created {filename} with {len(region_list)} datasets")


def process_documentation(session: ace_lib.SingleSession, knowledge_dir: str):
    """Fetch tutorials and pages, save one JSON per page."""
    print("=== Processing Documentation ===")

    tutorials = fetch_tutorials(session)
    if not tutorials:
        print("No tutorials fetched; skipping documentation.")
        return

    print(f"Fetched {len(tutorials)} tutorials")

    page_count = 0
    seen_pages = set()
    
    for idx, tutorial in enumerate(tutorials, start=1):
        tutorial_id = _extract_page_id(tutorial) or f"tutorial_{idx}"
        tutorial_title = tutorial.get("title") or tutorial_id

        page_candidates = []
        if isinstance(tutorial.get("pages"), list):
            page_candidates.extend(tutorial["pages"])
        if tutorial_id:
            try:
                page_candidates.extend(fetch_tutorial_pages(session, tutorial_id))
            except Exception as exc:
                print(f"[{idx:03d}] failed to fetch pages for {tutorial_id}: {exc}")

        if not page_candidates and tutorial_id:
            page_candidates.append({"id": tutorial_id, "title": tutorial_title})

        for page_entry in page_candidates:
            page_id = _extract_page_id(page_entry)
            if not page_id or page_id in seen_pages:
                continue
            seen_pages.add(page_id)
            
            try:
                page = fetch_page(session, page_id)
            except Exception as exc:
                print(f"[{idx:03d}] page {page_id} failed: {exc}")
                continue
            
            page_count += 1
            page_title = page.get("title") or page_entry.get("title") or page_id
            
            # Save each page as individual JSON
            filename = safe_filename(f"{idx:03d}_{page_title}", "_documentation.json")
            filepath = os.path.join(knowledge_dir, filename)
            
            with open(filepath, "w", encoding="utf-8") as f:
                json.dump(to_jsonable(page), f, ensure_ascii=False, indent=2)
            
            print(f"[{idx:03d}] ✓ Created {filename}")
    
    print(f"✓ Total: {page_count} documentation pages saved")


def main():
    print("=== BRAIN Knowledge Base Processor ===")
    print("Starting operator processing...\n")
    
    # Get credentials
    email, password = prompt_credentials()
    ace_lib.get_credentials = lambda: (email, password)
    
    print("Logging in to BRAIN platform...")
    try:
        session = ace_lib.start_session()
        print("✓ Login successful\n")
    except Exception as exc:
        print(f"✗ Login failed: {exc}")
        return
    
    # Ensure knowledge directory exists
    knowledge_dir = ensure_knowledge_dir()
    # dataset_dir = knowledge_dir  # Save datasets directly under knowledge (disabled)
    print(f"Knowledge directory: {knowledge_dir}\n")

    # Process documentation (tutorials/pages)
    print("\nStarting documentation processing...\n")
    try:
        process_documentation(session, knowledge_dir)
    except Exception as exc:
        print(f"✗ Failed to process documentation: {exc}")
        import traceback
        traceback.print_exc()
        return

    # Process operators
    try:
        process_operators(session, knowledge_dir)
    except Exception as exc:
        print(f"✗ Failed to process operators: {exc}")
        import traceback
        traceback.print_exc()
        return

    # Dataset processing disabled; re-enable by uncommenting the block below.
    # print("\nStarting dataset processing...\n")
    # try:
    #     process_datasets(session, dataset_dir)
    # except Exception as exc:
    #     print(f"✗ Failed to process datasets: {exc}")
    #     import traceback
    #     traceback.print_exc()
    #     return

    print("\n=== Processing Complete ===")


if __name__ == "__main__":
    main()