import os
import json
import hashlib
from typing import List, Dict, Tuple, Any
import pandas as pd
from sklearn.model_selection import train_test_split
from .schema import DocumentRecord

def compute_file_hash(filepath: str) -> str:
    hasher = hashlib.sha256()
    with open(filepath, "rb") as f:
        while chunk := f.read(8192):
            hasher.update(chunk)
    return hasher.hexdigest()

def build_stratified_splits(
    records: List[DocumentRecord],
    splits_dir: str,
    processed_dir: str,
    seed: int = 42,
    train_ratio: float = 0.70,
    val_ratio: float = 0.15,
    test_ratio: float = 0.15
) -> Dict[str, Any]:
    """
    Creates stratified 70/15/15 train/val/test splits, writes manifests as CSV,
    saves unified records to jsonl, and returns metadata including SHA256 hashes
    and class weight balancing strategy.
    """
    os.makedirs(splits_dir, exist_ok=True)
    os.makedirs(processed_dir, exist_ok=True)

    jsonl_path = os.path.join(processed_dir, "unified_documents.jsonl")
    with open(jsonl_path, "w", encoding="utf-8") as f_out:
        for r in records:
            f_out.write(r.model_dump_json() + "\n")

    # Build DataFrame for stratification
    df_meta = pd.DataFrame([
        {
            "doc_id": r.doc_id,
            "source_dataset": r.source_dataset,
            "doc_class": r.doc_class,
            "image_path": r.image_path,
            "ocr_token_count": len(r.ocr_words)
        }
        for r in records
    ])

    # Stratified split: train vs temp (val + test)
    temp_ratio = val_ratio + test_ratio
    train_df, temp_df = train_test_split(
        df_meta,
        test_size=temp_ratio,
        random_state=seed,
        stratify=df_meta["doc_class"]
    )

    # Split temp into val and test
    val_relative_ratio = val_ratio / temp_ratio
    val_df, test_df = train_test_split(
        temp_df,
        test_size=(1.0 - val_relative_ratio),
        random_state=seed,
        stratify=temp_df["doc_class"]
    )

    train_path = os.path.join(splits_dir, "train.csv")
    val_path = os.path.join(splits_dir, "val.csv")
    test_path = os.path.join(splits_dir, "test.csv")

    train_df.to_csv(train_path, index=False)
    val_df.to_csv(val_path, index=False)
    test_df.to_csv(test_path, index=False)

    # Class distribution and balanced weights
    class_counts = train_df["doc_class"].value_counts().to_dict()
    total_samples = len(train_df)
    n_classes = len(class_counts)
    class_weights = {
        cls: total_samples / (n_classes * count)
        for cls, count in class_counts.items()
    }

    manifest_hashes = {
        "train_manifest_hash": compute_file_hash(train_path),
        "val_manifest_hash": compute_file_hash(val_path),
        "test_manifest_hash": compute_file_hash(test_path),
        "jsonl_hash": compute_file_hash(jsonl_path)
    }

    return {
        "train_count": len(train_df),
        "val_count": len(val_df),
        "test_count": len(test_df),
        "class_counts_train": class_counts,
        "class_weights": class_weights,
        "hashes": manifest_hashes
    }