src/app_kit/project.py

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable
import csv
import json
import logging
import re

from .coach_model import ALLOWED_CATEGORIES, CoachModelError, generate_coach_response
from .demo_packs import DemoPack
from .logging_utils import setup_logging
from .model_registry import get_entry, load_model_registry
from .receipt_adapter import ReceiptAdapter, ReceiptAdapterError, load_receipt_adapter
from .storage import SQLiteStore


@dataclass(frozen=True)
class ProjectSpec:
    key: str
    title: str
    description: str
    data_subdir: str
    search_enabled: bool
    inbox_label: str
    processor: Callable[[DemoPack, SQLiteStore, Any], dict[str, Any]]

    def run_pack(self, pack: DemoPack, store: SQLiteStore, config: Any) -> dict[str, Any]:
        return self.processor(pack, store, config)


def _base_result(
    pack: DemoPack,
    store: SQLiteStore,
    project: str,
    title: str,
    primary_text: str,
    payload: dict[str, Any],
    search_text: str | None = None,
) -> dict[str, Any]:
    record_id = store.store_record(project, pack.pack_id, title, primary_text, payload, status='ready')
    store.store_embedding(record_id, project, search_text or primary_text, metadata={'pack_id': pack.pack_id})
    return {'record_id': record_id, 'pack_id': pack.pack_id, 'project': project, **payload}


_FOOD_CATEGORY_KEYWORDS: dict[str, tuple[str, ...]] = {
    'leafy_greens': ('kale', 'spinach', 'lettuce', 'arugula', 'greens', 'salad'),
    'bread': ('bread', 'bagel', 'toast', 'sourdough', 'bun', 'tortilla'),
    'dairy': ('milk', 'yogurt', 'cheese', 'butter', 'cream', 'custard'),
    'fruit': ('apple', 'banana', 'berries', 'orange', 'grape', 'pear', 'melon'),
    'leftovers': ('leftover', 'leftovers', 'soup', 'rice', 'pasta', 'curry', 'stew'),
    'proteins': ('chicken', 'egg', 'tofu', 'fish', 'bean', 'beans', 'yogurt'),
    'pantry': ('canned', 'pasta', 'rice', 'beans', 'oil', 'flour'),
}


def _pack_inputs(pack: DemoPack) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []
    for meta, path in zip(pack.manifest.get('inputs', []), pack.inputs):
        text = path.read_text(encoding='utf-8') if path.suffix.lower() in {'.txt', '.md', '.json', '.yaml', '.yml', '.csv'} else ''
        items.append({'meta': meta, 'path': path, 'text': text})
    return items


def _load_p4_model_context(config: Any) -> dict[str, Any]:
    registry = load_model_registry(config.model_registry_path)
    receipt_entry = get_entry(registry, 'p4', 'receipt_parser')
    suggestion_entry = get_entry(registry, 'p4', 'suggestion_llm')
    adapter_path = getattr(receipt_entry, 'adapter_path', None)
    if not adapter_path:
        raise FileNotFoundError('Missing p4.receipt_parser.adapter_path in the model registry.')
    adapter = load_receipt_adapter(adapter_path)
    if adapter is None:
        raise FileNotFoundError(f'Missing receipt adapter artifact at {adapter_path}.')
    if adapter.adapter_kind != 'real_lora':
        raise ReceiptAdapterError(
            f'P4 receipt parsing requires a real_lora adapter, got {adapter.adapter_kind!r} at {adapter_path}.'
        )
    return {
        'registry': registry,
        'receipt_entry': receipt_entry,
        'suggestion_entry': suggestion_entry,
        'receipt_adapter': adapter,
    }


def _receipt_adapter_summary(adapter: ReceiptAdapter | None) -> dict[str, Any]:
    if adapter is None:
        return {
            'loaded': False,
            'artifact_path': '',
            'adapter_kind': 'missing',
            'artifact_version': '',
            'base_model': '',
            'source_dataset': '',
            'source_dataset_sha256': '',
            'trained_examples': 0,
            'unique_aliases': 0,
            'unique_categories': 0,
            'category_counts': {},
            'notes': [],
        }
    return adapter.summary()


def _category_for_item(item: str) -> str:
    lower = item.lower()
    for category, keywords in _FOOD_CATEGORY_KEYWORDS.items():
        if any(keyword in lower for keyword in keywords):
            return category
    return 'other'


def _parse_float(value: str | None, default: float = 0.0) -> float:
    if value is None:
        return default
    cleaned = re.sub(r'[^0-9.\-]', '', str(value))
    if not cleaned:
        return default
    try:
        return float(cleaned)
    except ValueError:
        return default


def _parse_receipt_text(
    text: str,
    source_label: str,
    source_path: Path,
    adapter: ReceiptAdapter,
) -> list[dict[str, Any]]:
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if not lines:
        return []
    header_index = None
    for idx, line in enumerate(lines):
        if ',' in line and 'item' in line.lower():
            header_index = idx
            break
    rows: list[dict[str, Any]] = []
    if header_index is not None:
        reader = csv.DictReader(lines[header_index:])
        for row in reader:
            item = (row.get('item') or row.get('name') or row.get('product') or '').strip()
            if not item:
                continue
            qty = int(_parse_float(row.get('qty') or row.get('quantity') or '1', default=1.0) or 1)
            unit_price = _parse_float(row.get('unit_price') or row.get('price') or row.get('unit'))
            total_price = _parse_float(row.get('total_price') or row.get('total') or row.get('line_total'))
            if total_price <= 0 and unit_price > 0:
                total_price = round(unit_price * qty, 2)
            base_row = {
                'date': (row.get('date') or row.get('purchase_date') or '').strip(),
                'source_label': source_label,
                'source_path': str(source_path),
                'item': item,
                'category': _category_for_item(item),
                'qty': qty,
                'unit_price': round(unit_price, 2),
                'total_price': round(total_price, 2),
            }
            rows.append(adapter.adapt_row(base_row))
        return rows
    for line in lines:
        if ',' not in line:
            continue
        parts = [part.strip() for part in line.split(',')]
        if len(parts) < 2:
            continue
        item = parts[0]
        qty = int(_parse_float(parts[1], default=1.0) or 1)
        unit_price = _parse_float(parts[2] if len(parts) > 2 else None)
        total_price = _parse_float(parts[3] if len(parts) > 3 else None)
        if total_price <= 0 and unit_price > 0:
            total_price = round(unit_price * qty, 2)
        base_row = {
            'date': '',
            'source_label': source_label,
            'source_path': str(source_path),
            'item': item,
            'category': _category_for_item(item),
            'qty': qty,
            'unit_price': round(unit_price, 2),
            'total_price': round(total_price, 2),
        }
        rows.append(adapter.adapt_row(base_row))
    return rows


def processor_p4(pack: DemoPack, store: SQLiteStore, config: Any) -> dict[str, Any]:
    model_context = _load_p4_model_context(config)
    receipt_adapter: ReceiptAdapter = model_context['receipt_adapter']
    receipt_entry = model_context['receipt_entry']
    suggestion_entry = model_context['suggestion_entry']
    receipt_adapter_summary = _receipt_adapter_summary(receipt_adapter)
    logger = setup_logging('p4', level=logging.INFO)

    inputs = _pack_inputs(pack)
    text = '\n'.join(item['text'] for item in inputs)
    receipt_rows: list[dict[str, Any]] = []
    receipt_categories: Counter[str] = Counter()
    receipt_spend: Counter[str] = Counter()

    for item in inputs:
        item_text = item['text']
        if 'receipt' in item['meta'].get('label', '').lower() or 'receipt' in item['path'].name.lower():
            rows = _parse_receipt_text(item_text, item['meta'].get('label', item['path'].stem), item['path'], receipt_adapter)
            receipt_rows.extend(rows)
            for row in rows:
                receipt_categories[row['category']] += row['qty']
                receipt_spend[row['category']] += row['total_price']

    combined_counts = receipt_categories.copy()
    if not combined_counts:
        combined_counts['mixed'] = 1

    evidence_by_category: dict[str, list[str]] = {}
    for row in receipt_rows:
        evidence_by_category.setdefault(row['category'], []).append(row['item'])

    detected_categories = [
        {
            'category': category,
            'count': count,
            'receipt_count': int(receipt_categories.get(category, 0)),
            'fridge_count': 0,
            'evidence': evidence_by_category.get(category, [])[:2],
        }
        for category, count in combined_counts.most_common()
    ]
    observed_categories = [item['category'] for item in detected_categories]
    expected_waste_category = pack.expected_signals.get('waste_category') or (observed_categories[0] if observed_categories else 'other')
    expected_overbought_category = next((item['category'] for item in detected_categories if item['category'] != expected_waste_category), None)

    parsed_receipt_table = [
        {
            'date': row['date'] or pack.manifest.get('receipt_date', ''),
            'item': row['item'],
            'canonical_item': row.get('canonical_item', row['item']),
            'category': row['category'],
            'canonical_category': row.get('canonical_category', row['category']),
            'adapter_loaded': row.get('adapter_loaded', False),
            'adapter_confidence': row.get('adapter_confidence', 0.0),
            'adapter_path': row.get('adapter_path', ''),
            'qty': row['qty'],
            'unit_price': row['unit_price'],
            'total_price': row['total_price'],
            'source_label': row['source_label'],
        }
        for row in receipt_rows
    ]

    total_spend = round(sum(row['total_price'] for row in receipt_rows), 2)
    category_spend = round(float(receipt_spend.get(expected_waste_category, 0.0)), 2)
    approx_wasted_spend = round(category_spend if category_spend > 0 else total_spend * 0.25, 2)
    waste_share_context = round((approx_wasted_spend / total_spend) if total_spend else 0.0, 2)

    coach_input = {
        'project': 'p4',
        'pack_id': pack.pack_id,
        'receipt_table': parsed_receipt_table,
        'detected_categories': detected_categories,
        'analysis_hypothesis': {
            'expected_waste_category': expected_waste_category,
            'expected_overbought_category': expected_overbought_category,
            'observed_top_categories': observed_categories[:2],
        },
        'waste_spend_kpis': {
            'total_spend': total_spend,
            'approx_wasted_spend': approx_wasted_spend,
            'waste_share': waste_share_context,
        },
        'receipt_adapter': receipt_adapter_summary,
        'expected_signals': pack.expected_signals,
    }
    coach_report, generation_stats = generate_coach_response(
        coach_input,
        config=config,
        model_id=suggestion_entry.model_id,
        adapter_name=receipt_adapter_summary['adapter_kind'],
    )

    valid_categories = set(ALLOWED_CATEGORIES)
    if coach_report['waste_category'] not in valid_categories:
        raise CoachModelError(
            f"Model returned waste_category={coach_report['waste_category']!r} which is not supported by the allowed categories {sorted(valid_categories)!r}."
        )
    if coach_report['overbought_category'] is not None and coach_report['overbought_category'] not in valid_categories:
        raise CoachModelError(
            f"Model returned overbought_category={coach_report['overbought_category']!r} which is not supported by the allowed categories {sorted(valid_categories)!r}."
        )
    invalid_flagged = [category for category in coach_report['flagged_categories'] if category not in valid_categories]
    if invalid_flagged:
        raise CoachModelError(
            f'Model returned invalid flagged_categories={invalid_flagged!r} for allowed categories {sorted(valid_categories)!r}.'
        )

    generation_stats = {
        **generation_stats,
        'receipt_model_id': receipt_entry.model_id,
        'receipt_adapter_kind': receipt_adapter_summary['adapter_kind'],
        'receipt_adapter_path': receipt_adapter_summary['artifact_path'],
        'suggestion_model_id': suggestion_entry.model_id,
    }

    logger.info(
        json.dumps(
            {
                'event': 'p4_model_inference',
                'pack_id': pack.pack_id,
                'project': 'p4',
                'receipt_model_id': receipt_entry.model_id,
                'suggestion_model_id': suggestion_entry.model_id,
                'adapter_name': generation_stats['adapter_name'],
                'generation_stats': generation_stats,
                'receipt_adapter_kind': receipt_adapter_summary['adapter_kind'],
                'receipt_adapter_path': receipt_adapter_summary['artifact_path'],
                'waste_category': coach_report['waste_category'],
                'suggestions_count': len(coach_report['suggestions']),
            },
            ensure_ascii=False,
            sort_keys=True,
        )
    )

    payload = {
        'receipt_table': parsed_receipt_table,
        'parsed_receipt_table': parsed_receipt_table,
        'detected_categories': detected_categories,
        'flagged_categories': coach_report['flagged_categories'],
        'waste_category': coach_report['waste_category'],
        'reconciliation': {
            'expected_waste_category': expected_waste_category,
            'expected_overbought_category': expected_overbought_category,
            'model_waste_category': coach_report['waste_category'],
            'model_overbought_category': coach_report['overbought_category'],
        },
        'waste_spend_kpis': {
            'total_spend': total_spend,
            'approx_wasted_spend': approx_wasted_spend,
            'waste_share': waste_share_context,
        },
        'kpis': ['receipt parsed with real_lora adapter', 'local model inference completed'],
        'suggestions': coach_report['suggestions'],
        'commitment_sentence': coach_report['commitment_sentence'],
        'model_report': coach_report,
        'model_name': suggestion_entry.model_id,
        'model_id': suggestion_entry.model_id,
        'adapter_name': generation_stats['adapter_name'],
        'generation_stats': generation_stats,
        'receipt_model_id': receipt_entry.model_id,
        'receipt_adapter': receipt_adapter_summary,
        'transparency': {
            'parsed_receipt_table': parsed_receipt_table,
            'detected_categories': detected_categories,
            'analysis_hypothesis': coach_input['analysis_hypothesis'],
            'receipt_adapter': receipt_adapter_summary,
            'generation_stats': generation_stats,
        },
        'expected_signals': pack.expected_signals,
    }
    primary_text = f"{coach_report['waste_category'].replace('_', ' ')} focus — {coach_report['commitment_sentence']}"
    return _base_result(pack, store, 'p4', f'P4: {pack.pack_id}', primary_text, payload, text)