import hashlib
import json
import os
import shutil
import tempfile
import tarfile
import statistics
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Union

import pandas as pd

from marco.core.preprocessor import run_pipeline_dag, compute_output_hash_from_rows, compute_metrics_for_rows
from marco.core.locker import FileLock
from marco.token_analytics.distribution import compute_token_frequencies, save_token_frequencies

def init_repo(repo_path: Union[str, Path]):
    repo_path = Path(repo_path)
    marco_dir = repo_path / '.marco'
    
    if marco_dir.exists():
        raise RuntimeError("This directory is already a marco project. Run 'marco list' to see existing versions.")
        
    (marco_dir / 'versions').mkdir(parents=True, exist_ok=True)
    (marco_dir / 'locks').mkdir(parents=True, exist_ok=True)
    (repo_path / 'raw').mkdir(parents=True, exist_ok=True)
    
    refs_file = marco_dir / 'refs.json'
    if not refs_file.exists():
        refs_file.write_text('{}', encoding='utf-8')
        
    import uuid
    import sys
    from datetime import datetime, timezone
    
    root_id = f"marco_r_{uuid.uuid4().hex[:6]}"
    
    root_node = {
        "node_type": "ROOT",
        "root_id": root_id,
        "project_name": repo_path.resolve().name,
        "initialized_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "python_version": sys.version.split()[0],
        "marco_version": "0.1.37", 
        "children": []
    }
    
    registry = {
        "root_id": root_id,
        "total_versions": 0,
        "versions": []
    }
    
    chains = {
        "chains": {}
    }
    
    workspace = {
        "project_name": repo_path.resolve().name,
        "version": "1.0"
    }
    
    (marco_dir / 'root.json').write_text(json.dumps(root_node, indent=2), encoding='utf-8')
    (marco_dir / 'registry.json').write_text(json.dumps(registry, indent=2), encoding='utf-8')
    (marco_dir / 'chains.json').write_text(json.dumps(chains, indent=2), encoding='utf-8')
    (marco_dir / 'workspace.json').write_text(json.dumps(workspace, indent=2), encoding='utf-8')
    
    print(f"\n✅ Initialized marco project '{root_node['project_name']}'\n")
    print("  .marco/          project metadata")
    print("  .marco/versions/ dataset versions will be stored here")
    print("  raw/             place your raw data here\n")
    print("Next steps:\n  1. Add your raw data to raw/\n  2. Create a pipeline config\n  3. Run: marco upload raw/data.txt")
    
    return marco_dir

def canonical_json(obj) -> str:
    return json.dumps(obj, sort_keys=True, separators=(',', ':'))

def compute_raw_hash(raw_path: Path) -> str:
    hasher = hashlib.sha256()
    with raw_path.open('rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            hasher.update(chunk.replace(b"\r\n", b"\n"))
    return hasher.hexdigest()

def create_version(raw_path: Path, config: dict, repo_path: Path, user: str = "unknown", tags: List[str] = None, parents: List[str] = None, results_path: str = None):
    tags = tags or []
    parents = parents or []
    repo_path = Path(repo_path)
    marco_dir = repo_path / '.marco'
    versions_dir = marco_dir / 'versions'
    locks_dir = marco_dir / 'locks'
    
    if not marco_dir.exists():
        raise RuntimeError(f"Marco repository not initialized at {repo_path}. Run 'marco init' first.")
        
    raw_hash = compute_raw_hash(raw_path)
    canonical = canonical_json(config)
    version_id = hashlib.sha256((raw_hash + ':' + canonical).encode('utf-8')).hexdigest()
    version_folder = versions_dir / version_id
    
    with FileLock(locks_dir, "create_version"):
        # --- ROOT NODE & REGISTRY LOGIC ---
        root_path = marco_dir / 'root.json'
        registry_path = marco_dir / 'registry.json'
        chains_path = marco_dir / 'chains.json'
        
        # Auto-migrate old repos
        if not root_path.exists():
            import uuid, sys
            r_id = f"marco_r_{uuid.uuid4().hex[:6]}"
            root_path.write_text(json.dumps({
                "node_type": "ROOT", "root_id": r_id, "project_name": repo_path.resolve().name,
                "initialized_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
                "python_version": sys.version.split()[0], "marco_version": "0.1.37", "children": []
            }, indent=2), encoding='utf-8')
            registry_path.write_text(json.dumps({
                "root_id": r_id, "total_versions": 0, "versions": []
            }, indent=2), encoding='utf-8')
            chains_path.write_text(json.dumps({"chains": {}}, indent=2), encoding='utf-8')
            
        root_data = json.loads(root_path.read_text(encoding='utf-8'))
        registry = json.loads(registry_path.read_text(encoding='utf-8'))
        chains_data = json.loads(chains_path.read_text(encoding='utf-8'))
        
        root_id = root_data['root_id']
        chain_id = None
        
        for cid, cinfo in chains_data.get('chains', {}).items():
            if cinfo.get('raw_data_hash') == raw_hash:
                chain_id = cid
                break
                
        if not chain_id:
            import string
            chars = string.ascii_uppercase
            chain_idx = len(chains_data.get('chains', {}))
            letter = chars[chain_idx % 26] if chain_idx < 26 else f"{chars[chain_idx % 26]}{chain_idx // 26}"
            chain_id = f"chain_{letter}"
            if 'chains' not in chains_data: chains_data['chains'] = {}
            chains_data['chains'][chain_id] = {"raw_data_hash": raw_hash}
            
            if chain_id not in root_data.get('children', []):
                root_data.setdefault('children', []).append(chain_id)
                root_path.write_text(json.dumps(root_data, indent=2), encoding='utf-8')
            chains_path.write_text(json.dumps(chains_data, indent=2), encoding='utf-8')
            
        parent_version_id = None
        chain_versions = [v for v in registry.get('versions', []) if v.get('chain_id') == chain_id]
        if chain_versions:
            chain_versions.sort(key=lambda x: x.get('created_at', ''), reverse=True)
            parent_version_id = chain_versions[0].get('version_id')
            
        if parent_version_id and parent_version_id not in parents:
            parents.append(parent_version_id)
            
        if version_folder.exists():
            print(f"Version {version_id} already exists.")
            if tags:
                tag_version(version_id, tags, repo_path)
            return version_id
            
        text = raw_path.read_text(encoding='utf-8', errors='replace')
        text = text.replace('\r\n', '\n').replace('\r', '\n')
        lines = [ln for ln in text.split('\n') if ln.strip()]
        
        raw_rows = []
        for ln in lines:
            if '\t' in ln:
                lbl, txt = ln.split('\t', 1)
            else:
                lbl, txt = '', ln
            raw_rows.append((lbl.strip(), txt.strip()))

        doc_lengths = []
        vocab = set()
        empty_docs = 0
        label_dist = {}
        has_labels = False

        for lbl, txt in raw_rows:
            if lbl:
                has_labels = True
                label_dist[lbl] = label_dist.get(lbl, 0) + 1
            
            if not txt or not txt.strip():
                empty_docs += 1
                doc_lengths.append(0)
            else:
                tokens = txt.split()
                doc_lengths.append(len(tokens))
                vocab.update(tokens)
                
        if not has_labels:
            label_dist = None
            
        total_docs = len(doc_lengths)
        total_tokens = sum(doc_lengths)
        
        if total_docs > 0:
            avg_len = round(total_tokens / total_docs, 2)
            median_len = float(statistics.median(doc_lengths))
            min_len = min(doc_lengths)
            max_len = max(doc_lengths)
            std_len = round(statistics.stdev(doc_lengths), 2) if total_docs > 1 else 0.0
        else:
            avg_len = median_len = min_len = max_len = std_len = 0.0
            
        raw_stats = {
            "total_documents": total_docs,
            "total_tokens": total_tokens,
            "vocabulary_size": len(vocab),
            "avg_doc_length": avg_len,
            "median_doc_length": median_len,
            "std_doc_length": std_len,
            "min_doc_length": min_len,
            "max_doc_length": max_len,
            "empty_documents": empty_docs,
            "label_distribution": label_dist
        }

        pipeline_spec = config.get('dag') or {}
        if not pipeline_spec:
            raise ValueError("Config must include 'dag' mapping of nodes")

        final_rows, metrics_per_step, dag_structure = run_pipeline_dag(raw_rows, pipeline_spec)
        
        final_metrics = compute_metrics_for_rows(final_rows)
        output_hash = compute_output_hash_from_rows(final_rows)
        final_metrics['output_hash'] = output_hash
        final_metrics['created_by'] = user
        
        df = pd.DataFrame([{'label': l, 'text': t, 'n_tokens': len(t.split()) if t else 0} for l, t in final_rows])
        
        manifest = {
            'version_id': version_id,
            'root_id': root_id,
            'chain_id': chain_id,
            'parent_version_id': parent_version_id,
            'created_at': datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            'created_by': user,
            'raw_hash': raw_hash,
            'config_hash': hashlib.sha256(canonical.encode('utf-8')).hexdigest(),
            'n_raw_lines': len(lines),
            'parents': parents,
            'output_hash': output_hash,
            'dag': dag_structure,
            'raw_stats_path': 'raw_stats.json',
            'post_stats_path': 'post_stats.json'
        }
        
        tmp = Path(tempfile.mkdtemp(prefix='version_tmp_', dir=versions_dir))
        try:
            all_tokens = []
            for lbl, txt in final_rows:
                if txt:
                    all_tokens.extend(txt.split())
            if all_tokens:
                freqs = compute_token_frequencies(all_tokens)
                save_token_frequencies(freqs, tmp)

            shutil.copy2(raw_path, tmp / 'raw.txt')
            if results_path:
                res_path = Path(results_path)
                if res_path.exists():
                    shutil.copy2(res_path, tmp / f'training_results{res_path.suffix}')
            (tmp / 'config.json').write_text(canonical, encoding='utf-8')
            (tmp / 'raw_hash.txt').write_text(raw_hash, encoding='utf-8')
            df.to_csv(tmp / 'preprocessed.tsv', sep='\t', index=False, header=True, encoding='utf-8', lineterminator='\n')
            final_metrics['metrics_per_step'] = metrics_per_step
            (tmp / 'raw_stats.json').write_text(json.dumps(raw_stats, indent=2), encoding='utf-8')
            (tmp / 'metrics.json').write_text(json.dumps(final_metrics, indent=2), encoding='utf-8')
            
            final_doc_lengths = [len(txt.split()) if txt and txt.strip() else 0 for lbl, txt in final_rows]
            std_len = statistics.stdev(final_doc_lengths) if len(final_doc_lengths) > 1 else 0.0
            
            raw_docs = raw_stats.get('total_documents', 0)
            raw_tokens = raw_stats.get('total_tokens', 0)
            raw_vocab = raw_stats.get('vocabulary_size', 0)
            fin_docs = final_metrics.get('n_documents', 0)
            fin_tokens = final_metrics.get('n_tokens', 0)
            fin_vocab = final_metrics.get('vocab_size', 0)
            
            post_stats = {
                "total_documents": fin_docs,
                "documents_removed": raw_docs - fin_docs,
                "documents_removed_pct": round(((raw_docs - fin_docs) / raw_docs) * 100, 2) if raw_docs else 0.0,
                "total_tokens": fin_tokens,
                "token_reduction_pct": round(((raw_tokens - fin_tokens) / raw_tokens) * 100, 2) if raw_tokens else 0.0,
                "vocabulary_size": fin_vocab,
                "vocabulary_reduction_pct": round(((raw_vocab - fin_vocab) / raw_vocab) * 100, 2) if raw_vocab else 0.0,
                "avg_doc_length": round(final_metrics.get('avg_tokens_per_doc', 0.0), 2),
                "median_doc_length": round(final_metrics.get('median_tokens_per_doc', 0.0), 2),
                "std_doc_length": round(std_len, 2),
                "min_doc_length": min(final_doc_lengths) if final_doc_lengths else 0,
                "max_doc_length": max(final_doc_lengths) if final_doc_lengths else 0
            }
            (tmp / 'post_stats.json').write_text(json.dumps(post_stats, indent=2), encoding='utf-8')

            (tmp / 'dag_structure.json').write_text(json.dumps(dag_structure, indent=2), encoding='utf-8')
            (tmp / 'manifest.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')
            
            v_idx = registry['total_versions'] + 1
            registry['versions'].append({
                "version_id": version_id,
                "chain_id": chain_id,
                "raw_data_hash": raw_hash,
                "created_at": manifest['created_at'],
                "label": f"V{v_idx}"
            })
            registry['total_versions'] = len(registry['versions'])
            registry_path.write_text(json.dumps(registry, indent=2), encoding='utf-8')
            
            os.replace(str(tmp), str(version_folder))
        except Exception:
            shutil.rmtree(tmp, ignore_errors=True)
            raise

    if tags:
        tag_version(version_id, tags, repo_path)
        
    return version_id

def tag_version(version_id: str, tags: List[str], repo_path: Path):
    repo_path = Path(repo_path)
    refs_file = repo_path / '.marco' / 'refs.json'
    locks_dir = repo_path / '.marco' / 'locks'
    
    with FileLock(locks_dir, "refs_update"):
        refs = {}
        if refs_file.exists():
            try:
                text = refs_file.read_text(encoding='utf-8').strip()
                if text:
                    refs = json.loads(text)
            except json.JSONDecodeError:
                pass
        for t in tags:
            refs[t] = version_id
        refs_file.write_text(json.dumps(refs, indent=2), encoding='utf-8')

def list_tags(repo_path: Path) -> Dict[str, str]:
    refs_file = Path(repo_path) / '.marco' / 'refs.json'
    if refs_file.exists():
        try:
            text = refs_file.read_text(encoding='utf-8').strip()
            if text:
                return json.loads(text)
        except json.JSONDecodeError:
            pass
    return {}

def resolve_version(ref: str, repo_path: Path) -> str:
    repo_path = Path(repo_path)
    versions_dir = repo_path / '.marco' / 'versions'
    
    if (versions_dir / ref).exists():
        return ref
        
    tags = list_tags(repo_path)
    if ref in tags:
        return tags[ref]
        
    matches = [d.name for d in versions_dir.iterdir() if d.is_dir() and d.name.startswith(ref)]
    if len(matches) == 1:
        return matches[0]
    elif len(matches) > 1:
        raise ValueError(f"Ambiguous short hash '{ref}'. Matches: {matches}")
        
    raise ValueError(f"Version '{ref}' not found.")

def list_versions(repo_path: Path) -> List[dict]:
    versions_dir = Path(repo_path) / '.marco' / 'versions'
    if not versions_dir.exists():
        return []
    
    tags = list_tags(repo_path)
    tag_map = {}
    for t, v in tags.items():
        tag_map.setdefault(v, []).append(t)
        
    out = []
    for d in versions_dir.iterdir():
        if d.is_dir() and (d / 'manifest.json').exists():
            man = json.loads((d / 'manifest.json').read_text(encoding='utf-8'))
            man['tags'] = tag_map.get(d.name, [])
            out.append(man)
    out.sort(key=lambda x: x.get('created_at', ''), reverse=True)
    return out

def export_version(version_ref: str, dest_dir: Path, repo_path: Path):
    vid = resolve_version(version_ref, repo_path)
    version_dir = Path(repo_path) / '.marco' / 'versions' / vid
    dest_dir = Path(dest_dir)
    dest_dir.mkdir(parents=True, exist_ok=True)
    
    tar_path = dest_dir / f"marco_version_{vid[:8]}.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tar:
        tar.add(version_dir, arcname=vid)
    return tar_path

def import_version(tar_path: Path, repo_path: Path):
    tar_path = Path(tar_path)
    repo_path = Path(repo_path)
    versions_dir = repo_path / '.marco' / 'versions'
    locks_dir = repo_path / '.marco' / 'locks'
    
    if not versions_dir.exists():
        raise RuntimeError("Marco repo not initialized.")
        
    with FileLock(locks_dir, "import_version"):
        with tempfile.TemporaryDirectory() as tmpdir:
            with tarfile.open(tar_path, "r:gz") as tar:
                for member in tar.getmembers():
                    if member.name.startswith('/') or '..' in member.name:
                        raise ValueError("Tarball contains illegal paths.")
                tar.extractall(path=tmpdir)
                
            extracted = list(Path(tmpdir).iterdir())
            if len(extracted) != 1 or not extracted[0].is_dir() or len(extracted[0].name) != 64:
                raise ValueError("Invalid marco tarball format.")
                
            vid = extracted[0].name
            dest_dir = versions_dir / vid
            if dest_dir.exists():
                print(f"Version {vid} already exists.")
                return vid
                
            shutil.copytree(extracted[0], dest_dir)
            return vid

def restore_version(version_ref: str, dest_path: Path, repo_path: Path):
    vid = resolve_version(version_ref, repo_path)
    version_dir = Path(repo_path) / '.marco' / 'versions' / vid
    dest_path = Path(dest_path)
    
    if not dest_path.parent.exists():
        dest_path.parent.mkdir(parents=True, exist_ok=True)
        
    source_processed = version_dir / 'preprocessed.tsv'
    if not source_processed.exists():
        raise FileNotFoundError(f"Processed file not found in version {vid}")
        
    shutil.copy2(source_processed, dest_path)
    return dest_path

def attach_results(version_ref: str, results_path: str, repo_path: Path):
    vid = resolve_version(version_ref, repo_path)
    version_dir = Path(repo_path) / '.marco' / 'versions' / vid
    res_path = Path(results_path)
    
    if not res_path.exists():
        raise FileNotFoundError(f"Results file {results_path} not found.")

    locks_dir = Path(repo_path) / '.marco' / 'locks'
    
    with FileLock(locks_dir, f"attach_results_{vid}"):
        dest_file = version_dir / f"training_results{res_path.suffix}"
        shutil.copy2(res_path, dest_file)
    
    return vid, dest_file

def delete_version(version_ref: str, repo_path: Path):
    vid = resolve_version(version_ref, repo_path)
    repo_path = Path(repo_path)
    marco_dir = repo_path / '.marco'
    versions_dir = marco_dir / 'versions'
    locks_dir = marco_dir / 'locks'
    refs_file = marco_dir / 'refs.json'
    
    version_dir = versions_dir / vid
    if not version_dir.exists():
        raise ValueError(f"Version {vid} not found for deletion.")

    with FileLock(locks_dir, "delete_version"):
        # Get parents of the version to be deleted
        manifest_path = version_dir / 'manifest.json'
        deleted_parents = []
        if manifest_path.exists():
            manifest = json.loads(manifest_path.read_text(encoding='utf-8'))
            deleted_parents = manifest.get('parents', [])

        # Find children and update their parent lists
        for child_dir in versions_dir.iterdir():
            if child_dir.is_dir() and child_dir.name != vid:
                child_manifest_path = child_dir / 'manifest.json'
                if child_manifest_path.exists():
                    child_manifest = json.loads(child_manifest_path.read_text(encoding='utf-8'))
                    child_parents = child_manifest.get('parents', [])
                    
                    if vid in child_parents:
                        # Remove deleted version from child's parents
                        child_parents.remove(vid)
                        # Add deleted version's parents to child's parents (avoid duplicates)
                        for p in deleted_parents:
                            if p not in child_parents:
                                child_parents.append(p)
                        
                        child_manifest['parents'] = child_parents
                        child_manifest_path.write_text(json.dumps(child_manifest, indent=2), encoding='utf-8')

        # Remove the version directory
        shutil.rmtree(version_dir, ignore_errors=True)

        # Remove any tags associated with this version
        if refs_file.exists():
            with FileLock(locks_dir, "refs_update"):
                refs = {}
                try:
                    text = refs_file.read_text(encoding='utf-8').strip()
                    if text:
                        refs = json.loads(text)
                except json.JSONDecodeError:
                    pass
                tags_to_delete = [tag for tag, v in refs.items() if v == vid]
                if tags_to_delete:
                    for tag in tags_to_delete:
                        del refs[tag]
                    refs_file.write_text(json.dumps(refs, indent=2), encoding='utf-8')
    
    return vid

def get_lineage(repo_path: Path) -> str:
    repo_path = Path(repo_path)
    marco_dir = repo_path / '.marco'
    root_path = marco_dir / 'root.json'
    registry_path = marco_dir / 'registry.json'
    
    if not root_path.exists() or not registry_path.exists():
        return "No lineage found. Are you in a completely initialized marco repository?"
        
    root_data = json.loads(root_path.read_text(encoding='utf-8'))
    registry = json.loads(registry_path.read_text(encoding='utf-8'))
    
    out = []
    d = root_data.get('initialized_at', '')[:10]
    out.append(f"PROJECT: {root_data.get('project_name', 'unknown')}")
    out.append(f"ROOT:    {root_data.get('root_id', 'unknown')}  [initialized {d}]")
    out.append("─────────────────────────────────────────────────────")
    out.append("│")
    
    chains = set(v.get('chain_id') for v in registry.get('versions', []) if v.get('chain_id'))
    sorted_chains = sorted(list(chains))
    
    for i, chain_id in enumerate(sorted_chains):
        chain_versions = [v for v in registry.get('versions', []) if v.get('chain_id') == chain_id]
        if not chain_versions: continue
        
        raw_hash = chain_versions[0].get('raw_data_hash', 'unknown')[:6]
        prefix = "├──" if i < len(sorted_chains) - 1 else "└──"
        out.append(f"{prefix} {chain_id.upper().replace('_', ' ')}  ──  raw: {raw_hash}")
        
        bar = "│" if i < len(sorted_chains) - 1 else " "
        
        def sort_key(v):
            lbl = v.get('label', '')
            if lbl.startswith('V') and lbl[1:].isdigit():
                return int(lbl[1:])
            return 9999
        chain_versions.sort(key=sort_key)
        
        for j, v in enumerate(chain_versions):
            out.append(f"{bar}    │")
            d_v = v.get('created_at', '')[:10]
            label = v.get('label', 'V?')
            v_id = v.get('version_id', 'unknown')[:6]
            out.append(f"{bar}    ●  {label:<3} {v_id}  [{d_v}]")
            
        out.append(f"{bar}")
        
    return "\n".join(out)
