\
from __future__ import annotations
import os, json, hashlib, datetime, re
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Tuple

def _read_json(path: str) -> dict:
    with open(path, "r", encoding="latin-1") as f:
        return json.load(f)

def _sha1_of_file(path: str, max_bytes: int = 4 * 1024 * 1024) -> str:
    """Fast-ish hash for large files: hash first max_bytes + file size."""
    h = hashlib.sha1()
    size = os.path.getsize(path)
    h.update(str(size).encode("utf-8"))
    with open(path, "rb") as f:
        h.update(f.read(max_bytes))
    return h.hexdigest()

def _infer_engine(config: dict, case_dir: str) -> str:
    # Prefer explicit version field
    for k in ["version", "engine", "tough_version"]:
        if k in config and isinstance(config[k], str) and config[k].strip():
            v = config[k].lower()
            if "react" in v: return "toughreact"
            if "tough3" in v: return "tough3"
            if "tough2" in v: return "tough2"
            return config[k]
    # Infer from EOS
    eos = next((v for k, v in config.items() if k.lower() == "eos"), None)
    if isinstance(eos, str):
        e = eos.lower()
        if "eco2" in e or "ewas" in e:  # usually TOUGH2/REACT, but keep generic
            return "tough-family"
    # If files contain "treact" exe etc.
    for fn in os.listdir(case_dir):
        low = fn.lower()
        if "treact" in low or "toughreact" in low:
            return "toughreact"
    return "unknown"

def _parse_mesh_counts(mesh_path: str) -> Tuple[Optional[int], Optional[int]]:
    """Count ELEME and CONNE records in a TOUGH MESH file.

    This is a heuristic parser: it scans blocks that start with 'ELEME' and 'CONNE'.
    Returns (n_elements, n_connections) or (None, None) if not detected.
    """
    if not os.path.isfile(mesh_path):
        return None, None
    n_el = n_co = 0
    state = None
    try:
        with open(mesh_path, "r", encoding="latin-1", errors="ignore") as f:
            for line in f:
                tag = line[:5].upper()
                if tag == "ELEME":
                    state = "ELEME"
                    continue
                if tag == "CONNE":
                    state = "CONNE"
                    continue
                if tag.strip() in {"", "END"}:
                    continue
                if state == "ELEME":
                    # element records often have 5-char name at col 1-5
                    if len(line.strip()) >= 5:
                        n_el += 1
                elif state == "CONNE":
                    if len(line.strip()) >= 10:
                        n_co += 1
        if n_el == 0 and n_co == 0:
            return None, None
        return n_el, n_co
    except Exception:
        return None, None

_VAR_RE = re.compile(r'VARIABLES\s*=\s*(.*)', re.IGNORECASE)

def _parse_tecplot_variables(path: str) -> List[str]:
    """Try to extract variable names from a Tecplot .tec file."""
    if not os.path.isfile(path):
        return []
    try:
        with open(path, "r", encoding="latin-1", errors="ignore") as f:
            for _ in range(60):  # usually appears early
                line = f.readline()
                if not line:
                    break
                m = _VAR_RE.search(line)
                if m:
                    rest = m.group(1)
                    # collect quoted tokens
                    vars_ = re.findall(r'"([^"]+)"', rest)
                    # sometimes variables continue next lines; read a few more
                    for _ in range(5):
                        pos = f.tell()
                        nxt = f.readline()
                        if not nxt:
                            break
                        if "ZONE" in nxt.upper():
                            break
                        vars_.extend(re.findall(r'"([^"]+)"', nxt))
                    return [v.strip() for v in vars_ if v.strip()]
        return []
    except Exception:
        return []

def _io_signature(case_dir: str, config: dict) -> str:
    flags = []
    # file presence flags
    def has(name: str) -> bool:
        return os.path.exists(os.path.join(case_dir, name))
    for f in ["MESH","INCON","SAVE","MINC","ROCKS","GENER","ELEME","CONNE"]:
        if has(f): flags.append(f"has_{f.lower()}")
    # configured inputs/outputs
    in_files = [os.path.basename(p) for p in config.get("input_files", []) if isinstance(p, str)]
    out_files = [os.path.basename(p) for p in config.get("output_files", []) if isinstance(p, str)]
    if any(fn.lower().endswith(".tec") for fn in out_files): flags.append("has_tecplot")
    if any(fn.lower().endswith(".csv") for fn in out_files): flags.append("has_csv")
    # options
    for opt in ["minc","ngv","trap"]:
        v = next((v for k, v in config.items() if k.lower() == opt), False)
        if bool(v): flags.append(opt)
    eos = next((v for k, v in config.items() if k.lower() == "eos"), None)
    if isinstance(eos, str) and eos.strip():
        flags.append(f"eos:{eos.strip()}")
    return "+".join(sorted(flags)) if flags else "unknown"

@dataclass
class CaseMeta:
    case_id: str
    case_dir: str
    case_name: str
    engine: str
    eos: str
    version: str
    io_signature: str
    mesh_type: str
    n_elements: Optional[int]
    n_connections: Optional[int]
    fields: List[str]
    fields_hash: str
    files_hash: str
    last_seen: str
    source: str = "unknown"
    notes: str = ""
    status: str = "unknown"
    tested_on: str = ""
    toughanimator_commit: str = ""

def scan_case(case_dir: str) -> CaseMeta:
    cfg_path = os.path.join(case_dir, "config.json")
    config = _read_json(cfg_path) if os.path.isfile(cfg_path) else {}
    case_name = config.get("case_name") or os.path.basename(case_dir.rstrip(os.sep))
    eos = next((v for k, v in config.items() if k.lower() == "eos"), "") or ""
    version = config.get("version","") if isinstance(config.get("version",""), str) else ""
    engine = _infer_engine(config, case_dir)
    sig = _io_signature(case_dir, config)

    mesh_path = os.path.join(case_dir, "MESH")
    n_el, n_co = _parse_mesh_counts(mesh_path)
    mesh_type = config.get("mesh_type","unknown") if isinstance(config.get("mesh_type","unknown"), str) else "unknown"

    # fields: union from output .tec files
    fields: List[str] = []
    for fn in config.get("output_files", []) if isinstance(config.get("output_files", []), list) else []:
        if isinstance(fn, str) and fn.lower().endswith(".tec"):
            fields.extend(_parse_tecplot_variables(os.path.join(case_dir, fn)))
    # de-dup preserving order
    seen=set(); fields=[x for x in fields if not (x in seen or seen.add(x))]
    fields_hash = hashlib.sha1(("|".join(sorted(fields))).encode("utf-8")).hexdigest() if fields else ""

    # files hash: hash key files (config + mesh + incon + save if exist)
    key_files = ["config.json","MESH","INCON","SAVE"]
    h = hashlib.sha1()
    for kf in key_files:
        p = os.path.join(case_dir, kf)
        if os.path.isfile(p):
            h.update(kf.encode("utf-8"))
            h.update(_sha1_of_file(p).encode("utf-8"))
    files_hash = h.hexdigest()

    # case_id: stable hash based on files_hash + case_name
    case_id = hashlib.sha1((case_name+"|"+files_hash).encode("utf-8")).hexdigest()[:12]

    now = datetime.datetime.now().strftime("%Y-%m-%d")
    return CaseMeta(
        case_id=case_id,
        case_dir=os.path.abspath(case_dir),
        case_name=case_name,
        engine=engine,
        eos=str(eos),
        version=str(version),
        io_signature=sig,
        mesh_type=mesh_type,
        n_elements=n_el,
        n_connections=n_co,
        fields=fields,
        fields_hash=fields_hash,
        files_hash=files_hash,
        last_seen=now,
        source=config.get("source","unknown") if isinstance(config.get("source","unknown"), str) else "unknown",
        notes=config.get("notes","") if isinstance(config.get("notes",""), str) else "",
        status=config.get("status","unknown") if isinstance(config.get("status","unknown"), str) else "unknown",
        tested_on=config.get("tested_on","") if isinstance(config.get("tested_on",""), str) else "",
        toughanimator_commit=config.get("toughanimator_commit","") if isinstance(config.get("toughanimator_commit",""), str) else ""
    )

def scan_cases(cases_root: str) -> List[CaseMeta]:
    """Scan cases_root for subfolders containing config.json."""
    metas: List[CaseMeta] = []
    for root, dirs, files in os.walk(cases_root):
        if "config.json" in files:
            metas.append(scan_case(root))
            # don't recurse into nested case dirs under this case unless you want
            dirs[:] = []
    return metas
