"""
Data ingestion pipeline for the agricultural treatment assistant system.

Reads local product profile JSON files and label-chunk JSONL files, then loads
them into the configured backends:

- **Weaviate** (``ProductProfile``) — semantic ``retrieval_summary`` for profiles
- **SQL Server** — ``products``, ``product_uses`` (profiles), ``label_chunks`` (chunks)

The pipeline is idempotent: a hash manifest (``.processed_manifest.json``)
tracks which files have already been ingested (and their content hash), so
re-running the script skips unchanged files unless ``--force`` is passed.

Usage:
    python ingest.py
    python ingest.py --profiles-dir data/productProfiles
    python ingest.py --chunks-dir data/chunks
    python ingest.py --file data/productProfiles/ramin_sc_0916.json
    python ingest.py --chunk-file data/chunks/delan_pro_16562.jsonl
    python ingest.py --force
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import uuid
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, Optional

import pyodbc
import weaviate
from dotenv import load_dotenv

from setup_weaviate import ensure_product_profile_collection

COLLECTION_NAME = "ProductProfile"
MANIFEST_PATH = Path(".processed_manifest.json")
DEFAULT_PROFILES_DIR = Path("data/productProfiles")
DEFAULT_CHUNKS_DIR = Path("data/chunks")

# Namespace used to derive deterministic Weaviate UUIDs from product_id, so
# re-ingesting the same product always maps to the same object.
WEAVIATE_UUID_NAMESPACE = uuid.UUID("1b7e2c3a-4f5d-4a8b-9c1e-6d2f3a8b9c0d")

ARRAY_FIELDS = (
    "active_ingredients",
    "frac_groups",
    "target_crops",
    "target_diseases",
    "action_type",
)

USE_FIELDS = (
    "crop",
    "disease",
    "setting",
    "dose_min",
    "dose_max",
    "dose_unit",
    "concentration_min",
    "concentration_max",
    "concentration_unit",
    "treatment_interval_min_days",
    "treatment_interval_max_days",
    "max_treatments_per_season",
    "pre_harvest_interval_days",
    "spray_volume_min",
    "spray_volume_max",
    "growth_stage_start",
    "growth_stage_end",
)

CHUNK_INSERT_COLUMNS = (
    "chunk_id",
    "product_id",
    "product_name",
    "target_crops",
    "target_diseases",
    "chunk_type",
    "chunk_text",
)

CHUNK_ARRAY_FIELDS = ("target_crops", "target_diseases")
CHUNK_REQUIRED_KEYS = ("chunk_id", "product_id", "chunk_text")


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

def load_config() -> Dict[str, str]:
    """Load and validate required environment variables from .env."""
    load_dotenv()

    required = [
        "GEMINI_API_KEY",
        "SQL_DRIVER",
        "SQL_SERVER",
        "SQL_DATABASE",
        "SQL_USERNAME",
        "SQL_PASSWORD",
    ]
    missing = [key for key in required if not os.getenv(key)]
    if missing:
        raise EnvironmentError(
            f"Missing required environment variable(s): {', '.join(missing)}. "
            "Copy .env.example to .env and fill in the values."
        )

    return {
        "GEMINI_API_KEY": os.getenv("GEMINI_API_KEY", ""),
        "WEAVIATE_HOST": os.getenv("WEAVIATE_HOST", "localhost"),
        "WEAVIATE_HTTP_PORT": os.getenv("WEAVIATE_HTTP_PORT", "8080"),
        "WEAVIATE_GRPC_PORT": os.getenv("WEAVIATE_GRPC_PORT", "50051"),
        "SQL_DRIVER": os.getenv("SQL_DRIVER", ""),
        "SQL_SERVER": os.getenv("SQL_SERVER", ""),
        "SQL_DATABASE": os.getenv("SQL_DATABASE", ""),
        "SQL_USERNAME": os.getenv("SQL_USERNAME", ""),
        "SQL_PASSWORD": os.getenv("SQL_PASSWORD", ""),
    }


# ---------------------------------------------------------------------------
# File discovery, hashing, and manifest tracking
# ---------------------------------------------------------------------------

def discover_files(data_dir: Path, pattern: str) -> list[Path]:
    """Recursively find files matching pattern under data_dir, sorted for stable order."""
    if not data_dir.exists():
        raise FileNotFoundError(f"Data directory not found: {data_dir}")
    return sorted(data_dir.rglob(pattern))


def discover_json_files(data_dir: Path) -> list[Path]:
    """Recursively find all *.json files under data_dir."""
    return discover_files(data_dir, "*.json")


def discover_jsonl_files(data_dir: Path) -> list[Path]:
    """Recursively find all *.jsonl files under data_dir."""
    return discover_files(data_dir, "*.jsonl")


def sha256_of(path: Path) -> str:
    """Compute the SHA-256 hex digest of a file's contents."""
    digest = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_manifest(manifest_path: Path = MANIFEST_PATH) -> Dict[str, str]:
    """Load the processed-files manifest (filename -> content hash)."""
    if not manifest_path.exists():
        return {}
    try:
        with manifest_path.open("r", encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError) as exc:
        print(f"[WARN] Could not read manifest at {manifest_path}: {exc}. Starting fresh.", file=sys.stderr)
        return {}


def save_manifest(manifest: Dict[str, str], manifest_path: Path = MANIFEST_PATH) -> None:
    """Persist the processed-files manifest to disk."""
    with manifest_path.open("w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, sort_keys=True)


def should_process(path: Path, file_hash: str, manifest: Dict[str, str], force: bool) -> bool:
    """Determine whether a file needs (re)processing based on the manifest."""
    if force:
        return True
    key = str(path.resolve())
    return manifest.get(key) != file_hash


def collect_files_to_process(
    candidates: list[Path],
    manifest: Dict[str, str],
    force: bool,
) -> list[tuple[Path, str]]:
    """Filter candidates through the manifest and return (path, hash) pairs to process."""
    to_process: list[tuple[Path, str]] = []
    for path in candidates:
        try:
            file_hash = sha256_of(path)
        except OSError as exc:
            print(f"[ERROR] Could not read {path}: {exc}", file=sys.stderr)
            continue
        if should_process(path, file_hash, manifest, force):
            to_process.append((path, file_hash))
        else:
            print(f"[SKIP] {path} already processed and unchanged.")
    return to_process


# ---------------------------------------------------------------------------
# Parsing
# ---------------------------------------------------------------------------

def parse_profile(path: Path) -> Dict[str, Any]:
    """Load and lightly validate a product profile JSON file."""
    with path.open("r", encoding="utf-8") as f:
        profile = json.load(f)

    required_keys = ("product_id", "product_name", "retrieval_summary")
    missing = [key for key in required_keys if key not in profile]
    if missing:
        raise ValueError(f"Profile {path} is missing required key(s): {', '.join(missing)}")

    return profile


def parse_chunks_file(path: Path) -> list[Dict[str, Any]]:
    """
    Load a JSONL label-chunk file, validating each non-empty line.

    Malformed lines are logged to stderr and skipped; returns all valid chunks.
    """
    chunks: list[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as f:
        for line_number, line in enumerate(f, start=1):
            stripped = line.strip()
            if not stripped:
                continue
            try:
                chunk = json.loads(stripped)
            except json.JSONDecodeError as exc:
                print(
                    f"[ERROR] {path} line {line_number}: invalid JSON: {exc}",
                    file=sys.stderr,
                )
                continue
            if not isinstance(chunk, dict):
                print(
                    f"[ERROR] {path} line {line_number}: expected a JSON object, got {type(chunk).__name__}",
                    file=sys.stderr,
                )
                continue
            missing = [key for key in CHUNK_REQUIRED_KEYS if key not in chunk]
            if missing:
                print(
                    f"[ERROR] {path} line {line_number}: missing required key(s): {', '.join(missing)}",
                    file=sys.stderr,
                )
                continue
            chunks.append(chunk)
    return chunks


# ---------------------------------------------------------------------------
# Connections
# ---------------------------------------------------------------------------

def connect_weaviate(cfg: Dict[str, str]) -> weaviate.WeaviateClient:
    """Open a connection to the local Weaviate instance with the Gemini API key header."""
    client = weaviate.connect_to_local(
        host=cfg["WEAVIATE_HOST"],
        port=int(cfg["WEAVIATE_HTTP_PORT"]),
        grpc_port=int(cfg["WEAVIATE_GRPC_PORT"]),
        headers={"X-Goog-Studio-Api-Key": cfg["GEMINI_API_KEY"]},
    )
    return client


def connect_sql(cfg: Dict[str, str]) -> pyodbc.Connection:
    """Open a connection to the Microsoft SQL Server database via pyodbc."""
    conn_str = (
        f"DRIVER={{{cfg['SQL_DRIVER']}}};"
        f"SERVER={cfg['SQL_SERVER']};"
        f"DATABASE={cfg['SQL_DATABASE']};"
        f"UID={cfg['SQL_USERNAME']};"
        f"PWD={cfg['SQL_PASSWORD']};"
    )
    conn = pyodbc.connect(conn_str, autocommit=False)
    return conn


# ---------------------------------------------------------------------------
# Weaviate insertion
# ---------------------------------------------------------------------------

def upsert_weaviate(client: weaviate.WeaviateClient, profile: Dict[str, Any]) -> None:
    """Insert or replace the semantic summary object for this product in Weaviate."""
    collection = client.collections.get(COLLECTION_NAME)

    product_id = profile["product_id"]
    object_uuid = uuid.uuid5(WEAVIATE_UUID_NAMESPACE, product_id)

    properties = {
        "product_id": product_id,
        "product_name": profile["product_name"],
        "retrieval_summary": profile["retrieval_summary"],
    }

    if collection.data.exists(object_uuid):
        collection.data.replace(uuid=object_uuid, properties=properties)
    else:
        collection.data.insert(properties=properties, uuid=object_uuid)


# ---------------------------------------------------------------------------
# SQL Server: products table
# ---------------------------------------------------------------------------

def upsert_product(cursor: pyodbc.Cursor, profile: Dict[str, Any]) -> None:
    """Upsert the global metadata row for this product into the products table."""
    array_json = {field: json.dumps(profile.get(field, [])) for field in ARRAY_FIELDS}

    params = (
        profile["product_id"],
        profile.get("product_name"),
        profile.get("registration_number"),
        profile.get("manufacturer"),
        int(bool(profile.get("organic_certified"))),
        profile.get("systemicity"),
        int(bool(profile.get("preventive_action"))),
        int(bool(profile.get("curative_action"))),
        int(bool(profile.get("eradicant_action"))),
        array_json["active_ingredients"],
        array_json["frac_groups"],
        array_json["target_crops"],
        array_json["target_diseases"],
        array_json["action_type"],
    )

    cursor.execute(
        """
        IF EXISTS (SELECT 1 FROM products WHERE product_id = ?)
            UPDATE products
            SET product_name = ?,
                registration_number = ?,
                manufacturer = ?,
                organic_certified = ?,
                systemicity = ?,
                preventive_action = ?,
                curative_action = ?,
                eradicant_action = ?,
                active_ingredients = ?,
                frac_groups = ?,
                target_crops = ?,
                target_diseases = ?,
                action_type = ?
            WHERE product_id = ?
        ELSE
            INSERT INTO products (
                product_id, product_name, registration_number, manufacturer,
                organic_certified, systemicity, preventive_action, curative_action,
                eradicant_action, active_ingredients, frac_groups, target_crops,
                target_diseases, action_type
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        # IF EXISTS check
        profile["product_id"],
        # UPDATE ... SET (13 params) + WHERE product_id
        *params[1:],
        profile["product_id"],
        # INSERT ... VALUES (14 params)
        *params,
    )


# ---------------------------------------------------------------------------
# SQL Server: product_uses table
# ---------------------------------------------------------------------------

def replace_uses(cursor: pyodbc.Cursor, profile: Dict[str, Any]) -> None:
    """Replace all use-rule rows for this product with the ones from the profile."""
    product_id = profile["product_id"]
    uses: Iterable[Dict[str, Any]] = profile.get("uses", [])

    cursor.execute("DELETE FROM product_uses WHERE product_id = ?", product_id)

    insert_sql = f"""
        INSERT INTO product_uses (
            product_id, {", ".join(USE_FIELDS)}
        )
        VALUES (?, {", ".join(["?"] * len(USE_FIELDS))})
    """

    for index, use in enumerate(uses):
        try:
            values = (product_id, *(use.get(field) for field in USE_FIELDS))
            cursor.execute(insert_sql, values)
        except pyodbc.Error as exc:
            print(
                f"[ERROR] Failed to insert use rule #{index} for product '{product_id}': {exc}",
                file=sys.stderr,
            )
            raise


# ---------------------------------------------------------------------------
# SQL Server: label_chunks table
# ---------------------------------------------------------------------------

def chunk_insert_values(chunk: Dict[str, Any]) -> tuple[Any, ...]:
    """Build parameterized insert values for a single label chunk row."""
    values: list[Any] = []
    for column in CHUNK_INSERT_COLUMNS:
        if column in CHUNK_ARRAY_FIELDS:
            values.append(json.dumps(chunk.get(column, [])))
        else:
            values.append(chunk.get(column))
    return tuple(values)


def replace_label_chunks(
    cursor: pyodbc.Cursor,
    product_id: str,
    chunks: Iterable[Dict[str, Any]],
) -> None:
    """Replace all label-chunk rows for a product with the provided set."""
    cursor.execute("DELETE FROM label_chunks WHERE product_id = ?", product_id)

    insert_sql = f"""
        INSERT INTO label_chunks (
            {", ".join(CHUNK_INSERT_COLUMNS)}
        )
        VALUES ({", ".join(["?"] * len(CHUNK_INSERT_COLUMNS))})
    """

    for chunk in chunks:
        chunk_id = chunk.get("chunk_id", "<unknown>")
        try:
            cursor.execute(insert_sql, chunk_insert_values(chunk))
        except pyodbc.Error as exc:
            print(
                f"[ERROR] Failed to insert label chunk '{chunk_id}' for product '{product_id}': {exc}",
                file=sys.stderr,
            )
            raise


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------

def process_profile_file(
    path: Path,
    weaviate_client: weaviate.WeaviateClient,
    sql_conn: pyodbc.Connection,
    manifest: Dict[str, str],
    file_hash: str,
) -> bool:
    """
    Process a single product profile JSON file across Weaviate and SQL sinks.

    Returns True on full success, False if any sink failed. On SQL failure the
    transaction for this file is rolled back so no partial rows are committed.
    """
    try:
        profile = parse_profile(path)
    except (json.JSONDecodeError, ValueError, OSError) as exc:
        print(f"[ERROR] Failed to parse {path}: {exc}", file=sys.stderr)
        return False

    product_id = profile.get("product_id", "<unknown>")

    weaviate_ok = True
    try:
        upsert_weaviate(weaviate_client, profile)
    except Exception as exc:  # noqa: BLE001 - log and continue to SQL sinks
        weaviate_ok = False
        print(f"[ERROR] Weaviate upsert failed for '{product_id}' ({path}): {exc}", file=sys.stderr)

    sql_ok = True
    cursor = sql_conn.cursor()
    try:
        upsert_product(cursor, profile)
        replace_uses(cursor, profile)
        sql_conn.commit()
    except pyodbc.Error as exc:
        sql_ok = False
        sql_conn.rollback()
        print(f"[ERROR] SQL upsert failed for '{product_id}' ({path}): {exc}. Transaction rolled back.", file=sys.stderr)
    finally:
        cursor.close()

    success = weaviate_ok and sql_ok
    if success:
        manifest[str(path.resolve())] = file_hash
        print(f"[OK] Processed profile '{product_id}' from {path}")
    else:
        print(f"[SKIP] Profile '{product_id}' from {path} not fully ingested; manifest not updated.", file=sys.stderr)

    return success


def process_chunks_file(
    path: Path,
    sql_conn: pyodbc.Connection,
    manifest: Dict[str, str],
    file_hash: str,
) -> bool:
    """
    Process a single label-chunk JSONL file into the label_chunks table.

    Returns True on success. On SQL failure the transaction is rolled back.
    """
    try:
        chunks = parse_chunks_file(path)
    except OSError as exc:
        print(f"[ERROR] Failed to read {path}: {exc}", file=sys.stderr)
        return False

    if not chunks:
        print(f"[ERROR] No valid chunks found in {path}; manifest not updated.", file=sys.stderr)
        return False

    chunks_by_product: dict[str, list[Dict[str, Any]]] = defaultdict(list)
    for chunk in chunks:
        chunks_by_product[chunk["product_id"]].append(chunk)

    product_ids = ", ".join(sorted(chunks_by_product))
    cursor = sql_conn.cursor()
    try:
        for product_id, product_chunks in sorted(chunks_by_product.items()):
            replace_label_chunks(cursor, product_id, product_chunks)
        sql_conn.commit()
    except pyodbc.Error as exc:
        sql_conn.rollback()
        print(
            f"[ERROR] SQL label_chunks upsert failed for product(s) [{product_ids}] ({path}): {exc}. "
            "Transaction rolled back.",
            file=sys.stderr,
        )
        return False
    finally:
        cursor.close()

    manifest[str(path.resolve())] = file_hash
    print(f"[OK] Processed {len(chunks)} label chunk(s) for product(s) [{product_ids}] from {path}")
    return True


def run(
    profiles_dir: Path,
    chunks_dir: Path,
    single_profile: Optional[Path],
    single_chunk: Optional[Path],
    force: bool,
) -> int:
    """Main pipeline execution. Returns a process exit code."""
    cfg = load_config()
    manifest = load_manifest()

    if single_profile is not None:
        profile_candidates = [single_profile]
    else:
        profile_candidates = discover_json_files(profiles_dir)

    if single_chunk is not None:
        chunk_candidates = [single_chunk]
    elif chunks_dir.exists():
        chunk_candidates = discover_jsonl_files(chunks_dir)
    else:
        chunk_candidates = []
        if single_profile is None:
            print(f"[WARN] Chunks directory not found: {chunks_dir}; skipping chunk ingestion.")

    profiles_to_process = collect_files_to_process(profile_candidates, manifest, force)
    chunks_to_process = collect_files_to_process(chunk_candidates, manifest, force)

    if not profiles_to_process and not chunks_to_process:
        print("Nothing to process.")
        return 0

    weaviate_client: Optional[weaviate.WeaviateClient] = None
    sql_conn: Optional[pyodbc.Connection] = None
    failures = 0

    try:
        sql_conn = connect_sql(cfg)

        if profiles_to_process:
            weaviate_client = connect_weaviate(cfg)
            ensure_product_profile_collection(weaviate_client)

        for path, file_hash in profiles_to_process:
            try:
                success = process_profile_file(path, weaviate_client, sql_conn, manifest, file_hash)
            except Exception as exc:  # noqa: BLE001 - never let one file crash the batch
                success = False
                print(f"[ERROR] Unexpected failure processing profile {path}: {exc}", file=sys.stderr)
            if not success:
                failures += 1

        for path, file_hash in chunks_to_process:
            try:
                success = process_chunks_file(path, sql_conn, manifest, file_hash)
            except Exception as exc:  # noqa: BLE001
                success = False
                print(f"[ERROR] Unexpected failure processing chunks {path}: {exc}", file=sys.stderr)
            if not success:
                failures += 1

        save_manifest(manifest)
    finally:
        if weaviate_client is not None:
            try:
                weaviate_client.close()
            except Exception as exc:  # noqa: BLE001
                print(f"[WARN] Error closing Weaviate client: {exc}", file=sys.stderr)
        if sql_conn is not None:
            try:
                sql_conn.close()
            except Exception as exc:  # noqa: BLE001
                print(f"[WARN] Error closing SQL connection: {exc}", file=sys.stderr)

    total_files = len(profiles_to_process) + len(chunks_to_process)
    if failures:
        print(f"Completed with {failures} failure(s) out of {total_files} file(s).", file=sys.stderr)
        return 1

    print(f"Completed successfully. {total_files} file(s) processed.")
    return 0


def parse_args(argv: Optional[list[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Ingest product profile JSON and label-chunk JSONL files into Weaviate and SQL Server."
    )
    parser.add_argument(
        "--profiles-dir",
        type=Path,
        default=DEFAULT_PROFILES_DIR,
        help="Directory to scan recursively for *.json profile files (default: data/productProfiles)",
    )
    parser.add_argument(
        "--chunks-dir",
        type=Path,
        default=DEFAULT_CHUNKS_DIR,
        help="Directory to scan recursively for *.jsonl chunk files (default: data/chunks)",
    )
    parser.add_argument(
        "--file",
        type=Path,
        default=None,
        help="Process a single profile JSON file instead of scanning --profiles-dir.",
    )
    parser.add_argument(
        "--chunk-file",
        type=Path,
        default=None,
        help="Process a single label-chunk JSONL file instead of scanning --chunks-dir.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Reprocess files even if unchanged since the last successful run.",
    )
    return parser.parse_args(argv)


def main() -> None:
    args = parse_args()
    try:
        exit_code = run(
            args.profiles_dir,
            args.chunks_dir,
            args.file,
            args.chunk_file,
            args.force,
        )
    except (EnvironmentError, FileNotFoundError) as exc:
        print(f"[FATAL] {exc}", file=sys.stderr)
        sys.exit(1)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
