"""Flat file I/O — write fixed-length records from COBOL FD definitions"""
import re, struct
from pathlib import Path
from typing import Any

def analyze_fd_layout(source_text: str, copybook_dirs: list[str] = None) -> dict[str, dict]:
    """From COBOL source, extract FD file layouts."""
    from .read import preprocess, parse_file_control, parse_file_section, parse_data_division, extract_data_division, scan_open_statements

    pp = preprocess(source_text, extra_search_paths=copybook_dirs)
    fc = parse_file_control(pp) if pp else {}
    fs = parse_file_section(pp) if pp else {}
    ops = scan_open_statements(pp) if pp else {}
    dd = extract_data_division(pp)
    all_fields = parse_data_division(dd) if dd else []

    layouts = {}
    for fd_name, rec_names in fs.items():
        records = []
        for rec_name in rec_names:
            children = []
            found = False
            rec_level = None
            offset = 0
            for f in all_fields:
                if f.name == rec_name:
                    found = True
                    rec_level = f.level
                    rec_field_obj = f
                    continue
                if found:
                    if f.level is not None and f.level <= rec_level:
                        break
                    if f.is_88 or f.is_filler:
                        continue
                    pi = f.pic_info
                    if pi:
                        length = (pi.digits + pi.decimal) if pi.type == "numeric" else (pi.length or 0)
                    else:
                        length = 0
                    ftype = pi.type if pi else "unknown"
                    usage = f.usage if f.usage else None
                    children.append({
                        "name": f.name, "pic": str(f.pic or ""),
                        "type": ftype, "length": length, "offset": offset,
                        "usage": usage,
                        "pic_info": {
                            "type": f.pic_info.type if f.pic_info else "unknown",
                            "digits": f.pic_info.digits if f.pic_info else 0,
                            "decimal": f.pic_info.decimal if f.pic_info else 0,
                            "length": f.pic_info.length if f.pic_info else 0,
                            "signed": f.pic_info.signed if f.pic_info else False,
                        } if f.pic_info else None,
                    })
                    offset += length
            # If record has no elementary children but the 01-level has PIC info
            # (e.g. 01 SYSINREC PIC X(080)), use it as a single opaque field
            if not children and rec_field_obj and rec_field_obj.pic_info:
                pi = rec_field_obj.pic_info
                length = pi.length or 0
                if length > 0:
                    children.append({
                        "name": rec_field_obj.name,
                        "pic": str(rec_field_obj.pic or ""),
                        "type": "alphanumeric",
                        "length": length,
                        "offset": 0,
                        "usage": None,
                        "pic_info": {
                            "type": "alphanumeric",
                            "digits": 0,
                            "decimal": 0,
                            "length": length,
                            "signed": False,
                        },
                    })
                    offset = length
            records.append({"record_name": rec_name, "fields": children, "record_length": offset})

        assign_to = fc.get(fd_name, {}).get("assign", fd_name)
        layouts[assign_to] = {
            "fd_name": fd_name, "records": records,
            "direction": ops.get(fd_name, "INPUT"),
        }
    return layouts


def select_records_for_file(records: list[dict], layout: dict) -> list[dict]:
    """Extract and route only the fields belonging to this file layout."""
    if not layout or not layout.get("records"):
        return records
    field_names = set()
    for rec in layout["records"]:
        for f in rec["fields"]:
            field_names.add(f["name"])
    result = []
    for rec in records:
        row = {k: v for k, v in rec.items() if k in field_names}
        if row:
            result.append(row)
    return result if result else records


def _format_value(value: Any, field: dict) -> bytes:
    """Format a value for COBOL fixed-length storage."""
    from . import file_io

    usage = field.get("usage")
    if usage in ("COMP", "COMP-3", "BINARY", "PACKED-DECIMAL"):
        pic_info = field.get("pic_info") or {}
        packed = file_io.pack_value(str(value) if value is not None else "", {
            "usage": usage,
            "pic_info": pic_info,
        })
        want_len = file_io.get_storage_length({
            "usage": usage,
            "pic_info": pic_info,
        })
        if len(packed) < want_len:
            packed = packed.rjust(want_len, b'\x00')
        return packed[:want_len]

    ftype = field["type"]
    length = field["length"]
    val = str(value) if value is not None else ""

    if ftype == "numeric":
        sval = str(val) if val is not None else ""
        # Preserve spaces for DP#7 (R01EMP-ID = SPACE) and similar COBOL checks
        if sval and all(c == ' ' for c in sval):
            return (' ' * length).encode("ascii")
        try:
            num = int(float(val)) if val else 0
        except (ValueError, TypeError):
            num = 0
        num = abs(num)
        max_val = 10 ** length - 1
        if num > max_val:
            num = max_val
        s = str(num).zfill(length)
        if len(s) > length:
            s = s[-length:]
        return s.encode("ascii")
    else:
        s = val.ljust(length)[:length]
        return s.encode("ascii", errors="replace")


def write_flat_file(records: list[dict], layout: dict, outpath: Path, field_filter: set = None):
    """Write records as a COBOL-compatible fixed-length flat file.

    Supports multi-record FDs: uses the longest record layout (most fields)
    to maximize compatible field coverage.
    """
    outpath = Path(outpath)
    if not layout or not layout.get("records"):
        return
    # Pick the record with the most fields (best coverage for multi-record FDs)
    rec = max(layout["records"], key=lambda r: (len(r["fields"]), r["record_length"]))
    if rec["record_length"] == 0:
        return

    rec_fields = rec["fields"]
    if field_filter:
        rec_fields = [f for f in rec_fields if f["name"] in field_filter]

    with open(outpath, "wb") as f:
        for row in records:
            buf = bytearray()
            for field in rec_fields:
                val = row.get(field["name"], "")
                formatted = _format_value(val, field)
                buf.extend(formatted)
            f.write(buf)


def write_all_files(records: list[dict], source_text: str, outdir: Path, prefix: str = "", copybook_dirs: list[str] = None):
    """Analyze source, write flat files for all INPUT FDs."""
    outdir = Path(outdir)
    layouts = analyze_fd_layout(source_text, copybook_dirs=copybook_dirs)
    written = []
    for filename, layout in layouts.items():
        if layout["direction"] == "OUTPUT":
            continue
        # Skip SYSIN files — handled separately by write_sysin_file
        if layout["fd_name"] == "SYSINFILE":
            continue
        fnames = set()
        for rec in layout["records"]:
            for f in rec["fields"]:
                fnames.add(f["name"])
        if not fnames:
            continue
        # Filter generated records to only include fields from this FD
        filtered = [{k: v for k, v in r.items() if k in fnames} for r in records]
        has_data = any(v for row in filtered for v in row.values())
        if not has_data:
            # Fallback: one zero-filled record from FD layout
            fallback = {}
            for rec in layout["records"]:
                for f in rec["fields"]:
                    fallback[f["name"]] = 0 if f["type"] == "numeric" else " "
            filtered = [fallback] if fallback else []
        if filtered:
            outpath = outdir / (prefix + filename)
            write_flat_file(filtered, layout, outpath)
            written.append((filename, outpath, len(filtered)))
    return written


def write_sysin_file(records: list[dict], source_text: str, outdir: Path, prefix: str = "", copybook_dirs: list[str] = None, run_cfg: dict | None = None):
    """Generate SYSIN configuration card file from FD layout + generated records.
    
    Args:
        run_cfg: Scenario sysin override dict, e.g. {"period": "202607", "include_invalid_period": True, "modes": ["NORMAL", "RESET"]}.
                 If None or empty, uses defaults (existing behavior).
    """
    outdir = Path(outdir)
    layouts = analyze_fd_layout(source_text, copybook_dirs=copybook_dirs)
    sysin_filename = None
    sysin_layout = None
    for filename, layout in layouts.items():
        if layout["fd_name"] == "SYSINFILE":
            sysin_filename = filename
            sysin_layout = layout
            break
    if not sysin_layout:
        return None

    # Determine record length from layout
    rec_length = 0
    for rec in sysin_layout["records"]:
        if rec["record_length"] > rec_length:
            rec_length = rec["record_length"]
    if rec_length == 0:
        rec_length = 80

    # Extract employee IDs from the generated R01 flat file if it exists,
    # falling back to JSON record fields for backward compatibility.
    r01_path = outdir / "KIN08R01"
    emp_ids_ordered = []
    emp_set = set()
    if r01_path.exists():
        rec_size = 200
        data = r01_path.read_bytes()
        num_recs = len(data) // rec_size
        for i in range(num_recs):
            off = i * rec_size
            eid = data[off:off+8].decode("ascii", errors="replace").strip()
            if eid and eid != "00000000":
                emp_ids_ordered.append(eid)
    else:
        for r in records:
            eid = r.get("R01EMP-ID", "") or r.get("R01INNREC", {}).get("R01EMP-ID", "")
            if not eid or eid == "00000000":
                line = r.get("R01LINE", "")
                if line:
                    eid = line.split(",")[0].strip()
            if eid and eid != "00000000":
                emp_ids_ordered.append(eid)
                if eid not in emp_set:
                    emp_set.add(eid)
                else:
                    emp_ids_ordered.append(eid)

    # Build SYSIN card records
    # Card format: position 1 = type, position 2 = space (ignored), position 3+ = data
    if run_cfg is None:
        run_cfg = {}
    lines = ["* GENERATED TEST DATA"]
    sysin = run_cfg
    period = sysin.get("period", "202607")
    if period is not None:
        lines.append(f"P YEAR-MONTH={period}")
    if sysin.get("include_invalid_period", True):
        lines.append("P YEAR-MONTH=000000")
    for mode in (sysin.get("modes") or ["NORMAL"]):
        lines.append(f"M MODE={mode}")

    # Detect duplicate EMP_IDs → need RESET mode
    seen = set()
    has_dups = False
    for eid in emp_ids_ordered:
        if eid in seen:
            has_dups = True
            break
        seen.add(eid)

    if has_dups:
        lines.append("M MODE=RESET")
        chunks = [emp_ids_ordered[i:i+8] for i in range(0, len(emp_ids_ordered), 8)]
        for chunk in chunks:
            lines.append(f"T {','.join(chunk)}")
        lines.append("M MODE=NORMAL")

    # Always add a RESET-mode batch with duplicate EMP-ID for UPDATE path
    unique_ids = list(dict.fromkeys(emp_ids_ordered))
    if not unique_ids:
        unique_ids = ["EMP00001", "EMP00002", "EMP00003", "EMP00004", "EMP00005"]
    # Inject duplicates: use first EMP-ID twice to trigger RESET/UPDATE
    if len(unique_ids) >= 1:
        dup_id = unique_ids[0]
        reset_ids = [dup_id, dup_id]
        lines.append("M MODE=RESET")
        lines.append(f"T {','.join(reset_ids)}")
        lines.append("M MODE=NORMAL")
    # Unique T card for NORMAL mode
    if unique_ids:
        chunks = [unique_ids[i:i+8] for i in range(0, len(unique_ids), 8)]
        for chunk in chunks:
            lines.append(f"T {','.join(chunk)}")

    # Unknown card type to cover IF WRK-CARD-TYPE '*' ELSE branch (DP#8 F)
    lines.append("X UNKNOWN")

    # End with RESET mode so 3000STPSOR runs the RESET path (DP#22-#23)
    lines.append("M MODE=RESET")

    # Write as fixed-length flat file
    outpath = outdir / (prefix + sysin_filename)
    with open(outpath, "wb") as f:
        for line in lines:
            buf = line.encode("ascii", errors="replace")
            if len(buf) < rec_length:
                buf = buf.ljust(rec_length, b" ")
            f.write(buf[:rec_length])
    return outpath