import argparse
import ast
import csv
import json
import os
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Set, Tuple


NULL_LITERALS = {"NULL", "null", ""}


def _is_null(value: Optional[str]) -> bool:
    if value is None:
        return True
    return value.strip() in NULL_LITERALS


def _json_loads_best_effort(raw: str) -> Any:
    s = raw.strip()
    if not s:
        raise ValueError("empty")
    if s[0] in ("'", '"') and s[-1] == s[0]:
        s = s[1:-1].strip()
    try:
        return json.loads(s)
    except Exception:
        pass
    if s.startswith("{") or s.startswith("["):
        try:
            return ast.literal_eval(s)
        except Exception:
            pass
    if '""' in s:
        try:
            return json.loads(s.replace('""', '"'))
        except Exception:
            pass
    raise ValueError("not json")


def _stringify_json_value(v: Any) -> str:
    if isinstance(v, (dict, list)):
        return json.dumps(v, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
    if v is None:
        return ""
    return str(v)


def _flatten_top_level(d: Any, prefix: str) -> Dict[str, str]:
    if not isinstance(d, dict):
        return {prefix: _stringify_json_value(d)}
    out: Dict[str, str] = {}
    for k, v in d.items():
        out[f"{prefix}.{k}"] = _stringify_json_value(v)
    return out


_RE_KV_TOKEN = re.compile(
    r"(?P<key>[A-Za-z_][A-Za-z0-9_()]*)\s*=\s*(?P<value>\[[^\]]*\]|\{[^\}]*\}|[^\s\}]+)"
)


def _extract_service_state_metrics(service_state_string: str) -> Dict[str, str]:
    s = service_state_string
    metrics: Dict[str, str] = {}

    for key in (
        "mVoiceRegState",
        "mDataRegState",
        "mChannelNumber",
        "duplexMode()",
        "mCellBandwidths",
        "mOperatorAlphaLong",
        "mOperatorAlphaShort",
        "mIsEmergencyOnly",
        "isUsingCarrierAggregation",
        "mArfcnRsrpBoost",
    ):
        m = re.search(rf"{re.escape(key)}=([^,\}}]+)", s)
        if m:
            metrics[f"service_state.metric.{key}"] = m.group(1).strip()

    for key in ("getRilVoiceRadioTechnology", "getRilDataRadioTechnology"):
        m = re.search(rf"{re.escape(key)}=([^,\}}]+)", s)
        if m:
            metrics[f"service_state.metric.{key}"] = m.group(1).strip()

    return metrics


def _split_top_level(s: str, sep: str) -> List[str]:
    out: List[str] = []
    buf: List[str] = []
    depth = 0
    for ch in s:
        if ch in "{[(":
            depth += 1
        elif ch in "}])":
            if depth > 0:
                depth -= 1
        if ch == sep and depth == 0:
            part = "".join(buf).strip()
            if part:
                out.append(part)
            buf = []
            continue
        buf.append(ch)
    tail = "".join(buf).strip()
    if tail:
        out.append(tail)
    return out


_RE_SIMPLE_KEY_EQ = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)=")


def _kv_by_top_level_keyeq(s: str) -> Dict[str, str]:
    out: Dict[str, str] = {}
    depth = 0
    i = 0
    starts: List[Tuple[str, int, int]] = []
    while i < len(s):
        if depth == 0:
            m = _RE_SIMPLE_KEY_EQ.match(s, i)
            if m:
                starts.append((m.group(1), m.start(), m.end()))
                i = m.end()
                continue
        ch = s[i]
        if ch in "{[(":
            depth += 1
        elif ch in "}])":
            if depth > 0:
                depth -= 1
        i += 1

    for idx, (k, _ks, vs) in enumerate(starts):
        ve = starts[idx + 1][1] if idx + 1 < len(starts) else len(s)
        v = s[vs:ve].strip().strip(",")
        out[k] = v
    return out


def _extract_braced_content(text: str) -> Optional[str]:
    start = text.find("{")
    if start < 0:
        return None
    depth = 0
    for i in range(start, len(text)):
        ch = text[i]
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return text[start + 1 : i]
    return None


def _extract_network_registration_infos(value: str) -> List[str]:
    s = value.strip()
    if s.startswith("[") and s.endswith("]"):
        s = s[1:-1].strip()
    out: List[str] = []
    i = 0
    while i < len(s):
        start = s.find("NetworkRegistrationInfo{", i)
        if start < 0:
            break
        j = start + len("NetworkRegistrationInfo")
        brace_start = s.find("{", j)
        if brace_start < 0:
            break
        depth = 0
        end = None
        for k in range(brace_start, len(s)):
            ch = s[k]
            if ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    end = k
                    break
        if end is None:
            break
        out.append(s[brace_start + 1 : end].strip())
        i = end + 1
    return out


def _extract_service_state_struct(service_state_string: str) -> Dict[str, str]:
    s = service_state_string.strip()
    if s.startswith("{") and s.endswith("}"):
        s = s[1:-1].strip()

    out: Dict[str, str] = {}
    for part in _split_top_level(s, ","):
        if "=" not in part:
            continue
        k, v = part.split("=", 1)
        k = k.strip()
        v = v.strip()
        if not k:
            continue
        if k != "mNetworkRegistrationInfos":
            out[f"service_state.ss.{k}"] = v

        if k == "mNetworkRegistrationInfos":
            nris = _extract_network_registration_infos(v)
            out["service_state.ss.mNetworkRegistrationInfos.count"] = str(len(nris))
            for idx, nri in enumerate(nris):
                kv = _kv_by_top_level_keyeq(nri)
                for kk, vv in kv.items():
                    if kk in {"cellIdentity", "dataSpecificInfo", "voiceSpecificInfo"}:
                        continue
                    out[f"service_state.nri{idx}.{kk}"] = vv

                cell_id = kv.get("cellIdentity")
                if cell_id and cell_id != "null":
                    tech_m = re.search(r"CellIdentity([A-Za-z0-9_]+):\{", cell_id)
                    if tech_m:
                        out[f"service_state.nri{idx}.cellIdentity.tech"] = tech_m.group(1)
                    body = _extract_braced_content(cell_id)
                    if body:
                        for kk, vv in _parse_kv_tokens(body).items():
                            out[f"service_state.nri{idx}.cellIdentity.{kk}"] = vv

                for info_key, prefix in (
                    ("dataSpecificInfo", "dataSpecific"),
                    ("voiceSpecificInfo", "voiceSpecific"),
                ):
                    info = kv.get(info_key)
                    if not info or info == "null":
                        continue
                    body = _extract_braced_content(info)
                    if not body:
                        continue
                    for kk, vv in _parse_kv_tokens(body).items():
                        out[f"service_state.nri{idx}.{prefix}.{kk}"] = vv

    out.update(_extract_service_state_metrics(service_state_string))
    return out


def _parse_service_state_string(
    service_state_string: str,
) -> Tuple[Dict[str, str], Dict[str, str], List[Dict[str, Dict[str, str]]]]:
    s = service_state_string.strip()
    if s.startswith("{") and s.endswith("}"):
        s = s[1:-1].strip()

    ss_out: Dict[str, str] = {}
    metrics_out = _extract_service_state_metrics(service_state_string)
    nris_out: List[Dict[str, Dict[str, str]]] = []

    for part in _split_top_level(s, ","):
        if "=" not in part:
            continue
        k, v = part.split("=", 1)
        k = k.strip()
        v = v.strip()
        if not k:
            continue

        if k != "mNetworkRegistrationInfos":
            ss_out[f"service_state.ss.{k}"] = v
            continue

        nris = _extract_network_registration_infos(v)
        ss_out["service_state.ss.mNetworkRegistrationInfos.count"] = str(len(nris))

        for nri in nris:
            kv = _kv_by_top_level_keyeq(nri)

            fields: Dict[str, str] = {}
            for kk, vv in kv.items():
                if kk in {"cellIdentity", "dataSpecificInfo", "voiceSpecificInfo"}:
                    continue
                fields[f"service_state.nri.{kk}"] = vv

            cell_identity: Dict[str, str] = {}
            cell_id = kv.get("cellIdentity")
            if cell_id and cell_id != "null":
                tech_m = re.search(r"CellIdentity([A-Za-z0-9_]+):\\{", cell_id)
                if tech_m:
                    cell_identity["service_state.nri.cellIdentity.tech"] = tech_m.group(1)
                body = _extract_braced_content(cell_id)
                if body:
                    for kk, vv in _parse_kv_tokens(body).items():
                        cell_identity[f"service_state.nri.cellIdentity.{kk}"] = vv

            data_specific: Dict[str, str] = {}
            info = kv.get("dataSpecificInfo")
            if info and info != "null":
                body = _extract_braced_content(info)
                if body:
                    for kk, vv in _parse_kv_tokens(body).items():
                        data_specific[f"service_state.nri.dataSpecific.{kk}"] = vv

            voice_specific: Dict[str, str] = {}
            info = kv.get("voiceSpecificInfo")
            if info and info != "null":
                body = _extract_braced_content(info)
                if body:
                    for kk, vv in _parse_kv_tokens(body).items():
                        voice_specific[f"service_state.nri.voiceSpecific.{kk}"] = vv

            nris_out.append(
                {
                    "fields": fields,
                    "cell_identity": cell_identity,
                    "data_specific": data_specific,
                    "voice_specific": voice_specific,
                }
            )

    return ss_out, metrics_out, nris_out


@dataclass(frozen=True)
class _CellInfoExtract:
    tech: str
    identity: Dict[str, str]
    signal: Dict[str, str]
    registered: Optional[str]


_RE_CELLINFO_TECH = re.compile(r"CellInfo(?P<tech>[A-Za-z0-9_]+):\{(?P<body>.*)\}\s*\Z", re.DOTALL)
_RE_IDENTITY_BLOCK = re.compile(r"CellIdentity(?P<tech>[A-Za-z0-9_]+):\{\s*(?P<body>.*?)\s*\}", re.DOTALL)
_RE_SIGNAL_LTE = re.compile(
    r"CellSignalStrengthLte:\s*(?P<body>.*?)(?:\s+android\.telephony\.CellConfigLte|\s*\}\s*\Z)",
    re.DOTALL,
)
_RE_SIGNAL_NR = re.compile(r"CellSignalStrengthNr:\{?\s*(?P<body>.*?)\s*\}", re.DOTALL)
_RE_REGISTERED = re.compile(r"mRegistered=(?P<value>[A-Z_]+)")


def _parse_kv_tokens(body: str) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for m in _RE_KV_TOKEN.finditer(body):
        k = m.group("key").strip()
        v = m.group("value").strip().rstrip(",")
        out[k] = v
    return out


def _extract_first_cellinfo(cell_info_text: str) -> Optional[_CellInfoExtract]:
    s = cell_info_text.strip()
    if not s:
        return None
    if s.startswith("[") and s.endswith("]"):
        s = s[1:-1].strip()
    if not s:
        return None

    m = _RE_CELLINFO_TECH.match(s)
    if not m:
        return None

    tech = m.group("tech")
    body = m.group("body")

    reg_m = _RE_REGISTERED.search(body)
    registered = reg_m.group("value") if reg_m else None

    identity: Dict[str, str] = {}
    identity_m = _RE_IDENTITY_BLOCK.search(body)
    if identity_m:
        identity = _parse_kv_tokens(identity_m.group("body"))

    signal: Dict[str, str] = {}
    if tech.lower() == "lte":
        sig_m = _RE_SIGNAL_LTE.search(body)
        if sig_m:
            signal = _parse_kv_tokens(sig_m.group("body"))
    elif tech.lower() == "nr":
        sig_m = _RE_SIGNAL_NR.search(body)
        if sig_m:
            signal = _parse_kv_tokens(sig_m.group("body"))

    return _CellInfoExtract(tech=tech, identity=identity, signal=signal, registered=registered)


def _extract_cell_info_metrics(cell_info_json: Any) -> Dict[str, str]:
    if not isinstance(cell_info_json, dict):
        return {}

    cell_info_raw = cell_info_json.get("cellInfo")
    if not isinstance(cell_info_raw, str):
        return {}

    parsed = _extract_first_cellinfo(cell_info_raw)
    if not parsed:
        return {}

    out: Dict[str, str] = {"cell_info_changed.metric.tech": parsed.tech}
    if parsed.registered is not None:
        out["cell_info_changed.metric.mRegistered"] = parsed.registered

    for k, v in parsed.identity.items():
        out[f"cell_info_changed.identity.{k}"] = v
    for k, v in parsed.signal.items():
        out[f"cell_info_changed.signal.{k}"] = v

    return out


def _collect_dynamic_columns(csv_path: str, json_cols: List[str]) -> Tuple[List[str], Dict[str, Set[str]]]:
    base_headers: List[str] = []
    dynamic_cols: Dict[str, Set[str]] = {"__metrics__": set()}

    with open(csv_path, "r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        if reader.fieldnames:
            base_headers = list(reader.fieldnames)
        for row in reader:
            for col in json_cols:
                raw = row.get(col)
                if _is_null(raw):
                    continue
                try:
                    parsed = _json_loads_best_effort(raw or "")
                except Exception:
                    continue

                flat = _flatten_top_level(parsed, col)
                for k in flat.keys():
                    dynamic_cols.setdefault(col, set()).add(k)

                if col == "service_state" and isinstance(parsed, dict):
                    ss = parsed.get("serviceStateString")
                    if isinstance(ss, str) and ss.strip():
                        for mk in _extract_service_state_struct(ss).keys():
                            dynamic_cols["__metrics__"].add(mk)

                if col == "cell_info_changed":
                    for mk in _extract_cell_info_metrics(parsed).keys():
                        dynamic_cols["__metrics__"].add(mk)

    return base_headers, dynamic_cols


def _build_output_headers(
    base_headers: List[str],
    dynamic_cols: Dict[str, Set[str]],
    json_cols: List[str],
    keep_original_json_cols: bool,
) -> List[str]:
    out: List[str] = []
    for h in base_headers:
        if (not keep_original_json_cols) and (h in json_cols):
            continue
        out.append(h)

    expanded: List[str] = []
    for col in json_cols:
        keys = sorted(dynamic_cols.get(col, set()))
        expanded.extend(keys)
    expanded.extend(sorted(dynamic_cols.get("__metrics__", set())))

    existing = set(out)
    for h in expanded:
        if h not in existing:
            out.append(h)
            existing.add(h)

    return out


def _collect_headers_for_one_output(
    input_csv: str,
    col: str,
    base_cols: List[str],
    drop_raw: bool,
) -> List[str]:
    base_headers, dynamic_cols = _collect_dynamic_columns(input_csv, [col])

    actual_base_cols = [c for c in base_cols if c in base_headers]
    out_headers: List[str] = []
    out_headers.extend(actual_base_cols)
    dynamic = set(dynamic_cols.get(col, set()))
    if drop_raw:
        if col == "service_state":
            dynamic.discard("service_state.serviceStateString")
        if col == "cell_info_changed":
            dynamic.discard("cell_info_changed.cellInfo")
    out_headers.extend(sorted(dynamic))
    out_headers.extend(sorted(dynamic_cols.get("__metrics__", set())))

    seen: Set[str] = set()
    deduped: List[str] = []
    for h in out_headers:
        if h not in seen:
            deduped.append(h)
            seen.add(h)
    return deduped


def _collect_service_state_long_headers(input_csv: str, base_cols: List[str], drop_raw: bool) -> List[str]:
    base_headers: List[str] = []
    top_keys: Set[str] = set()
    ss_keys: Set[str] = set()
    metric_keys: Set[str] = set()
    nri_keys: Set[str] = set()
    cell_id_keys: Set[str] = set()
    data_spec_keys: Set[str] = set()
    voice_spec_keys: Set[str] = set()

    with open(input_csv, "r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        if reader.fieldnames:
            base_headers = list(reader.fieldnames)
        for row in reader:
            raw_service = row.get("service_state")
            if _is_null(raw_service):
                continue
            try:
                parsed = _json_loads_best_effort(raw_service or "")
            except Exception:
                continue
            if not isinstance(parsed, dict):
                continue

            for k in _flatten_top_level(parsed, "service_state").keys():
                top_keys.add(k)

            ss = parsed.get("serviceStateString")
            if not isinstance(ss, str) or not ss.strip():
                continue

            ss_kv, metrics_kv, nris_kv = _parse_service_state_string(ss)
            ss_keys.update(ss_kv.keys())
            metric_keys.update(metrics_kv.keys())
            for nri in nris_kv:
                nri_keys.update(nri.get("fields", {}).keys())
                cell_id_keys.update(nri.get("cell_identity", {}).keys())
                data_spec_keys.update(nri.get("data_specific", {}).keys())
                voice_spec_keys.update(nri.get("voice_specific", {}).keys())

    actual_base_cols = [c for c in base_cols if c in base_headers]
    headers: List[str] = []
    headers.extend(actual_base_cols)

    top_sorted = sorted(top_keys)
    if drop_raw:
        top_sorted = [k for k in top_sorted if k != "service_state.serviceStateString"]
    headers.extend(top_sorted)

    headers.extend(sorted(ss_keys))
    headers.extend(sorted(metric_keys))
    headers.append("service_state.nri.index")
    headers.extend(sorted(nri_keys))
    headers.extend(sorted(cell_id_keys))
    headers.extend(sorted(data_spec_keys))
    headers.extend(sorted(voice_spec_keys))

    seen: Set[str] = set()
    deduped: List[str] = []
    for h in headers:
        if h not in seen:
            deduped.append(h)
            seen.add(h)
    return deduped


def parse_split_csv(
    input_csv: str,
    out_prefix: str,
    base_cols: List[str],
    drop_raw: bool,
    service_format: str,
) -> Tuple[str, str]:
    service_out_csv = f"{out_prefix}.service_state.csv"
    cell_out_csv = f"{out_prefix}.cell_info_changed.csv"

    if service_format == "wide":
        service_headers = _collect_headers_for_one_output(
            input_csv=input_csv, col="service_state", base_cols=base_cols, drop_raw=drop_raw
        )
    else:
        service_headers = _collect_service_state_long_headers(
            input_csv=input_csv, base_cols=base_cols, drop_raw=drop_raw
        )
    cell_headers = _collect_headers_for_one_output(
        input_csv=input_csv, col="cell_info_changed", base_cols=base_cols, drop_raw=drop_raw
    )

    with open(input_csv, "r", encoding="utf-8", newline="") as fin, open(
        service_out_csv, "w", encoding="utf-8-sig", newline=""
    ) as f_service, open(cell_out_csv, "w", encoding="utf-8-sig", newline="") as f_cell:
        reader = csv.DictReader(fin)
        service_writer = csv.DictWriter(f_service, fieldnames=service_headers, extrasaction="ignore")
        cell_writer = csv.DictWriter(f_cell, fieldnames=cell_headers, extrasaction="ignore")
        service_writer.writeheader()
        cell_writer.writeheader()

        for row in reader:
            raw_service = row.get("service_state")
            if not _is_null(raw_service):
                try:
                    parsed = _json_loads_best_effort(raw_service or "")
                except Exception:
                    parsed = None

                if service_format == "wide":
                    out_row: Dict[str, str] = {k: (row.get(k) or "") for k in service_headers if k in row}
                    if parsed is not None:
                        out_row.update(_flatten_top_level(parsed, "service_state"))
                        if isinstance(parsed, dict):
                            ss = parsed.get("serviceStateString")
                            if isinstance(ss, str) and ss.strip():
                                out_row.update(_extract_service_state_struct(ss))
                    service_writer.writerow(out_row)
                else:
                    if isinstance(parsed, dict):
                        common: Dict[str, str] = {k: (row.get(k) or "") for k in service_headers if k in row}
                        top = _flatten_top_level(parsed, "service_state")
                        if drop_raw:
                            top.pop("service_state.serviceStateString", None)
                        common.update(top)

                        ss_kv: Dict[str, str] = {}
                        metrics_kv: Dict[str, str] = {}
                        nris_kv: List[Dict[str, Dict[str, str]]] = []
                        ss = parsed.get("serviceStateString")
                        if isinstance(ss, str) and ss.strip():
                            ss_kv, metrics_kv, nris_kv = _parse_service_state_string(ss)
                        common.update(ss_kv)
                        common.update(metrics_kv)

                        if not nris_kv:
                            out_row = dict(common)
                            out_row["service_state.nri.index"] = ""
                            service_writer.writerow(out_row)
                        else:
                            for idx, nri in enumerate(nris_kv):
                                out_row = dict(common)
                                out_row["service_state.nri.index"] = str(idx)
                                out_row.update(nri.get("fields", {}))
                                out_row.update(nri.get("cell_identity", {}))
                                out_row.update(nri.get("data_specific", {}))
                                out_row.update(nri.get("voice_specific", {}))
                                service_writer.writerow(out_row)

            raw_cell = row.get("cell_info_changed")
            if not _is_null(raw_cell):
                out_row = {k: (row.get(k) or "") for k in cell_headers if k in row}
                try:
                    parsed = _json_loads_best_effort(raw_cell or "")
                except Exception:
                    parsed = None
                if parsed is not None:
                    out_row.update(_flatten_top_level(parsed, "cell_info_changed"))
                    out_row.update(_extract_cell_info_metrics(parsed))
                cell_writer.writerow(out_row)

    return service_out_csv, cell_out_csv


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input_csv", help="输入 CSV 路径")
    parser.add_argument(
        "--out-prefix",
        default=None,
        help="输出前缀（默认: <input 去掉扩展名>）",
    )
    parser.add_argument(
        "--base-cols",
        default="client_dt,slot_id",
        help="每行都保留的基础列（逗号分隔），例如 client_dt,slot_id,imei",
    )
    parser.add_argument(
        "--keep-raw",
        action="store_true",
        help="保留原始大字段（service_state.serviceStateString / cell_info_changed.cellInfo）。默认不保留。",
    )
    parser.add_argument(
        "--service-format",
        choices=("long", "wide"),
        default="long",
        help="service_state 输出格式：long 更适合表格查看（默认），wide 会生成 nri0/nri1... 列",
    )
    args = parser.parse_args()

    input_csv = args.input_csv
    base_cols = [c.strip() for c in str(args.base_cols).split(",") if c.strip()]

    out_prefix = args.out_prefix
    if not out_prefix:
        out_prefix, _ = os.path.splitext(input_csv)

    service_out, cell_out = parse_split_csv(
        input_csv=input_csv,
        out_prefix=out_prefix,
        base_cols=base_cols,
        drop_raw=not bool(args.keep_raw),
        service_format=str(args.service_format),
    )

    print(service_out)
    print(cell_out)


if __name__ == "__main__":
    main()