#!/usr/bin/env python3
"""parse_json_logs.py

Parse JSON logs into a Pandas DataFrame, normalize nested fields, extract
structured values from free-text messages with regex, and run validation checks.

Examples:
  python parse_json_logs.py --input logs.jsonl --output parsed.csv
  python parse_json_logs.py --input logs.json --output parsed.parquet
  python parse_json_logs.py --input logs.jsonl --message-column message --preview 10
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Iterable, List, Optional

import pandas as pd


DEFAULT_REGEX = (
    r"user=(?P<user>\w+)"
    r".*?req_id=(?P<request_id>[A-Za-z0-9_-]+)"
    r".*?status=(?P<status>\d+)"
    r"(?:.*?code=(?P<error_code>[A-Za-z0-9_-]+))?"
)


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Parse JSON logs with Pandas and regex.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--input", required=True, help="Path to a JSON or JSON Lines log file.")
    parser.add_argument(
        "--output",
        help="Optional output path. Use .csv or .parquet. If omitted, prints a preview to stdout.",
    )
    parser.add_argument(
        "--lines",
        action="store_true",
        help="Treat input as JSON Lines. Use this for one JSON object per line.",
    )
    parser.add_argument(
        "--message-column",
        default="message",
        help="Column containing free-text log messages for regex extraction.",
    )
    parser.add_argument(
        "--regex",
        default=DEFAULT_REGEX,
        help="Regex pattern with named capture groups for message extraction.",
    )
    parser.add_argument(
        "--preview",
        type=int,
        default=5,
        help="Number of rows to display when no output file is provided.",
    )
    return parser


def validate_input_path(path_str: str) -> Path:
    path = Path(path_str)
    if not path.exists():
        raise FileNotFoundError(f"Input file not found: {path}")
    if not path.is_file():
        raise ValueError(f"Input path is not a file: {path}")
    return path


def load_json_records(path: Path, lines: bool) -> pd.DataFrame:
    try:
        if lines:
            return pd.read_json(path, lines=True)
        return pd.read_json(path)
    except ValueError:
        records: List[dict] = []
        with path.open("r", encoding="utf-8") as f:
            for line_no, line in enumerate(f, start=1):
                stripped = line.strip()
                if not stripped:
                    continue
                try:
                    records.append(json.loads(stripped))
                except json.JSONDecodeError as exc:
                    raise ValueError(f"Invalid JSON on line {line_no}: {exc}") from exc
        if not records:
            raise ValueError("No JSON records were loaded from the input file.")
        return pd.json_normalize(records)


def normalize_if_needed(df: pd.DataFrame) -> pd.DataFrame:
    # If nested dictionaries are present as object cells, json_normalize on records
    # is usually better. This function is intentionally conservative.
    return df.copy()


def extract_message_fields(df: pd.DataFrame, message_column: str, regex: str) -> pd.DataFrame:
    if message_column not in df.columns:
        return df

    pattern = re.compile(regex)
    extracted = df[message_column].astype("string").str.extract(pattern)
    if extracted is None:
        return df
    return pd.concat([df, extracted], axis=1)


def coerce_common_types(df: pd.DataFrame) -> pd.DataFrame:
    for col in ["status"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64")
    return df


def validate_dataframe(df: pd.DataFrame, message_column: str) -> None:
    if df.empty:
        raise ValueError("Parsed DataFrame is empty.")

    print(f"Rows: {len(df)}")
    print(f"Columns: {len(df.columns)}")
    print("Missing values by column:")
    print(df.isna().sum().to_string())
    print("Data types:")
    print(df.dtypes.to_string())

    if message_column in df.columns:
        non_null_messages = df[message_column].notna().sum()
        if non_null_messages == 0:
            raise ValueError(f"Column '{message_column}' exists but contains no usable values.")


def write_output(df: pd.DataFrame, output: str) -> None:
    out = Path(output)
    suffix = out.suffix.lower()
    if suffix == ".csv":
        df.to_csv(out, index=False)
    elif suffix == ".parquet":
        df.to_parquet(out, index=False)
    else:
        raise ValueError("Unsupported output format. Use .csv or .parquet.")


def main(argv: Optional[Iterable[str]] = None) -> int:
    parser = build_arg_parser()
    args = parser.parse_args(argv)

    try:
        input_path = validate_input_path(args.input)
        df = load_json_records(input_path, args.lines)
        df = normalize_if_needed(df)
        df = extract_message_fields(df, args.message_column, args.regex)
        df = coerce_common_types(df)
        validate_dataframe(df, args.message_column)

        if args.output:
            write_output(df, args.output)
            print(f"Saved parsed output to: {args.output}")
        else:
            preview_rows = max(args.preview, 0)
            print(df.head(preview_rows).to_string(index=False))

        return 0
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())