# scripts/pipeline_dag.py
import json
import hashlib
from collections import defaultdict, deque
from typing import List, Tuple, Dict, Callable, Any

# allowed libs - used elsewhere
import pandas as pd

# Type: data is list of (label, text)
DataRows = List[Tuple[str, str]]
NodeSpec = Dict[str, Any]
PipelineSpec = Dict[str, NodeSpec]

ENGLISH_STOPWORDS = {
    "a", "about", "above", "after", "again", "against", "all", "am", "an", "and", "any", "are", "aren't", "as", "at",
    "be", "because", "been", "before", "being", "below", "between", "both", "but", "by", "can't", "cannot", "could",
    "couldn't", "did", "didn't", "do", "does", "doesn't", "doing", "don't", "down", "during", "each", "few", "for",
    "from", "further", "had", "hadn't", "has", "hasn't", "have", "haven't", "having", "he", "he'd", "he'll", "he's",
    "her", "here", "here's", "hers", "herself", "him", "himself", "his", "how", "how's", "i", "i'd", "i'll", "i'm",
    "i've", "if", "in", "into", "is", "isn't", "it", "it's", "its", "itself", "let's", "me", "more", "most", "mustn't",
    "my", "myself", "no", "nor", "not", "of", "off", "on", "once", "only", "or", "other", "ought", "our", "ours",
    "ourselves", "out", "over", "own", "same", "shan't", "she", "she'd", "she'll", "she's", "should", "shouldn't", "so",
    "some", "such", "than", "that", "that's", "the", "their", "theirs", "them", "themselves", "then", "there", "there's",
    "these", "they", "they'd", "they'll", "they're", "they've", "this", "those", "through", "to", "too", "under", "until",
    "up", "very", "was", "wasn't", "we", "we'd", "we'll", "we're", "we've", "were", "weren't", "what", "what's", "when",
    "when's", "where", "where's", "which", "while", "who", "who's", "whom", "why", "why's", "with", "won't", "would",
    "wouldn't", "you", "you'd", "you'll", "you're", "you've", "your", "yours", "yourself", "yourselves"
}

# -------------------------
# Node implementations
# -------------------------
# Each node function takes DataRows and params, returns DataRows.
# Implementations MUST be deterministic.

def normalize_newlines_node(data: DataRows, params: dict) -> DataRows:
    out = []
    for lbl, txt in data:
        t = txt.replace('\r\n', '\n').replace('\r', '\n')
        out.append((lbl, t))
    return out

def lowercase_node(data: DataRows, params: dict) -> DataRows:
    enabled = params.get('enabled', True)
    if not enabled:
        return data
    return [(lbl, txt.lower()) for lbl, txt in data]

def uppercase_custom_node(data: DataRows, params: dict) -> DataRows:
    return [(lbl, txt.upper()) for lbl, txt in data]

def tokenize_node(data: DataRows, params: dict) -> DataRows:
    # Deterministic whitespace tokenization; store tokens as space-joined text
    method = params.get('method', 'whitespace')
    out = []
    for lbl, txt in data:
        if method == 'whitespace':
            tokens = [t for t in txt.strip().split() if t]
        else:
            # fallback: whitespace
            tokens = [t for t in txt.strip().split() if t]
        out.append((lbl, ' '.join(tokens)))
    return out

def filter_length_node(data: DataRows, params: dict) -> DataRows:
    min_t = int(params.get('min_tokens', 0))
    max_t = int(params.get('max_tokens', 10_000_000))
    out = []
    for lbl, txt in data:
        n = len(txt.split()) if txt else 0
        if min_t <= n <= max_t:
            out.append((lbl, txt))
    return out

def remove_stopwords_node(data: DataRows, params: dict) -> DataRows:
    language = params.get('language', 'english')
    # only support english built-in for now
    if language != 'english':
        return data
        
    out = []
    for lbl, txt in data:
        tokens = txt.split()
        filtered = [t for t in tokens if t.lower() not in ENGLISH_STOPWORDS]
        out.append((lbl, ' '.join(filtered)))
    return out

def deduplicate_node(data: DataRows, params: dict) -> DataRows:
    # Keep first occurrence deterministically
    seen = set()
    out = []
    method = params.get('method', 'exact')
    for lbl, txt in data:
        key = txt if method == 'exact' else hashlib.sha256(txt.encode('utf-8')).hexdigest()
        if key in seen:
            continue
        seen.add(key)
        out.append((lbl, txt))
    return out

# node registry maps string name to function
_NODE_REGISTRY: Dict[str, Callable[[DataRows, dict], DataRows]] = {
    'normalize_newlines': normalize_newlines_node,
    'lowercase': lowercase_node,
    'uppercase_custom': uppercase_custom_node,
    'tokenize': tokenize_node,
    'filter_length': filter_length_node,
    'remove_stopwords': remove_stopwords_node,
    'deduplicate': deduplicate_node,
}

# -------------------------
# DAG helpers
# -------------------------
def validate_dag(spec: PipelineSpec):
    # Check for cycles using DFS detection
    visited = {}
    def dfs(node):
        if node not in spec:
            raise ValueError(f"Node '{node}' is declared as dependency but missing from pipeline spec")
        if visited.get(node) == 'visiting':
            raise ValueError(f"Cycle detected at node: {node}")
        if visited.get(node) == 'visited':
            return
        visited[node] = 'visiting'
        for dep in spec[node].get('depends_on', []):
            dfs(dep)
        visited[node] = 'visited'
    for n in spec:
        if n not in visited:
            dfs(n)
    return True

def topological_sort_order(spec: PipelineSpec) -> List[str]:
    # Kahn's algorithm (deterministic: sort nodes to break ties)
    indeg = {n: 0 for n in spec}
    deps = defaultdict(list)
    for n, s in spec.items():
        for d in s.get('depends_on', []):
            indeg[n] += 1
            deps[d].append(n)
    # start with indeg 0 nodes sorted
    q = deque(sorted([n for n, d in indeg.items() if d == 0]))
    order = []
    while q:
        n = q.popleft()
        order.append(n)
        for child in sorted(deps.get(n, [])):
            indeg[child] -= 1
            if indeg[child] == 0:
                q.append(child)
    if len(order) != len(spec):
        raise ValueError("Cycle or missing nodes detected in pipeline spec")
    return order

# -------------------------
# Metrics helper
# -------------------------
def compute_metrics_for_rows(rows: DataRows) -> dict:
    n_docs = len(rows)
    total_tokens = 0
    vocab = {}
    n_empty = 0
    for lbl, txt in rows:
        tokens = txt.split() if txt else []
        if not tokens:
            n_empty += 1
        total_tokens += len(tokens)
        for t in tokens:
            vocab[t] = vocab.get(t, 0) + 1
    return {
        'n_documents': int(n_docs),
        'n_tokens': int(total_tokens),
        'vocab_size': int(len(vocab)),
        'avg_tokens_per_doc': float(total_tokens / n_docs) if n_docs else 0.0,
        'median_tokens_per_doc': float(pd.Series([len(txt.split()) for _, txt in rows]).median()) if n_docs else 0.0,
        'n_empty_docs': int(n_empty),
        'top_tokens': [{'token': t, 'count': c} for t, c in sorted(vocab.items(), key=lambda x: -x[1])[:50]],
    }

# -------------------------
# Runner
# -------------------------
def run_pipeline_dag(raw_rows: DataRows, pipeline_spec: PipelineSpec) -> (DataRows, dict, dict):
    """
    raw_rows: list of (label, text)
    pipeline_spec: dict mapping node_name -> { 'func': <string key in registry>, 'params': {...}, 'depends_on': [...] }
    Returns:
      final_rows: DataRows (concatenation of final nodes' outputs)
      metrics_per_step: dict mapping node_name -> metrics
      dag_structure: minimal structure (nodes + depends_on)
    """

    # Validate and get order
    validate_dag(pipeline_spec)
    order = topological_sort_order(pipeline_spec)

    results: Dict[str, DataRows] = {}
    metrics_per_step = {}

    # raw is the canonical input name
    results['raw'] = raw_rows

    # run nodes in deterministic order
    for node in order:
        node_spec = pipeline_spec[node]
        func_key = node_spec.get('func')
        params = node_spec.get('params', {})
        depends = node_spec.get('depends_on', [])

        # gather inputs deterministically: concat in order of depends list; if no deps use raw
        if depends:
            inputs = []
            for d in depends:
                # If dependency not run (shouldn't happen), raise
                if d not in results:
                    raise RuntimeError(f"Dependency output missing for {d}")
                inputs.extend(results[d])
        else:
            inputs = results['raw']

        # get function
        if func_key not in _NODE_REGISTRY:
            raise ValueError(f"Unknown node func '{func_key}' for node '{node}'")
        func = _NODE_REGISTRY[func_key]
        out_rows = func(inputs, params)

        # store result and metrics
        results[node] = out_rows
        metrics_per_step[node] = compute_metrics_for_rows(out_rows)

    # Determine final nodes: nodes that are not depended upon by any other node
    depended = set()
    for n, s in pipeline_spec.items():
        for d in s.get('depends_on', []):
            depended.add(d)
    final_nodes = sorted([n for n in pipeline_spec.keys() if n not in depended])
    # Deterministic final output: concat outputs of final nodes in sorted order
    final_output: DataRows = []
    if final_nodes:
        for fn in final_nodes:
            final_output.extend(results[fn])
    else:
        # fallback: last node in order
        final_output = results[order[-1]]

    dag_structure = {n: {'depends_on': pipeline_spec[n].get('depends_on', []), 'func': pipeline_spec[n].get('func'), 'params': pipeline_spec[n].get('params', {})} for n in pipeline_spec}

    return final_output, metrics_per_step, dag_structure

# -------------------------
# Serialization helpers for canonical bytes (for output hash)
# -------------------------
def canonical_preprocessed_bytes_from_rows(rows: DataRows) -> bytes:
    # deterministic columns: label, text, n_tokens
    import io
    import csv
    buf = io.StringIO()
    writer = csv.writer(buf, delimiter='\t', lineterminator='\n', quoting=csv.QUOTE_MINIMAL)
    writer.writerow(['label', 'text', 'n_tokens'])
    for lbl, txt in rows:
        writer.writerow([lbl, txt, str(len(txt.split()) if txt else 0)])
    return buf.getvalue().encode('utf-8')

def compute_output_hash_from_rows(rows: DataRows) -> str:
    b = canonical_preprocessed_bytes_from_rows(rows)
    return hashlib.sha256(b).hexdigest()
