#!/usr/bin/env python3
"""
Generic bulk API runner for the Orchard platform.

Reads a CSV, calls the same API endpoint for every row using a JSON request
template. Fields in the template like {vendor_uuid} are filled from CSV row
data (after optional Marshmallow schema validation/coercion).

Auth headers are loaded from .env (TOKEN, ORCHARD_IDENTITY_ID, ORCHARD_IDENTITY_UUID,
ORCHARD_PROFILE_ID, ORCHARD_PROFILE_UUID, ORCHARD_PROFILE_TYPE).

Usage:
  python bulk_runner.py --config request.json --csv vendors.csv \\
      --base-url https://ows-account.theorchard.io

  # With schema validation + response consumer
  python bulk_runner.py --config request.json --csv vendors.csv \\
      --schema mymodule.schemas:VendorSchema \\
      --consumer mymodule.handlers:on_response

  # Dry run
  python bulk_runner.py --config request.json --csv vendors.csv --dry-run

request.json format:
  {
    "method": "PUT",
    "path": "/v2/vendors/{vendor_uuid}/internal-staff",
    "payload": {
      "product_manager": "{product_manager}",
      "status": "deletion"
    }
  }

  String values that are exactly "{field}" are replaced with the typed value
  from the row (preserving int/bool). Partial strings like "prefix-{field}"
  are always rendered as strings.
"""

import argparse
import csv
import importlib
import json
import logging
import os
import re
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

from dotenv import load_dotenv

try:
    import httpx
except ImportError:
    print('pip install httpx', file=sys.stderr)
    sys.exit(1)

log = logging.getLogger(__name__)

VALID_METHODS = {'GET', 'POST', 'PUT', 'PATCH', 'DELETE'}
MAX_RETRIES = 5
BASE_BACKOFF = 1.0  # seconds

# ── Request config ────────────────────────────────────────────────────────────


@dataclass
class RequestConfig:
    method: str
    path: str
    payload: dict[str, Any]

    @classmethod
    def from_file(cls, path: str) -> 'RequestConfig':
        p = Path(path)
        if not p.exists():
            log.error('Config file not found: %s', path)
            sys.exit(1)
        try:
            data = json.loads(p.read_text())
        except json.JSONDecodeError as e:
            log.error('Invalid JSON in config: %s', e)
            sys.exit(1)
        cfg = cls(
            method=str(data.get('method', '')).upper(),
            path=str(data.get('path', '')),
            payload=data.get('payload') or {},
        )
        cfg.validate()
        return cfg

    def validate(self) -> None:
        if self.method not in VALID_METHODS:
            log.error(
                'Invalid method %r. Must be one of: %s',
                self.method,
                ', '.join(VALID_METHODS),
            )
            sys.exit(1)
        if not self.path.startswith('/'):
            log.error("path must start with '/', got: %r", self.path)
            sys.exit(1)


# ── Template rendering ────────────────────────────────────────────────────────

_EXACT_PLACEHOLDER = re.compile(r'^\{(\w+)\}$')
_PARTIAL_PLACEHOLDER = re.compile(r'\{(\w+)\}')


def _render_value(value: Any, row: dict[str, Any]) -> Any:
    if not isinstance(value, str):
        return value  # int/bool/None literals pass through unchanged

    exact = _EXACT_PLACEHOLDER.match(value)
    if exact:
        key = exact.group(1)
        if key not in row:
            raise KeyError(f"Template key '{key}' not found in row")
        val = row[key]
        if isinstance(val, str) and val.lower() in ('true', 'false'):
            return val.lower() == 'true'
        return val

    def _sub(m: re.Match[str]) -> str:
        key = m.group(1)
        if key not in row:
            raise KeyError(f"Template key '{key}' not found in row")
        return str(row[key])

    return _PARTIAL_PLACEHOLDER.sub(_sub, value)


def render_template(template: Any, row: dict[str, Any]) -> Any:
    if isinstance(template, dict):
        return {k: render_template(v, row) for k, v in template.items()}
    if isinstance(template, list):
        return [render_template(item, row) for item in template]
    return _render_value(template, row)


# ── Import loading ────────────────────────────────────────────────────────────


def load_import(import_path: str) -> Any:
    if ':' not in import_path:
        raise ValueError(f"Import path must be 'module.path:Name', got: {import_path!r}")
    module_path, attr_name = import_path.rsplit(':', 1)
    try:
        module = importlib.import_module(module_path)
    except ModuleNotFoundError as e:
        raise ImportError(f"Cannot import module '{module_path}': {e}") from e
    if not hasattr(module, attr_name):
        raise ImportError(f"Module '{module_path}' has no attribute '{attr_name}'")
    return getattr(module, attr_name)


def load_schema(import_path: str | None) -> Any:
    if not import_path:
        return None
    try:
        import marshmallow  # noqa: F401
    except ImportError:
        log.error('pip install marshmallow to use --schema')
        sys.exit(1)
    try:
        cls = load_import(import_path)
        return cls()
    except (ImportError, ValueError) as e:
        log.error('Failed to load schema %r: %s', import_path, e)
        sys.exit(1)


def load_consumer(import_path: str | None) -> Callable[..., Any] | None:
    if not import_path:
        return None
    try:
        fn: Callable[..., Any] = load_import(import_path)
        return fn
    except (ImportError, ValueError) as e:
        log.error('Failed to load consumer %r: %s', import_path, e)
        sys.exit(1)


# ── HTTP engine ───────────────────────────────────────────────────────────────


def build_client() -> httpx.Client:
    return httpx.Client(
        headers={
            'Content-Type': 'application/json',
            # Required by the prod graphql-router (require_apollo_client_name plugin);
            # ignored by ows services, so safe to send on every request.
            'apollographql-client-name': 'bulk-api-caller',
        },
        timeout=30,
    )


def load_env_headers() -> dict[str, str]:
    load_dotenv()
    required = {
        'TOKEN': 'Authorization',
        'ORCHARD_IDENTITY_ID': 'Orchard-Identity-Id',
        'ORCHARD_PROFILE_ID': 'Orchard-Profile-Id',
        'ORCHARD_PROFILE_TYPE': 'Orchard-Profile-Type',
    }
    optional = {
        'ORCHARD_IDENTITY_UUID': 'Orchard-Identity-Uuid',
        'ORCHARD_PROFILE_UUID': 'Orchard-Profile-Uuid',
    }
    missing = [k for k in required if not os.environ.get(k)]
    if missing:
        log.error('Missing required .env vars: %s', ', '.join(missing))
        sys.exit(1)
    headers: dict[str, str] = {}
    for env_key, header_name in required.items():
        val = os.environ[env_key]
        if env_key == 'TOKEN':
            headers['Authorization'] = f'Bearer {val}'
        else:
            headers[header_name] = val
    for env_key, header_name in optional.items():
        opt_val = os.environ.get(env_key)
        if opt_val is not None:
            headers[header_name] = opt_val
    return headers


def execute_request(
    client: httpx.Client,
    base_url: str,
    config: RequestConfig,
    env_headers: dict[str, str],
    row: dict[str, Any],
) -> dict[str, Any]:
    rendered_path = render_template(config.path, row)
    rendered_payload = render_template(config.payload, row)
    url = base_url.rstrip('/') + rendered_path

    for attempt in range(1, MAX_RETRIES + 1):
        resp = client.request(
            method=config.method,
            url=url,
            json=rendered_payload or None,
            headers=env_headers,
        )
        if resp.status_code == 429:
            retry_after = float(resp.headers.get('Retry-After', BASE_BACKOFF * (2 ** (attempt - 1))))
            log.warning(
                'Rate limited (attempt %d/%d). Sleeping %.1fs',
                attempt,
                MAX_RETRIES,
                retry_after,
            )
            time.sleep(retry_after)
            continue
        resp.raise_for_status()
        result: dict[str, Any] = resp.json()
        return result

    raise httpx.HTTPStatusError(
        f'Exceeded {MAX_RETRIES} retries on 429',
        request=resp.request,
        response=resp,
    )


# ── Row result ────────────────────────────────────────────────────────────────


@dataclass
class RowResult:
    row_num: int
    row: dict[str, Any]
    success: bool
    status_code: int | None = None
    response_body: dict[str, Any] | None = None
    error: str | None = None


# ── Row worker ────────────────────────────────────────────────────────────────


def process_row(
    row_num: int,
    raw_row: dict[str, Any],
    schema: Any,
    config: RequestConfig,
    env_headers: dict[str, str],
    client: httpx.Client,
    base_url: str,
    delay: float,
    consumer: Callable[..., Any] | None,
    dry_run: bool,
) -> RowResult:
    # Schema validation + coercion
    if schema is not None:
        import marshmallow

        try:
            row = schema.load(raw_row)
        except marshmallow.ValidationError as e:
            log.warning('Row %d validation failed: %s', row_num, e.messages)
            return RowResult(
                row_num=row_num,
                row=raw_row,
                success=False,
                error=f'validation: {e.messages}',
            )
    else:
        row = raw_row

    if dry_run:
        log.info('[DRY RUN] Row %d: %s %s', row_num, config.method, config.path)
        return RowResult(row_num=row_num, row=row, success=True)

    try:
        response = execute_request(client, base_url, config, env_headers, row)
        if consumer is not None:
            try:
                consumer(row, response)
            except Exception as e:
                log.warning('Row %d consumer error: %s', row_num, e)
        if delay:
            time.sleep(delay)
        return RowResult(row_num=row_num, row=row, success=True, response_body=response)
    except httpx.HTTPStatusError as e:
        status = e.response.status_code
        body = e.response.text[:500]
        log.warning('Row %d HTTP %s: %s', row_num, status, body)
        return RowResult(row_num=row_num, row=raw_row, success=False, status_code=status, error=body)
    except Exception as e:
        log.warning('Row %d error: %s', row_num, e)
        return RowResult(row_num=row_num, row=raw_row, success=False, error=str(e))


# ── Runner ────────────────────────────────────────────────────────────────────


def run(args: argparse.Namespace) -> None:
    csv_path = Path(args.csv)
    if not csv_path.exists():
        log.error('CSV not found: %s', csv_path)
        sys.exit(1)

    repo_root = Path(__file__).resolve().parent.parent
    output_dir = Path(args.output_dir) if args.output_dir else repo_root / 'data' / 'output'
    output_dir.mkdir(parents=True, exist_ok=True)

    config = RequestConfig.from_file(args.config)
    schema = load_schema(args.schema)
    consumer = load_consumer(args.consumer)

    with open(csv_path, newline='', encoding='utf-8-sig') as f:
        rows = list(csv.DictReader(f))

    log.info('Loaded %d rows from %s', len(rows), csv_path)

    if args.skip:
        rows = rows[args.skip :]
        log.info('Skipping first %d rows; %d remaining', args.skip, len(rows))

    if args.dry_run:
        log.info('DRY RUN — no requests will be made')

    client = build_client()
    env_headers = load_env_headers()

    failures: list[dict[str, Any]] = []
    responses: list[dict[str, Any]] = []
    failure_lock = threading.Lock()
    response_lock = threading.Lock()
    succeeded = 0
    failed = 0

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {
            pool.submit(
                process_row,
                i + args.skip + 1,
                row,
                schema,
                config,
                env_headers,
                client,
                args.base_url,
                args.delay,
                consumer,
                args.dry_run,
            ): row
            for i, row in enumerate(rows)
        }

        for future in as_completed(futures):
            result: RowResult = future.result()
            if result.success:
                succeeded += 1
                if result.response_body is not None:
                    with response_lock:
                        responses.append({'row': result.row_num, 'response': result.response_body})
            else:
                failed += 1
                with failure_lock:
                    failures.append(
                        {
                            'row': result.row_num,
                            'data': result.row,
                            'status': result.status_code,
                            'error': result.error,
                        }
                    )

            total = succeeded + failed
            if total % 500 == 0:
                log.info(
                    'Progress: %d/%d (succeeded=%d, failed=%d)',
                    total,
                    len(rows),
                    succeeded,
                    failed,
                )

    log.info('Done. %d succeeded, %d failed out of %d', succeeded, failed, len(rows))

    stem = csv_path.stem
    if responses:
        out = output_dir / f'{stem}_responses.json'
        out.write_text(json.dumps(responses, indent=2, default=str), encoding='utf-8')
        log.info('Responses written to %s', out)
    if failures:
        out = output_dir / f'{stem}_failures.json'
        out.write_text(json.dumps(failures, indent=2, default=str), encoding='utf-8')
        log.info('Failures written to %s', out)


# ── CLI ───────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description='Generic bulk API runner for the Orchard platform',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  python bulk_runner.py --config request.json --csv vendors.csv \\
      --base-url https://ows-account.theorchard.io

  python bulk_runner.py --config request.json --csv vendors.csv \\
      --schema mymodule.schemas:VendorSchema \\
      --consumer mymodule.handlers:on_response \\
      --workers 20 --delay 0.02

  python bulk_runner.py --config request.json --csv vendors.csv --dry-run

Auth headers (TOKEN, ORCHARD_IDENTITY_ID, ORCHARD_PROFILE_ID, ORCHARD_PROFILE_TYPE)
must be set in .env.
        """,
    )

    parser.add_argument(
        '--config',
        required=True,
        metavar='request.json',
        help='Path to request template JSON file',
    )
    parser.add_argument('--csv', required=True, help='Path to input CSV file')

    parser.add_argument(
        '--base-url',
        default='https://ows-grass.theorchard.io/account',
        help='API base URL (default: ows-account prod)',
    )

    parser.add_argument(
        '--schema',
        metavar='module.path:ClassName',
        help='Marshmallow schema for CSV row validation + coercion',
    )
    parser.add_argument(
        '--consumer',
        metavar='module.path:fn_name',
        help='Callable (row, response) -> None invoked per successful response',
    )

    parser.add_argument(
        '--workers',
        type=int,
        default=10,
        help='Concurrent worker threads (default: 10)',
    )
    parser.add_argument(
        '--delay',
        type=float,
        default=0.05,
        help='Per-worker delay after each request in seconds (default: 0.05)',
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Log what would be sent without making requests',
    )
    parser.add_argument(
        '--skip',
        type=int,
        default=0,
        help='Skip first N rows (for resuming a partial run)',
    )
    parser.add_argument(
        '--output-dir',
        help='Directory for output JSON files (default: <repo>/data/output)',
    )

    args = parser.parse_args()

    run(args)


if __name__ == '__main__':
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s %(levelname)s %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler('bulk_runner.log'),
        ],
    )
    main()