from __future__ import annotations

import json
from dataclasses import asdict, is_dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Union


def ensure_workspace_dir(workspace_dir: Union[str, Path]) -> Path:
    """
    Ensure a local training workspace exists.

    Workspace layout:
      - examples.jsonl  (one JSON object per line; training examples)
      - contexts.jsonl  (one JSON object per line; training contexts)
      - workspace.json  (lightweight metadata)
    """
    ws = Path(workspace_dir).expanduser().resolve()
    ws.mkdir(parents=True, exist_ok=True)

    examples = ws / "examples.jsonl"
    contexts = ws / "contexts.jsonl"
    meta = ws / "workspace.json"

    if not examples.exists():
        examples.write_text("", encoding="utf-8")
    if not contexts.exists():
        contexts.write_text("", encoding="utf-8")
    if not meta.exists():
        meta.write_text(json.dumps({"version": 1}, indent=2), encoding="utf-8")

    return ws


def _to_jsonable(obj: Any) -> Any:
    if is_dataclass(obj):
        return asdict(obj)
    return obj


def append_jsonl(path: Union[str, Path], obj: Any) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    payload = _to_jsonable(obj)
    with open(p, "a", encoding="utf-8") as f:
        f.write(json.dumps(payload, ensure_ascii=False) + "\n")


def read_jsonl(path: Union[str, Path]) -> List[Dict[str, Any]]:
    p = Path(path)
    if not p.exists():
        return []
    out: List[Dict[str, Any]] = []
    with open(p, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            out.append(json.loads(line))
    return out


def add_example(
    *,
    workspace_dir: Union[str, Path],
    question: str,
    answer: str,
    domain: str = "general",
    task_type: str = "qa",
    source: str = "user",
    context: Optional[Dict[str, Any]] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> Path:
    ws = ensure_workspace_dir(workspace_dir)
    example = {
        "input_text": question,
        "expected_output": answer,
        "domain": domain,
        "task_type": task_type,
        "source": source,
        "context": context,
        "metadata": metadata or {},
    }
    append_jsonl(ws / "examples.jsonl", example)
    return ws / "examples.jsonl"


def add_examples_from_jsonl(
    *,
    workspace_dir: Union[str, Path],
    jsonl_path: Union[str, Path],
    default_domain: str = "general",
    default_task_type: str = "qa",
    default_source: str = "user",
) -> int:
    ws = ensure_workspace_dir(workspace_dir)
    raw = read_jsonl(jsonl_path)
    n = 0
    for row in raw:
        # Accept common keys:
        q = row.get("question") or row.get("input") or row.get("input_text")
        a = row.get("answer") or row.get("output") or row.get("expected_output")
        if not q or not a:
            continue
        add_example(
            workspace_dir=ws,
            question=str(q),
            answer=str(a),
            domain=str(row.get("domain") or default_domain),
            task_type=str(row.get("task_type") or default_task_type),
            source=str(row.get("source") or default_source),
            context=row.get("context"),
            metadata=row.get("metadata") if isinstance(row.get("metadata"), dict) else {},
        )
        n += 1
    return n


def add_contexts(
    *,
    workspace_dir: Union[str, Path],
    contexts: Iterable[Dict[str, Any]],
) -> int:
    ws = ensure_workspace_dir(workspace_dir)
    out_path = ws / "contexts.jsonl"
    n = 0
    for c in contexts:
        if not c:
            continue
        text = c.get("text") if isinstance(c, dict) else None
        if not text:
            continue
        append_jsonl(out_path, c)
        n += 1
    return n


def load_workspace_context_texts(workspace_dir: Union[str, Path], *, limit: Optional[int] = None) -> List[str]:
    ws = ensure_workspace_dir(workspace_dir)
    rows = read_jsonl(ws / "contexts.jsonl")
    texts: List[str] = []
    for r in rows:
        t = r.get("text")
        if isinstance(t, str) and t.strip():
            texts.append(t.strip())
        if limit is not None and len(texts) >= limit:
            break
    return texts


def load_workspace_examples(workspace_dir: Union[str, Path], *, limit: Optional[int] = None) -> List[Dict[str, Any]]:
    ws = ensure_workspace_dir(workspace_dir)
    rows = read_jsonl(ws / "examples.jsonl")
    out: List[Dict[str, Any]] = []
    for r in rows:
        q = r.get("input_text") or r.get("question") or r.get("input")
        a = r.get("expected_output") or r.get("answer") or r.get("output")
        if not q or not a:
            continue
        out.append(
            {
                "input_text": str(q),
                "expected_output": str(a),
                "domain": str(r.get("domain") or "general"),
                "task_type": str(r.get("task_type") or "qa"),
                "source": str(r.get("source") or "user"),
                "context": r.get("context"),
                "metadata": r.get("metadata") if isinstance(r.get("metadata"), dict) else {},
            }
        )
        if limit is not None and len(out) >= limit:
            break
    return out


def build_static_rules(
    *,
    examples: List[Dict[str, Any]],
    contexts: List[str],
    max_training_contexts: int = 20,
    max_example_phrases: int = 20,
) -> Dict[str, Any]:
    """
    Derive static training_contexts, key_rules, and example_phrases
    from workspace examples + contexts for injection into memory_state.json.
    """
    training_contexts: List[str] = []
    example_phrases: Dict[str, str] = {}

    # Use doc contexts as high-level training contexts (trim long ones).
    for text in contexts:
        if not text:
            continue
        snippet = text.strip()
        if len(snippet) > 400:
            snippet = snippet[:400] + "..."
        training_contexts.append(snippet)
        if len(training_contexts) >= max_training_contexts:
            break

    # Derive example phrases from Q&A pairs.
    for ex in examples:
        q = ex.get("input_text") or ""
        a = ex.get("expected_output") or ""
        if not q or not a:
            continue
        key = q[:80].strip()
        phrase = a.strip()
        if len(phrase) > 200:
            phrase = phrase[:200] + "..."
        example_phrases[key] = phrase
        if len(example_phrases) >= max_example_phrases:
            break

    # Simple aggregated key_rules summary.
    key_rules: Dict[str, str] = {}
    if training_contexts:
        key_rules["doc_training_summary"] = (
            f"{len(training_contexts)} training contexts derived from workspace documents."
        )
    if examples:
        key_rules["qa_training_summary"] = (
            f"{len(examples)} curated Q&A examples captured from workspace."
        )

    return {
        "training_contexts": training_contexts,
        "key_rules": key_rules,
        "example_phrases": example_phrases,
    }

