#!/usr/bin/env python3
"""
corpus.py - Ginnie Corpus Capture (WO#GINNIE_09)
==================================================

Capture high-quality traces from Ginnie sessions for future training.

Every session leaves behind:
- Work order request
- Plan generated
- Steps executed
- Helix/symbolic metrics
- Success/failure info

This builds the dataset for "train our own codex later."
"""

import json
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional


@dataclass
class CorpusRecord:
    """A single training record from a Ginnie session."""
    wo_id: str
    request: str
    backend: str
    request_type: str
    plan: Dict[str, Any]
    steps: List[Dict[str, Any]]
    final_summary: str
    success: bool
    status: str
    duration_sec: float
    scars: int
    reflex_edits: int
    helix_metrics: Dict[str, Any] = field(default_factory=dict)
    symbolic_metrics: Dict[str, Any] = field(default_factory=dict)
    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    errors: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def to_jsonl(self) -> str:
        """Convert to JSONL format (single line)."""
        return json.dumps(self.to_dict(), separators=(',', ':'))


class CorpusWriter:
    """
    Write Ginnie session data to corpus files.
    """

    def __init__(self, corpus_dir: Optional[Path] = None):
        self.corpus_dir = corpus_dir or Path.home() / "helix-cdc/corpus/ginnie"
        self.corpus_dir.mkdir(parents=True, exist_ok=True)

    def _get_corpus_file(self, dated: bool = True) -> Path:
        """Get corpus file path."""
        if dated:
            date_str = datetime.now().strftime("%Y-%m-%d")
            return self.corpus_dir / f"ginnie_corpus_{date_str}.jsonl"
        return self.corpus_dir / "ginnie_corpus.jsonl"

    def append(self, record: CorpusRecord, dated: bool = True) -> Path:
        """
        Append a record to the corpus.

        Args:
            record: CorpusRecord to append
            dated: If True, use dated file (one per day)

        Returns:
            Path to corpus file
        """
        corpus_file = self._get_corpus_file(dated)

        with open(corpus_file, 'a') as f:
            f.write(record.to_jsonl() + '\n')

        return corpus_file

    def count_records(self, dated: bool = True) -> int:
        """Count records in corpus file."""
        corpus_file = self._get_corpus_file(dated)
        if not corpus_file.exists():
            return 0

        with open(corpus_file) as f:
            return sum(1 for _ in f)

    def get_stats(self) -> Dict[str, Any]:
        """Get corpus statistics."""
        stats = {
            "corpus_dir": str(self.corpus_dir),
            "files": [],
            "total_records": 0,
        }

        for jsonl_file in sorted(self.corpus_dir.glob("*.jsonl")):
            with open(jsonl_file) as f:
                count = sum(1 for _ in f)

            stats["files"].append({
                "name": jsonl_file.name,
                "records": count,
            })
            stats["total_records"] += count

        return stats


def record_session(
    wo_id: str,
    request: str,
    backend: str,
    request_type: str,
    plan: Dict[str, Any],
    steps: List[Dict[str, Any]],
    result: Dict[str, Any],
    helix_metrics: Optional[Dict[str, Any]] = None,
    symbolic_metrics: Optional[Dict[str, Any]] = None,
) -> CorpusRecord:
    """
    Create a CorpusRecord from session data.

    Args:
        wo_id: Work order ID
        request: Original request string
        backend: LLM backend used
        request_type: Classified request type
        plan: Plan dict (from planner)
        steps: List of step results
        result: Final work order result
        helix_metrics: Optional Helix tool metrics
        symbolic_metrics: Optional symbolic metrics

    Returns:
        CorpusRecord ready to append
    """
    # Count reflex edits
    reflex_edits = 0
    if plan.get("reflex_applied"):
        reflex_edits = len(plan.get("notes", []))

    return CorpusRecord(
        wo_id=wo_id,
        request=request,
        backend=backend,
        request_type=request_type,
        plan=plan,
        steps=steps,
        final_summary=result.get("description", ""),
        success=result.get("status") == "success",
        status=result.get("status", "unknown"),
        duration_sec=result.get("duration_s", 0),
        scars=result.get("scars", 0),
        reflex_edits=reflex_edits,
        helix_metrics=helix_metrics or {},
        symbolic_metrics=symbolic_metrics or {},
        errors=result.get("errors", []),
    )


def append_to_corpus(
    record: CorpusRecord,
    corpus_dir: Optional[Path] = None,
    dated: bool = True,
) -> Path:
    """
    Append a record to the corpus.

    Convenience wrapper around CorpusWriter.append().
    """
    writer = CorpusWriter(corpus_dir)
    return writer.append(record, dated)


# Global writer for convenience
_writer: Optional[CorpusWriter] = None


def get_corpus_writer() -> CorpusWriter:
    """Get or create global corpus writer."""
    global _writer
    if _writer is None:
        _writer = CorpusWriter()
    return _writer


def corpus_stats() -> Dict[str, Any]:
    """Get corpus statistics."""
    return get_corpus_writer().get_stats()
