import difflib
import json
from pathlib import Path
from datetime import datetime, timezone

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)

import re

def generate_diff_report(v1_ref: str, v2_ref: str, repo_path: Path, target: str, diff_text: str, save_dir: Path) -> Path:
    from marco.core.repository import resolve_version, list_tags
    
    repo_path = Path(repo_path)
    save_dir = Path(save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    
    v1_hash = resolve_version(v1_ref, repo_path)
    v2_hash = resolve_version(v2_ref, repo_path)
    
    tags = list_tags(repo_path)
    hash_to_tags = {}
    for t, h in tags.items():
        hash_to_tags.setdefault(h, []).append(t)
        
    def format_version(ref, h):
        tag_list = hash_to_tags.get(h, [])
        tag_to_use = ref if ref in tag_list else (tag_list[0] if tag_list else None)
        if tag_to_use:
            return f"{tag_to_use} ({h[:8]})"
        return f"{h[:8]}"
        
    v1_display = format_version(v1_ref, v1_hash)
    v2_display = format_version(v2_ref, v2_hash)
    
    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
    clean_diff = ansi_escape.sub('', diff_text)
    
    diff_lines = clean_diff.splitlines()
    
    added_lines = 0
    removed_lines = 0
    added_docs = []
    removed_docs = []
    
    formatted_diff_lines = []
    
    for line in diff_lines:
        if line.startswith('+++') or line.startswith('---'):
            pass
        elif line.startswith('+') and not line.startswith('+++'):
            added_lines += 1
            content = line[1:]
            
            if ' | ' in content and target == "preprocessed":
                parts = [p.strip() for p in content.split(' | ')]
                label = parts[1] if len(parts) > 1 else ""
                text = parts[2] if len(parts) > 2 else ""
            else:
                parts = content.split('\t', 1)
                if len(parts) > 1:
                    label = parts[0]
                    text = parts[1]
                else:
                    label = "`None`" if target == "config" else ""
                    text = parts[0]
                    
            preview = text[:60] + "..." if len(text) > 60 else text
            added_docs.append({'label': label, 'preview': preview})
            
        elif line.startswith('-') and not line.startswith('---'):
            removed_lines += 1
            content = line[1:]
            
            if ' | ' in content and target == "preprocessed":
                parts = [p.strip() for p in content.split(' | ')]
                label = parts[1] if len(parts) > 1 else ""
                text = parts[2] if len(parts) > 2 else ""
            else:
                parts = content.split('\t', 1)
                if len(parts) > 1:
                    label = parts[0]
                    text = parts[1]
                else:
                    label = "`None`" if target == "config" else ""
                    text = parts[0]
                    
            preview = text[:60] + "..." if len(text) > 60 else text
            removed_docs.append({'label': label, 'preview': preview})

        if len(line) > 90:
            formatted_line = line[:90] + "..."
        else:
            formatted_line = line
            
        formatted_diff_lines.append(formatted_line)
        if not line.startswith(('---', '+++')):
            formatted_diff_lines.append("")

    if formatted_diff_lines and formatted_diff_lines[-1] == "":
        formatted_diff_lines.pop()

    diff_target_name = "Raw Data (`raw.txt`)"
    if target == "preprocessed":
        diff_target_name = "Preprocessed Data (`preprocessed.tsv`)"
    elif target == "config":
        diff_target_name = "Config (`config.json`)"
        
    now_utc = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
    
    net_change = added_lines - removed_lines
    net_change_str = f"+{net_change}" if net_change > 0 else str(net_change)
    net_change_str += " documents"
    
    md = []
    md.append("# 📄 Diff Report")
    md.append("")
    md.append("---")
    md.append("")
    md.append("## Overview")
    md.append("")
    md.append("| | |")
    md.append("|---|---|")
    md.append(f"| **Base Version** | `{v1_display}` |")
    md.append(f"| **Target Version** | `{v2_display}` |")
    md.append(f"| **Diff Target** | {diff_target_name} |")
    md.append(f"| **Generated At** | {now_utc} |")
    md.append("")
    md.append("---")
    md.append("")
    md.append("## Summary of Changes")
    md.append("")
    md.append("| Metric | Value |")
    md.append("|---|---|")
    md.append(f"| ➕ Lines Added | {added_lines} |")
    md.append(f"| ➖ Lines Removed | {removed_lines} |")
    md.append(f"| 🔄 Net Change | {net_change_str} |")
    md.append("")
    md.append("---")
    md.append("")
    md.append("## Line-by-Line Diff")
    md.append("")
    md.append("```diff")
    if formatted_diff_lines:
        md.append("\n".join(formatted_diff_lines))
    md.append("```")
    md.append("")
    md.append("---")
    
    if added_docs:
        md.append("")
        md.append("## Added Documents")
        md.append("")
        md.append("| # | Label | Preview |")
        md.append("|---|---|---|")
        for i, doc in enumerate(added_docs, 1):
            lbl = doc['label'] or 'none'
            md.append(f"| {i} | `{lbl}` | {doc['preview']} |")
        md.append("")
        md.append("---")
        
    if removed_docs:
        md.append("")
        md.append("## Removed Documents")
        md.append("")
        md.append("| # | Label | Preview |")
        md.append("|---|---|---|")
        for i, doc in enumerate(removed_docs, 1):
            lbl = doc['label'] or 'none'
            md.append(f"| {i} | `{lbl}` | {doc['preview']} |")
        md.append("")
        md.append("---")
        
    md.append("")
    
    save_flag = f"--save {save_dir}" if str(save_dir) != "." else "--save"
    cmd = f"marco diff {v1_ref} {v2_ref} --target {target} {save_flag}"
    md.append(f"*Generated by **marco-dvcs** \u00b7 `{cmd}`*")
    
    filename = save_dir / f"diff_{v1_hash[:8]}_{v2_hash[:8]}.md"
    filename.write_text("\n".join(md), encoding="utf-8")
    
    return filename
