import json
import re
import warnings
import argparse
from pathlib import Path
from collections import Counter

import pandas as pd
import numpy as np
import frontmatter
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import MiniBatchKMeans
from sklearn.decomposition import TruncatedSVD

import common

# warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# --- GLOBAL CONFIGURATION ---
TARGET_BRANCHING_FACTOR = 7  # The "Rule of 7"
GOLD_PAN_SIZE = 5            # Top articles kept at hub level
NAVGRAPH_FILE = "navgraph.json"
GRAPH_FILE = "graph.json"
LLMS_TXT_FILE = "llms.txt"
SITEMAP_CORE_FILE = "sitemap-core.xml"
SITEMAP_HUBS_FILE = "sitemap-hubs.xml"

# --- 1. UNIFIED DATA INGESTION ---

def slugify(text):
    if not text: return ""
    text = str(text).lower()
    text = re.sub(r'[^a-z0-9\s-]', '', text)
    text = re.sub(r'\s+', '-', text)
    return text.strip('-')

def load_enriched_shards(context_dir, posts_dir):
    """
    Deep Ingestion: Reads JSON shards AND Markdown Frontmatter.
    This ensures the Graph and the Nav use the exact same Titles and Permalinks.
    """
    shards = []
    if not context_dir.exists():
         print(f"⚠️ Context dir {context_dir} does not exist.")
         return pd.DataFrame()

    files = list(context_dir.glob("*.json"))
    print(f"💎 Loading {len(files)} shards from {context_dir}...")
    
    for f in files:
        try:
            with open(f, 'r', encoding='utf-8') as file:
                data = json.load(file)
            
            md_path = posts_dir / f"{f.stem}.md"
            if not md_path.exists():
                continue

            post = frontmatter.load(md_path)
            
            # --- DEFENSIVE TITLE EXTRACTION ---
            title = post.metadata.get('title')
            if not title:
                title = data.get('t', 'Untitled')
            if not title:
                title = "Untitled"
            # ----------------------------------

            # Weighted Soup: Title gets 3x weight
            soup = (
                (str(title) + " ") * 3 + 
                (" ".join(data.get('kw', [])) + " ") * 2 + 
                " ".join(data.get('sub', []))
            )
            
            date_val = post.metadata.get('date', data.get('d', ''))
            
            shards.append({
                "id": f.stem,
                "title": str(title),
                "permalink": post.metadata.get('permalink', f"/{f.stem}/"),
                "description": post.metadata.get('description', data.get('s', '')),
                "date": str(date_val), 
                "soup": soup,
                "keywords": data.get('kw', []) + data.get('sub', [])
            })

        except Exception as e:
            print(f"⚠️ Error loading {f.name}: {e}")
            
    return pd.DataFrame(shards)

def load_velocity_data(directory=Path(".")):
    if not directory.exists(): directory = Path(__file__).parent
    velocity_file = directory / "gsc_velocity.json"
    if not velocity_file.exists(): return {}
    try:
        with open(velocity_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        slug_map = {}
        for key, metrics in data.items():
            if key.startswith("_"): continue
            slug = key.strip('/').split('/')[-1]
            slug_map[slug] = metrics
        return slug_map
    except: return {}

# --- 2. CANONICAL CLUSTERING LOGIC ---

def get_cluster_candidates(df_cluster):
    """Returns a list of (keyword, score) tuples sorted by relevance."""
    all_keywords = [kw for sublist in df_cluster['keywords'] for kw in sublist]
    if not all_keywords: return [("Misc", 0)]
    
    counts = Counter(all_keywords)
    return counts.most_common(10)

def calculate_node_gravity(label, keywords):
    """Calculates visual size (gravity) for D3."""
    # Without external market volume, gravity is determined by the raw number 
    # of semantic keywords attached to the node, granting weight to dense content.
    return 5 + len(keywords)

def build_canonical_tree(df_slice, current_node, current_depth, velocity_data, vectorizer=None):
    """
    The Single Logic Stream.
    Builds a recursive dictionary (Tree) that represents the Truth.
    """
    df = df_slice.copy()

    # Sort by GSC Clicks (High velocity content floats to top)
    df['sort_clicks'] = df['id'].apply(lambda x: velocity_data.get(re.sub(r'^\d{4}-\d{2}-\d{2}-', '', x), {}).get('total_clicks', 0))

    # THE FIX: Deterministic fallback. Sort by Clicks, then by Date (newest first).
    df = df.sort_values(by=['sort_clicks', 'date'], ascending=[False, False])

    def attach_article(row):
        # Calculate organic gravity
        grav = calculate_node_gravity(row['title'], row['keywords'])
        
        # Get status from GSC
        slug = re.sub(r'^\d{4}-\d{2}-\d{2}-', '', row['id'])
        gsc_meta = velocity_data.get(slug, {})
        
        article_node = {
            "type": "article",
            "id": row['id'],
            "title": str(row['title']),
            "permalink": row['permalink'],
            "date": row['date'],
            "gravity": grav,
            "status": gsc_meta.get("status", "unknown"),
            "velocity": gsc_meta.get("velocity", 0),
            "clicks": gsc_meta.get("total_clicks", 0)
        }
        current_node.setdefault('children_articles', []).append(article_node)

    # 1. Stop Condition
    if len(df) <= TARGET_BRANCHING_FACTOR:
        for _, row in df.iterrows(): attach_article(row)
        return

    # 2. Clustering (Evaluate ALL items to find the true leaders)
    if vectorizer is None:
        vectorizer = TfidfVectorizer(stop_words='english', max_features=1000)

    try:
        tfidf_matrix = vectorizer.fit_transform(df['soup'])
        n_components = min(5, len(df) - 1)
        if n_components > 1:
            svd = TruncatedSVD(n_components=n_components)
            matrix = svd.fit_transform(tfidf_matrix)
        else:
            matrix = tfidf_matrix

        kmeans = MiniBatchKMeans(n_clusters=TARGET_BRANCHING_FACTOR, random_state=42, n_init=10, batch_size=256)
        clusters = kmeans.fit_predict(matrix)
        df.loc[:, 'cluster'] = clusters

        # Collision Tracking (Scoped to this level of recursion)
        used_slugs = set()

        for cluster_id in range(TARGET_BRANCHING_FACTOR):
            cluster_data = df[df['cluster'] == cluster_id]
            if len(cluster_data) == 0: continue

            # 3. Flagship Promotion (The Semantic Gold Pan)
            # Extract the absolute best article from THIS specific cluster.
            flagship = cluster_data.iloc[0]
            attach_article(flagship)
            
            # The rest of the cluster gets pushed down into a sub-hub
            cluster_remainder = cluster_data.iloc[1:].copy()
            
            # Prevent "Ghost Hubs" if the cluster only had 1 article
            if len(cluster_remainder) == 0: continue

            # Semantic Labeling & Collision Resolution
            candidates = get_cluster_candidates(cluster_data)
            
            hub_label = "Misc"
            for kw, score in candidates:
                if not kw: continue
                test_slug = slugify(kw)
                if test_slug not in used_slugs:
                    hub_label = kw
                    break
            else:
                # Fallback: Append number
                top_kw = candidates[0][0]
                base_slug = slugify(top_kw)
                counter = 2
                while f"{base_slug}-{counter}" in used_slugs:
                    counter += 1
                hub_label = f"{top_kw} {counter}"

            slug = slugify(hub_label)
            used_slugs.add(slug)
            
            # Create Hub Node
            hub_gravity = calculate_node_gravity(hub_label, [hub_label])
            # Boost Hub gravity based on depth
            hub_val = max(10, 50 - (current_depth * 10)) + hub_gravity

            new_hub_node = {
                "type": "hub",
                "id": f"{current_node['id']}_{cluster_id}",
                "title": hub_label,
                "permalink": f"{current_node['permalink']}{slug}/",
                "blurb": f"Explore {len(cluster_remainder)} articles about {hub_label}.",
                "gravity": hub_val,
                "children_hubs": [],
                "children_articles": []
            }
            
            current_node.setdefault('children_hubs', []).append(new_hub_node)

            # Recurse (CRITICAL FIX: Pass ONLY the remainder)
            build_canonical_tree(
                cluster_remainder, new_hub_node, current_depth + 1,
                velocity_data
            )

    except Exception as e:
        print(f"⚠️ Clustering fallback at depth {current_depth}: {e}")
        for _, row in df.iterrows(): attach_article(row)

# --- 3. PROJECTORS ---

def project_d3_graph(tree_node, nodes, links, parent_id=None, depth=0):
    """
    Projector B: Flattens the Canonical Tree into D3 Nodes/Links.
    Now preserves hierarchy (parentId, depth) for D3 Stratify & Radial Layouts.
    """
    # Create the node for D3
    d3_node = {
        "id": tree_node['id'],
        "label": tree_node['title'],
        # Map canonical type to D3 group (ensures 'root' is preserved)
        "group": tree_node['type'], 
        "val": tree_node.get('gravity', 20),
        "status": "hub",
        "parentId": parent_id,
        "depth": depth
    }
    nodes.append(d3_node)

    # Process Articles (Leaves)
    for article in tree_node.get('children_articles', []):
        art_node = {
            "id": article['id'],
            "label": article['title'],
            "group": "article",
            "val": article.get('gravity', 5),
            "status": article.get('status', 'unknown'),
            "velocity": article.get('velocity', 0),
            "clicks": article.get('clicks', 0),
            "parentId": tree_node['id'],
            "depth": depth + 1
        }
        nodes.append(art_node)
        links.append({
            "source": tree_node['id'],
            "target": article['id'],
            "type": "article_link"
        })

    # Process Sub-Hubs (Recursion)
    for hub in tree_node.get('children_hubs', []):
        links.append({
            "source": tree_node['id'],
            "target": hub['id'],
            "type": "hub_link"
        })
        project_d3_graph(hub, nodes, links, parent_id=tree_node['id'], depth=depth + 1)

def project_llms_txt(tree_node, lines=None, level=0, base_url="https://mikelev.in"):
    """
    Projector C: Generates a Markdown Manifest (llms.txt) for AI Agents.
    Structure:
      - Intro / Brand / Instructions
      - Full Topology (Indented)
    """
    if lines is None:
        lines = []
        # --- PREAMBLE ---
        lines.append(f"# {tree_node['title']} - AI Context & Navigation Manifest")
        lines.append(f"> {tree_node['blurb']}")
        lines.append("")
        lines.append("## Information Architecture")
        lines.append("This site is organized using AI-driven K-Means clustering (Rule of 7).")
        lines.append("Content is load-balanced to prevent deep nesting.")
        lines.append("")
        lines.append("## Direct Data Access")
        lines.append(f"- **Full Graph Topology (JSON)**: [{base_url}/graph.json]({base_url}/graph.json) - Contains all nodes, links, and velocity data.")
        lines.append("- **Source Code**: Most articles offer `<link rel='alternate'>` to raw Markdown.")
        lines.append("")
        lines.append("## Site Map (High-Level)")
    
    indent = "  " * level
    
    # Hubs First (Navigation)
    for hub in tree_node.get('children_hubs', []):
        lines.append(f"{indent}- **[{hub['title']}]({base_url}{hub['permalink']}index.md?src=llms.txt)**")
        project_llms_txt(hub, lines, level + 1, base_url)

    # Articles (Content)
    for article in tree_node.get('children_articles', []):
        lines.append(f"{indent}- [{article['title']}]({base_url}{article['permalink']}index.md?src=llms.txt)")
        
    return lines

def project_xml_sitemaps(tree_node, base_url="https://mikelev.in", output_dir=Path(".")):
    """
    Projector D: Generates Top-Level Hub Sitemaps for load-balancing.
    """
    from datetime import datetime
    today = datetime.now().strftime('%Y-%m-%d')

    # 1. Core Sitemap (The VIP Lounge)
    core_urls = [
        f"  <url><loc>{base_url}/llms.txt</loc><lastmod>{today}</lastmod><changefreq>daily</changefreq><priority>1.0</priority></url>",
        f"  <url><loc>{base_url}/graph.json</loc><lastmod>{today}</lastmod><changefreq>daily</changefreq><priority>1.0</priority></url>",
        f"  <url><loc>{base_url}/feed.xml</loc><lastmod>{today}</lastmod><changefreq>daily</changefreq><priority>0.9</priority></url>",
        f"  <url><loc>{base_url}/</loc><lastmod>{today}</lastmod><changefreq>daily</changefreq><priority>0.9</priority></url>"
    ]
    
    hub_urls = []
    sitemap_index_urls = [f"  <sitemap><loc>{base_url}/sitemap-core.xml</loc></sitemap>"]
    sitemap_index_urls.append(f"  <sitemap><loc>{base_url}/sitemap-hubs.xml</loc></sitemap>")

    def traverse(node, post_list):
        # Process Hubs (Generated by K-Means)
        for hub in node.get('children_hubs', []):
            hub_urls.append(f"  <url><loc>{base_url}{hub['permalink']}</loc><lastmod>{today}</lastmod><changefreq>weekly</changefreq><priority>0.8</priority></url>")
            traverse(hub, post_list)
        
        # Process Articles (The Leaves)
        for article in node.get('children_articles', []):
            mod_date = article.get('date', today) or today
            post_list.append(f"  <url><loc>{base_url}{article['permalink']}</loc><lastmod>{mod_date}</lastmod><changefreq>monthly</changefreq><priority>0.7</priority></url>")

    def wrap_xml(urls):
        return '<?xml version="1.0" encoding="UTF-8"?>\n<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">\n' + "\n".join(urls) + '\n</urlset>'

    # 2. Generate Top-Level Hub Sitemaps
    for i, top_hub in enumerate(tree_node.get('children_hubs', [])):
        post_urls = []
        # Traverse this specific branch
        traverse(top_hub, post_urls)
        
        if post_urls:
            filename = f"sitemap-branch-{i}.xml"
            with open(output_dir / filename, 'w', encoding='utf-8') as f:
                f.write(wrap_xml(post_urls))
            sitemap_index_urls.append(f"  <sitemap><loc>{base_url}/{filename}</loc></sitemap>")

    def wrap_index(urls):
         return '<?xml version="1.0" encoding="UTF-8"?>\n<sitemapindex xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">\n' + "\n".join(urls) + '\n</sitemapindex>'

    return wrap_xml(core_urls), wrap_xml(hub_urls), wrap_index(sitemap_index_urls)

# --- MAIN EXECUTION ---

def main():
    print("🚀 Initializing Cartographer (Unified Graph Builder)...")
    parser = argparse.ArgumentParser()
    common.add_target_argument(parser)
    args = parser.parse_args()

    posts_dir = common.get_target_path(args)
    context_dir = posts_dir / "_context"

    if not context_dir.exists():
        print(f"❌ Context dir not found: {context_dir}")
        return

    # 1. LOAD DATA
    df = load_enriched_shards(context_dir, posts_dir)
    if df.empty:
        print("❌ No data found.")
        return
        
    velocity_data = load_velocity_data()

    # 2. BUILD CANONICAL TREE
    print(f"🧠 Clustering {len(df)} articles into Canonical Tree...")
    
    # Initialize Root Node
    canonical_tree = {
        "type": "root", 
        "id": "root",
        "title": "MikeLev.in", # Branding
        "permalink": "/",
        "blurb": "Python, NixOS, SEO, and AI Engineering.",
        "gravity": 60,
        "children_hubs": [],
        "children_articles": []
    }
    
    build_canonical_tree(df, canonical_tree, 0, velocity_data)

    # 3. EXPORT NAVGRAPH (JSON Tree for Jekyll)
    with open(NAVGRAPH_FILE, 'w', encoding='utf-8') as f:
        json.dump(canonical_tree, f, indent=2)
    print(f"✅ Generated NavGraph: {NAVGRAPH_FILE}")

    # 4. EXPORT GRAPH (Flat JSON for D3)
    nodes = []
    links = []
    project_d3_graph(canonical_tree, nodes, links)
    
    d3_data = {"nodes": nodes, "links": links}
    with open(GRAPH_FILE, 'w', encoding='utf-8') as f:
        json.dump(d3_data, f, indent=None) # Minified
    print(f"✅ Generated D3 Graph: {GRAPH_FILE} ({len(nodes)} nodes)")

    # 5. EXPORT LLMS.TXT (Markdown Manifest)
    # llms_lines = project_llms_txt(canonical_tree, base_url="https://mikelev.in")
    # with open(LLMS_TXT_FILE, 'w', encoding='utf-8') as f:
    #     f.write("\n".join(llms_lines))
    # print(f"✅ Generated LLMs.txt: {LLMS_TXT_FILE}")

    # 6. EXPORT XML SITEMAPS (The Sitrep)
    xml_core, xml_hubs, xml_index = project_xml_sitemaps(canonical_tree)
    with open(SITEMAP_CORE_FILE, 'w', encoding='utf-8') as f: f.write(xml_core)
    with open(SITEMAP_HUBS_FILE, 'w', encoding='utf-8') as f: f.write(xml_hubs)
    with open('sitemap.xml', 'w', encoding='utf-8') as f: f.write(xml_index)
    
    print(f"✅ Generated Sitemaps: Core, Hubs, Branches, and Root Index")

if __name__ == "__main__":
    main()
