import difflib
import json
from pathlib import Path
from datetime import datetime

import shutil
import textwrap

import pandas as pd
from marco.core.repository import resolve_version

def load_vocab_from_preprocessed(prep_path: Path):
    if not prep_path.exists():
        return set()
    df = pd.read_csv(prep_path, sep='\t', encoding='utf-8')
    tokens = set()
    for t in df['text'].fillna(''):
        tokens.update(t.split())
    return tokens

def compare_versions(v1_ref: str, v2_ref: str, repo_path: Path) -> dict:
    repo_path = Path(repo_path)
    v1 = resolve_version(v1_ref, repo_path)
    v2 = resolve_version(v2_ref, repo_path)
    
    v1_dir = repo_path / '.marco' / 'versions' / v1
    v2_dir = repo_path / '.marco' / 'versions' / v2
    
    m1 = json.loads((v1_dir / 'metrics.json').read_text(encoding='utf-8'))
    m2 = json.loads((v2_dir / 'metrics.json').read_text(encoding='utf-8'))
    
    report = {
        'v1_id': v1,
        'v2_id': v2,
        'v1_ref': v1_ref,
        'v2_ref': v2_ref,
        'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    # Metric diffs
    report['metrics'] = {}
    for k in ['n_documents', 'n_tokens', 'vocab_size']:
        a = m1.get(k, 0)
        b = m2.get(k, 0)
        diff = b - a
        pct = f"{(diff / a * 100):.2f}%" if a else "N/A"
        report['metrics'][k] = {'v1': a, 'v2': b, 'diff': diff, 'pct_change': pct}
        
    # Vocab overlap
    t1 = load_vocab_from_preprocessed(v1_dir / 'preprocessed.tsv')
    t2 = load_vocab_from_preprocessed(v2_dir / 'preprocessed.tsv')
    inter = t1 & t2
    union = t1 | t2
    jaccard = len(inter) / len(union) if union else 1.0
    
    report['vocab'] = {
        'v1_size': len(t1), 
        'v2_size': len(t2), 
        'intersection': len(inter), 
        'union': len(union), 
        'jaccard': f"{jaccard:.4f}"
    }
    
    # Configuration difference
    c1 = json.loads((v1_dir / 'config.json').read_text(encoding='utf-8'))
    c2 = json.loads((v2_dir / 'config.json').read_text(encoding='utf-8'))
    report['config_changed'] = c1 != c2
    
    return report

def generate_markdown_report(report: dict, out_dir: Path) -> Path:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    
    v1_short = report['v1_id'][:8]
    v2_short = report['v2_id'][:8]
    filename = out_dir / f"marco_diff_{v1_short}_vs_{v2_short}.md"
    
    md = [
        f"# Marco Dataset Comparison Report",
        f"**Generated:** {report['timestamp']}\n",
        f"| | Version 1 | Version 2 |",
        f"|---|---|---|",
        f"| **Reference** | `{report['v1_ref']}` | `{report['v2_ref']}` |",
        f"| **Hash ID** | `{v1_short}` | `{v2_short}` |\n",
        f"## 📊 Metrics Difference\n",
        f"| Metric | V1 | V2 | Diff | % Change |",
        f"|---|---|---|---|---|"
    ]
    
    for k, v in report['metrics'].items():
        name = k.replace('_', ' ').title()
        sign = "+" if v['diff'] > 0 else ""
        md.append(f"| **{name}** | {v['v1']} | {v['v2']} | {sign}{v['diff']} | {v['pct_change']} |")
        
    md.append(f"\n## 🔠 Vocabulary Overlap\n")
    vcb = report['vocab']
    md.extend([
        f"- **V1 Unique Terms**: {vcb['v1_size']}",
        f"- **V2 Unique Terms**: {vcb['v2_size']}",
        f"- **Shared Terms (Intersection)**: {vcb['intersection']}",
        f"- **Total Terms (Union)**: {vcb['union']}",
        f"- **Jaccard Similarity**: `{vcb['jaccard']}`\n"
    ])
    
    if report['config_changed']:
        md.append("## ⚠️ Warning: Pipeline Configurations Differ")
        md.append("The DAG preprocessing steps used to generate these two versions are different.")
    else:
        md.append("## ✅ Pipeline Configurations Match")
        md.append("Both versions were generated using the exact same DAG preprocessing steps.")
        
    filename.write_text("\n".join(md), encoding='utf-8')
    return filename


# ---------------------------------------------------------------------------
# Unified diff  (--- / +++ format)
# ---------------------------------------------------------------------------

_DIFF_TARGETS = {
    "raw":          "raw.txt",
    "preprocessed": "preprocessed.tsv",
    "config":       "config.json",
}

def _format_tabular_diff(f1: Path, f2: Path, v1_ref: str, v2_ref: str) -> str:
    text1 = f1.read_text(encoding="utf-8", errors="replace").splitlines()
    text2 = f2.read_text(encoding="utf-8", errors="replace").splitlines()

    if text1 and text2 and text1[0] == text2[0] and text1[0].startswith("label\ttext"):
        text1 = text1[1:]
        text2 = text2[1:]

    diff_generator = difflib.ndiff(text1, text2)

    term_width = shutil.get_terminal_size((80, 20)).columns
    label_width = 10
    tokens_width = 8
    text_width = term_width - label_width - tokens_width - 15  # 15 for index + decorators
    if text_width < 20:
        text_width = 20

    out = []
    header = f"{'Row':<6} | {'Label':<{label_width}} | {'Text':<{text_width}} | {'Tokens':<{tokens_width}}"
    
    # Prefix the table with exactly what files are being diffed
    out.append(f"--- {v1_ref} ({f1.parent.name[:8]})/{f1.name}")
    out.append(f"+++ {v2_ref} ({f2.parent.name[:8]})/{f2.name}")
    out.append("=" * len(header))
    out.append(header)
    out.append("-" * len(header))

    v1_idx = 1
    v2_idx = 1
    changes_found = False

    for line in diff_generator:
        prefix = line[:2]
        content = line[2:]

        if prefix == "  ":
            v1_idx += 1
            v2_idx += 1
            continue
        elif prefix == "? ":
            continue
            
        changes_found = True

        idx_str = f"-{v1_idx}" if prefix == "- " else f"+{v2_idx}"
        color_start = "\033[31m" if prefix == "- " else "\033[32m"
        color_end = "\033[0m"

        if prefix == "- ":
            v1_idx += 1
        elif prefix == "+ ":
            v2_idx += 1

        parts = content.split('\t')
        if len(parts) >= 3:
            label, text, tokens = parts[0], parts[1], parts[2]
        elif len(parts) == 2:
            label, text, tokens = parts[0], parts[1], ""
        else:
            label, text, tokens = "", content, ""

        wrapped_text = textwrap.wrap(text, width=text_width)
        if not wrapped_text:
            wrapped_text = [""]

        out.append(f"{idx_str:<6} | {color_start}{label[:label_width]:<{label_width}}{color_end} | {color_start}{wrapped_text[0]:<{text_width}}{color_end} | {color_start}{tokens[:tokens_width]:<{tokens_width}}{color_end}")

        for w in wrapped_text[1:]:
            out.append(f"{' ':<6} | {' ':<{label_width}} | {color_start}{w:<{text_width}}{color_end} | {' ':<{tokens_width}}")

    if not changes_found:
        return ""

    return "\n".join(out) + "\n"

def diff_versions_text(
    v1_ref: str,
    v2_ref: str,
    repo_path: Path,
    target: str = "raw",
    context: int = 3,
) -> str:
    """
    Return a unified diff string (--- / +++ format) comparing a specific file
    between two versions.

    Parameters
    ----------
    v1_ref : str
        First version reference (full hash, short hash, or tag).
    v2_ref : str
        Second version reference (full hash, short hash, or tag).
    repo_path : Path
        Root of the marco repository (directory containing .marco/).
    target : str
        Which file inside the version folder to diff.
        Choices: ``"raw"`` (raw.txt), ``"preprocessed"`` (preprocessed.tsv),
        ``"config"`` (config.json).  Defaults to ``"raw"``.
    context : int
        Number of unchanged lines to show around each hunk (default 3,
        identical to ``git diff``).

    Returns
    -------
    str
        The unified diff string, or an empty string if files are identical.
    """
    if target not in _DIFF_TARGETS:
        raise ValueError(
            f"Unknown diff target '{target}'. "
            f"Valid choices: {list(_DIFF_TARGETS.keys())}"
        )

    repo_path = Path(repo_path)
    v1 = resolve_version(v1_ref, repo_path)
    v2 = resolve_version(v2_ref, repo_path)

    v1_dir = repo_path / ".marco" / "versions" / v1
    v2_dir = repo_path / ".marco" / "versions" / v2

    filename = _DIFF_TARGETS[target]
    f1 = v1_dir / filename
    f2 = v2_dir / filename

    if target == "preprocessed":
        return _format_tabular_diff(f1, f2, v1_ref, v2_ref)

    if target == "config":
        obj1 = json.loads(f1.read_text(encoding="utf-8", errors="replace"))
        obj2 = json.loads(f2.read_text(encoding="utf-8", errors="replace"))
        lines1 = (json.dumps(obj1, indent=2) + "\n").splitlines(keepends=True)
        lines2 = (json.dumps(obj2, indent=2) + "\n").splitlines(keepends=True)
    else:
        lines1 = f1.read_text(encoding="utf-8", errors="replace").splitlines(keepends=True)
        lines2 = f2.read_text(encoding="utf-8", errors="replace").splitlines(keepends=True)

    # Human-readable labels for the diff header lines
    label1 = f"{v1_ref} ({v1[:8]})/{filename}"
    label2 = f"{v2_ref} ({v2[:8]})/{filename}"

    diff_lines = list(difflib.unified_diff(
        lines1,
        lines2,
        fromfile=label1,
        tofile=label2,
        n=context,
    ))

    # Apply basic ANSI terminal colors
    colored_diff = []
    for line in diff_lines:
        ending = "\n" if line.endswith("\n") else ""
        text = line.rstrip("\n")
        
        if text.startswith("--- ") or text.startswith("+++ ") or text.startswith("@@ "):
            colored_diff.append(f"\033[36m{text}\033[0m{ending}") # Cyan
        elif text.startswith("-"):
            colored_diff.append(f"\033[31m{text}\033[0m{ending}") # Red
        elif text.startswith("+"):
            colored_diff.append(f"\033[32m{text}\033[0m{ending}") # Green
        else:
            colored_diff.append(line)

    return "".join(colored_diff)
