# src/preprocessing/data_loader.py
"""
يحمّل داتاست PIE4Perf من الملفات المحلية.

الملفات المطلوبة في data/raw/:
    train.jsonl   للتدريب
    val.jsonl     للـ validation
    test.jsonl    للـ evaluation

تشغيل:
    py -3.11 -m src.preprocessing.data_loader
"""

import json
from pathlib import Path
from dataclasses import dataclass
from typing import List, Tuple
from src.utils.config import DATA_RAW_DIR
from src.utils.logger import get_logger

log = get_logger("data_loader")


@dataclass
class Sample:
    slow_code: str
    fast_code: str
    speedup: float

    def is_valid(self) -> bool:
        return (
            bool(self.slow_code.strip())
            and bool(self.fast_code.strip())
            and self.speedup > 0
        )


def _parse(record: dict) -> "Sample | None":
    try:
        slow = (record.get("input")
                or record.get("code_v0_no_empty_lines")
                or record.get("slow_code") or "")

        fast = (record.get("target")
                or record.get("code_v1_no_empty_lines")
                or record.get("fast_code") or "")

        cpu0 = record.get("cpu_time_v0")
        cpu1 = record.get("cpu_time_v1")
        if cpu0 and cpu1 and float(cpu1) > 0:
            speedup = float(cpu0) / float(cpu1)
        elif record.get("improvement_frac"):
            speedup = 1 + float(record["improvement_frac"]) / 100
        else:
            speedup = 1.0

        s = Sample(slow_code=slow.strip(), fast_code=fast.strip(), speedup=speedup)
        return s if s.is_valid() else None
    except Exception:
        return None


def load_raw(path: Path) -> List[Sample]:
    """يقرأ ملف JSONL واحد ويرجع list من الـ samples"""
    if not path.exists():
        log.warning(f"الملف مش موجود: {path.name}")
        return []

    samples = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
                s = _parse(record)
                if s:
                    samples.append(s)
            except json.JSONDecodeError:
                continue

    log.info(f"تحميل {path.name}  ←  {len(samples)} sample")
    return samples


def load_dataset() -> Tuple[List[Sample], List[Sample], List[Sample]]:
    """
    يقرأ train.jsonl و val.jsonl و test.jsonl من data/raw/
    ويرجع (train, val, test)
    """
    train = load_raw(DATA_RAW_DIR / "train.jsonl")
    val   = load_raw(DATA_RAW_DIR / "val.jsonl")
    test  = load_raw(DATA_RAW_DIR / "test.jsonl")

    # لو مفيش داتا خالص استخدم dummy
    if not train and not val and not test:
        log.warning("مفيش داتا في data/raw/  ← بستخدم dummy samples")
        dummy = _dummy()
        n = len(dummy)
        return dummy[:int(n*0.8)], dummy[int(n*0.8):int(n*0.9)], dummy[int(n*0.9):]

    log.info(f"Dataset: {len(train)} train / {len(val)} val / {len(test)} test")
    return train, val, test


def _dummy() -> List[Sample]:
    return [
        Sample("result=[]\nfor x in a:\n    result.append(x*2)",
               "result=[x*2 for x in a]", 1.3),
        Sample("s=''\nfor w in words:\n    s+=w+' '",
               "s=' '.join(words)", 5.0),
        Sample("def f(x=[]):\n    x.append(1)\n    return x",
               "def f(x=None):\n    if x is None: x=[]\n    x.append(1)\n    return x", 1.0),
    ]


if __name__ == "__main__":
    train, val, test = load_dataset()
    print(f"\nTrain: {len(train)}  Val: {len(val)}  Test: {len(test)}")
    if train:
        print("\nمثال من Train:")
        print("SLOW:", train[0].slow_code[:80])
        print("FAST:", train[0].fast_code[:80])
        print("Speedup:", round(train[0].speedup, 2))