import hashlib
import json
import os
import shutil
import tempfile
import tarfile
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

def init_repo(repo_path: Union[str, Path]):
    repo_path = Path(repo_path)
    marco_dir = repo_path / '.marco'
    (marco_dir / 'versions').mkdir(parents=True, exist_ok=True)
    (marco_dir / 'locks').mkdir(parents=True, exist_ok=True)
    refs_file = marco_dir / 'refs.json'
    if not refs_file.exists():
        refs_file.write_text('{}', encoding='utf-8')
    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):
    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"):
        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()))

        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,
            '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
        }
        
        tmp = Path(tempfile.mkdtemp(prefix='version_tmp_', dir=versions_dir))
        try:
            shutil.copy2(raw_path, tmp / 'raw.txt')
            (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 / 'metrics.json').write_text(json.dumps(final_metrics, 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')
            
            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():
            refs = json.loads(refs_file.read_text(encoding='utf-8'))
        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():
        return json.loads(refs_file.read_text(encoding='utf-8'))
    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
