main.py

#!/usr/bin/env python3
"""
Main CLI for Finland Job Market Visualizer data pipeline.
Orchestrates core data processing scripts.
"""

import argparse
import subprocess
import sys
import os
from pathlib import Path
from dotenv import load_dotenv

# Load .env file before checking environment variables
load_dotenv()

# Get the root directory (where this script is located)
ROOT_DIR = Path(__file__).parent.resolve()
SRC_DIR = ROOT_DIR / "src"


def run_script(script_path, args=None):
    """Run a Python script as a subprocess."""
    cmd = [sys.executable, str(script_path)]
    if args:
        cmd.extend(args)
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, cwd=ROOT_DIR)
    return result.returncode


def tree_level(value):
    """argparse type: a classification level (StatFin publishes levels 1-5)."""
    level = int(value)
    if not 1 <= level <= 5:
        raise argparse.ArgumentTypeError(f"{value} is not a level from 1 to 5")
    return level


def cmd_fetch(args):
    """Fetch data from StatFin and optionally Barometer."""
    print("=== Fetching StatFin data ===")
    statfin_args = ["--config", "data/statfin_config.json"]
    if args.config:
        statfin_args = ["--config", args.config]

    returncode = run_script(SRC_DIR / "fetch_statfin.py", statfin_args)
    if returncode != 0:
        print(f"Error: fetch_statfin.py failed with code {returncode}")
        return returncode

    if args.barometer_input:
        if not os.path.exists(args.barometer_input):
            print(f"Warning: Barometer CSV file not found: {args.barometer_input}")
            print("Barometer data is optional. Skipping barometer processing.")
            print("The pipeline will use existing extras.csv for outlook data.")
            # Skip barometer processing but continue with pipeline
        else:
            print("\n=== Fetching Barometer outlook ===")
            barometer_args = [
                "--input",
                args.barometer_input,
                "--output",
                args.barometer_output or "data/barometer_outlook.csv",
                "--code-column",
                args.code_column or "isco_code",
                "--outlook-desc-column",
                args.outlook_desc_column or "balance_label",
            ]
            if args.outlook_pct_column:
                barometer_args.extend(["--outlook-pct-column", args.outlook_pct_column])

            returncode = run_script(SRC_DIR / "fetch_barometer.py", barometer_args)
            if returncode != 0:
                print(f"Error: fetch_barometer.py failed with code {returncode}")
                return returncode

    return 0


def cmd_infer(args):
    """Score AI exposure (missing pay is estimated by build_site_data.py)."""
    # Use LLM scoring when any LLM backend is configured
    llm_configured = any(
        os.environ.get(var)
        for var in ("ANTHROPIC_API_KEY", "GOOGLE_API_KEY", "VERTEX_PROJECT")
    ) or os.environ.get("LLM_BACKEND", "auto") != "auto"
    if llm_configured:
        print("\n=== Scoring AI exposure with LLM ===")
        # score.py picks the default model for its backend (CLAUDE_MODEL / LLM_MODEL)
        score_args = []
        if getattr(args, "model", None):
            score_args = ["--model", args.model]
        returncode = run_script(SRC_DIR / "score.py", score_args)
        if returncode != 0:
            print(f"Error: score.py failed with code {returncode}")
            return returncode
    else:
        print("\n=== Scoring AI exposure (rule-based fallback) ===")
        print("Note: Set ANTHROPIC_API_KEY (or GOOGLE_API_KEY) in .env for LLM-based scoring")
        returncode = run_script(SRC_DIR / "infer_ai_exposure.py")
        if returncode != 0:
            print(f"Error: infer_ai_exposure.py failed with code {returncode}")
            return returncode

    return 0


def cmd_build(args):
    """Build site data."""
    print("=== Building site data ===")
    build_args = []
    # Without the flag, TREE_LEVEL from the environment / .env applies (as in score.py)
    if args.tree_level is not None:
        os.environ["TREE_LEVEL"] = str(args.tree_level)

    returncode = run_script(SRC_DIR / "build_site_data.py", build_args)
    return returncode


def cmd_process(args):
    """Process HTML files and generate prompts (deprecated US data)."""
    print("=== Processing HTML files to Markdown ===")
    process_args = []
    if args.force:
        process_args.append("--force")

    returncode = run_script(ROOT_DIR / "process.py", process_args)
    if returncode != 0:
        print(f"Error: process.py failed with code {returncode}")
        return returncode

    print("\n=== Generating LLM prompt ===")
    returncode = run_script(SRC_DIR / "make_prompt.py")
    return returncode


def cmd_all(args):
    """Run the complete pipeline."""
    print("Running complete data pipeline...")

    # Step 1: Fetch
    fetch_parser = argparse.ArgumentParser()
    fetch_args = fetch_parser.parse_args([])
    fetch_args.config = args.config
    fetch_args.barometer_input = args.barometer_input
    fetch_args.barometer_output = args.barometer_output
    fetch_args.code_column = args.code_column
    fetch_args.outlook_desc_column = args.outlook_desc_column
    fetch_args.outlook_pct_column = args.outlook_pct_column

    returncode = cmd_fetch(fetch_args)
    if returncode != 0:
        return returncode

    # Step 2: Intermediate build. Without an LLM key, the rule-based fallback
    # (infer_ai_exposure.py) reads site/data.json to find unscored occupations.
    # A second build (Step 4) produces the final site/data.json with the scores.
    print("\n=== Building initial site data (for unscored occupations) ===")
    build_parser = argparse.ArgumentParser()
    build_args = build_parser.parse_args([])
    build_args.tree_level = args.tree_level
    returncode = cmd_build(build_args)
    if returncode != 0:
        return returncode

    # Step 3: Score AI exposure
    infer_parser = argparse.ArgumentParser()
    infer_args = infer_parser.parse_args([])
    infer_args.model = getattr(args, "model", None)
    returncode = cmd_infer(infer_args)
    if returncode != 0:
        return returncode

    # Step 4: Build final (with the new scores)
    print("\n=== Building final site data ===")
    returncode = cmd_build(build_args)
    if returncode != 0:
        return returncode

    print("Pipeline completed successfully!")
    return 0


def main():
    parser = argparse.ArgumentParser(
        description="Finland Job Market Visualizer - Data Pipeline CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Fetch StatFin data only
  finlandjobs fetch
  
  # Fetch with Barometer outlook
  finlandjobs fetch --barometer-input barometer.csv
  
   # Score AI exposure with LLM
  finlandjobs infer
  finlandjobs infer --model claude-haiku-5-5
  
  # Build site data
  finlandjobs build
  
  # Run complete pipeline (set ANTHROPIC_API_KEY in .env for LLM scoring)
  finlandjobs all --barometer-input barometer.csv
  
  (Note: The finlandjobs script automatically detects and uses the project's virtual environment. Run from the project directory.)
        """,
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Fetch command
    fetch_parser = subparsers.add_parser(
        "fetch", help="Fetch data from StatFin and optionally Barometer"
    )
    fetch_parser.add_argument(
        "--config",
        help="Path to StatFin config JSON (default: data/statfin_config.json)",
    )
    fetch_parser.add_argument("--barometer-input", help="Path to Barometer CSV file")
    fetch_parser.add_argument(
        "--barometer-output", help="Output path for Barometer outlook CSV"
    )
    fetch_parser.add_argument(
        "--code-column", help="Column name for occupation codes in Barometer CSV"
    )
    fetch_parser.add_argument(
        "--outlook-desc-column", help="Column name for outlook description"
    )
    fetch_parser.add_argument(
        "--outlook-pct-column", help="Column name for outlook percentage"
    )
    fetch_parser.set_defaults(func=cmd_fetch)

    # Infer command
    infer_parser = subparsers.add_parser(
        "infer", help="Score AI exposure (LLM, or a rule-based fallback)"
    )
    infer_parser.add_argument(
        "--model",
        help="LLM model for AI exposure scoring (default: CLAUDE_MODEL / LLM_MODEL from .env)",
    )
    infer_parser.set_defaults(func=cmd_infer)

    # Build command
    build_parser = subparsers.add_parser("build", help="Build site data")
    build_parser.add_argument(
        "--tree-level",
        type=tree_level,
        default=None,
        help=(
            "Finest hierarchy level; levels 1..N are built "
            "(default: TREE_LEVEL from env, else 4)"
        ),
    )
    build_parser.set_defaults(func=cmd_build)

    # Process command (deprecated US data)
    process_parser = subparsers.add_parser(
        "process", help="Process HTML files and generate prompts (deprecated US data)"
    )
    process_parser.add_argument(
        "--force", action="store_true", help="Re-process even if .md exists"
    )
    process_parser.set_defaults(func=cmd_process)

    # All command
    all_parser = subparsers.add_parser(
        "all",
        help="Run complete pipeline (fetch + initial build + infer + final build)",
    )
    all_parser.add_argument("--config", help="Path to StatFin config JSON")
    all_parser.add_argument("--barometer-input", help="Path to Barometer CSV file")
    all_parser.add_argument(
        "--barometer-output", help="Output path for Barometer outlook CSV"
    )
    all_parser.add_argument("--code-column", help="Column name for occupation codes")
    all_parser.add_argument(
        "--outlook-desc-column", help="Column name for outlook description"
    )
    all_parser.add_argument(
        "--outlook-pct-column", help="Column name for outlook percentage"
    )
    all_parser.add_argument(
        "--tree-level",
        type=tree_level,
        default=None,
        help=(
            "Finest hierarchy level, built and LLM-scored "
            "(default: TREE_LEVEL from env, else 4)"
        ),
    )
    all_parser.add_argument(
        "--model",
        help="LLM model for AI exposure scoring (default: CLAUDE_MODEL / LLM_MODEL from .env)",
    )
    all_parser.set_defaults(func=cmd_all)

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return 1

    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())