scripts/triplets_synthesis.py

'''
Synthesize triplet and positive pair datasets from chunked code files.'''

import argparse
import json
import random
import hashlib
from pathlib import Path
from typing import Dict, List
from datetime import datetime
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


# ============================
# CONFIG
# ============================

MAX_DOCUMENTS = 200
POSITIVE_VARIANTS = 5
TFIDF_MAX_FEATURES = 5000
RANDOM_SEED = 42

BASE_OUTPUT_DIR = Path("data/synthetic")

random.seed(RANDOM_SEED)


# ============================
# UTILITIES
# ============================

def load_chunks(file_path):
    path = Path(file_path)

    if path.suffix == ".jsonl":
        chunks = []
        with open(path, "r", encoding="utf-8") as f:
            for line_no, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    chunks.append(json.loads(line))
                except json.JSONDecodeError as e:
                    raise ValueError(
                        f"Invalid JSON on line {line_no} in {path}"
                    ) from e
        return chunks

    elif path.suffix == ".json":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, list):
            raise ValueError(f"{path} must contain a list of chunks")
        return data

    else:
        raise ValueError(
            f"Unsupported file format {path.suffix}. Use .json or .jsonl"
        )



def save_jsonl(path: Path, records: List[Dict]):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def save_json(path: Path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def stable_document_id(chunk: Dict, idx: int) -> str:
    """
    Generate a canonical, stable document_id.
    """
    base = f"{chunk.get('file_path','unknown')}::{idx}"
    return "doc_" + hashlib.sha1(base.encode()).hexdigest()


def infer_framework(input_path: Path) -> str:
    """
    Infer framework from path (fallback-safe).
    """
    parts = [p.lower() for p in input_path.parts]
    for fw in ["crewai", "langchain", "langgraph", "autogen"]:
        if fw in parts:
            return fw
    return "unknown"


# ============================
# ANCHOR GENERATION (LLM PLACEHOLDER)
# ============================

def generate_anchor_questions(code: str, n: int) -> List[str]:
    """
    Deterministic placeholder (LLM-ready).
    """
    symbol = code.split("(")[0].replace("def ", "").replace("class ", "").strip()

    templates = [
        f"How does {symbol} work in Python?",
        f"How to implement {symbol}?",
        f"Example usage of {symbol}",
        f"Explain the {symbol} logic",
        f"Best practices for {symbol}",
    ]

    random.shuffle(templates)
    return templates[:n]


# ============================
# NEGATIVE MINING
# ============================

def build_tfidf(chunks: List[Dict]):
    corpus = [c["code"] for c in chunks]
    vectorizer = TfidfVectorizer(
        stop_words="english",
        max_features=TFIDF_MAX_FEATURES
    )
    matrix = vectorizer.fit_transform(corpus)
    return vectorizer, matrix


def mine_hard_negative(
    anchor: str,
    positive_idx: int,
    chunks: List[Dict],
    vectorizer,
    matrix,
) -> Dict:
    query_vec = vectorizer.transform([anchor])
    scores = cosine_similarity(query_vec, matrix)[0]

    ranked = sorted(
        [(i, s) for i, s in enumerate(scores)],
        key=lambda x: x[1],
        reverse=True,
    )

    for idx, _ in ranked:
        if idx != positive_idx:
            return chunks[idx]

    raise RuntimeError("No negative candidate found")


# ============================
# MAIN PIPELINE
# ============================

def generate_datasets(input_path: Path, run_name: str):
    output_dir = BASE_OUTPUT_DIR / run_name
    framework = infer_framework(input_path)

    chunks = load_chunks(input_path)
    # Filter only semantic code chunks
    chunks = [
        c for c in chunks
        if c.get("chunk_type") in {"class", "method", "function"}
        and "code" in c
    ]

    random.shuffle(chunks)
    chunks = chunks[:MAX_DOCUMENTS]

    # Assign canonical document_id
    for idx, c in enumerate(chunks):
        c["document_id"] = stable_document_id(c, idx)

    vectorizer, matrix = build_tfidf(chunks)

    positive_pairs = []
    triplets = []

    for idx, chunk in enumerate(chunks):
        code = chunk["code"]
        doc_id = chunk["document_id"]

        # -------- POSITIVE PAIRS --------
        anchors = generate_anchor_questions(code, POSITIVE_VARIANTS)
        for a in anchors:
            positive_pairs.append({
                "document_id": doc_id,
                "anchor": a,
                "positive": code,
                "framework": framework,
                "source": "synthetic_positive_v2",
            })

        # -------- TRIPLET --------
        anchor = anchors[0]
        negative_chunk = mine_hard_negative(
            anchor, idx, chunks, vectorizer, matrix
        )

        triplets.append({
            "document_id": doc_id,
            "anchor": anchor,
            "positive": code,
            "negative": negative_chunk["code"],
            "framework": framework,
            "source": "synthetic_triplet_v2",
        })

    # -------- SAVE --------
    save_jsonl(output_dir / "positive_pairs.jsonl", positive_pairs)
    save_jsonl(output_dir / "triplets.jsonl", triplets)

    save_json(output_dir / "positive_pairs.json", positive_pairs)
    save_json(output_dir / "triplets.json", triplets)

    metadata = {
        "name": run_name,
        "framework": framework,
        "input_file": str(input_path),
        "num_chunks": len(chunks),
        "positive_pairs": len(positive_pairs),
        "triplets": len(triplets),
        "created_at": datetime.utcnow().isoformat(),
        "random_seed": RANDOM_SEED,
    }

    save_json(output_dir / "metadata.json", metadata)

    print(f"✅ Dataset generated at: {output_dir}")


# ============================
# ENTRY POINT
# ============================

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="Chunked JSONL file")
    parser.add_argument("--name", required=True, help="Synthetic dataset name")

    args = parser.parse_args()

    generate_datasets(
        input_path=Path(args.input),
        run_name=args.name,
    )

# # For document id 

# document_id := sha1(
#     normalized_repo_path +
#     file_path +
#     top_level_symbol
# )