import argparse
import json
import os
from pathlib import Path
from typing import Dict, List

from .client import ToxoCloudClient


def _api_key(value: str | None) -> str | None:
    return value or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")


def _collect_docs(paths: List[str], recursive: bool, pattern: str | None) -> List[Path]:
    out: List[Path] = []
    for raw in paths:
        p = Path(raw)
        if p.is_file():
            out.append(p)
            continue
        if p.is_dir():
            itr = p.rglob(pattern or "*") if recursive else p.glob(pattern or "*")
            out.extend([x for x in itr if x.is_file()])
    return out


def _read_jsonl_examples(path: str | None) -> List[Dict[str, str]]:
    if not path:
        return []
    p = Path(path)
    if not p.exists():
        raise SystemExit(f"Examples file not found: {p}")
    examples: List[Dict[str, str]] = []
    for idx, raw in enumerate(p.read_text(encoding="utf-8").splitlines(), start=1):
        line = raw.strip()
        if not line:
            continue
        try:
            row = json.loads(line)
        except json.JSONDecodeError as exc:
            raise SystemExit(f"Invalid JSONL at line {idx} in {p}: {exc}") from exc
        inp = row.get("input") or row.get("question") or row.get("input_text")
        out = row.get("output") or row.get("answer") or row.get("expected_output")
        if not inp or not out:
            raise SystemExit(
                f"Invalid example at line {idx} in {p}: needs input/output (or question/answer)."
            )
        examples.append({"input": str(inp), "output": str(out)})
    return examples


def _read_contexts(path: str | None) -> List[str]:
    if not path:
        return []
    p = Path(path)
    if not p.exists():
        raise SystemExit(f"Contexts file not found: {p}")
    if p.suffix.lower() == ".json":
        data = json.loads(p.read_text(encoding="utf-8"))
        if not isinstance(data, list):
            raise SystemExit("Contexts JSON must be an array of strings.")
        return [str(x) for x in data if str(x).strip()]
    if p.suffix.lower() == ".jsonl":
        out: List[str] = []
        for idx, raw in enumerate(p.read_text(encoding="utf-8").splitlines(), start=1):
            line = raw.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise SystemExit(f"Invalid JSONL at line {idx} in {p}: {exc}") from exc
            txt = row.get("text") if isinstance(row, dict) else None
            if txt:
                out.append(str(txt))
        return out
    return [x for x in p.read_text(encoding="utf-8").splitlines() if x.strip()]


def main() -> None:
    parser = argparse.ArgumentParser(prog="toxo-cloud", description="TOXO Cloud CLI")
    sub = parser.add_subparsers(dest="command")

    root_p = sub.add_parser("root", help="Get service root info")

    health_p = sub.add_parser("health", help="Health check")

    query_p = sub.add_parser("query", help="Text query")
    query_p.add_argument("layer")
    query_p.add_argument("question")
    query_p.add_argument("--api-key")
    query_p.add_argument("--provider", default="gemini")
    query_p.add_argument("--model", default="gemini-2.5-flash-lite")
    query_p.add_argument("--response-depth", default="balanced")

    ask_doc_p = sub.add_parser("ask-doc", help="Multimodal document query")
    ask_doc_p.add_argument("layer")
    ask_doc_p.add_argument("document")
    ask_doc_p.add_argument("question")
    ask_doc_p.add_argument("--api-key")
    ask_doc_p.add_argument("--provider", default="gemini")
    ask_doc_p.add_argument("--model", default="gemini-2.5-flash-lite")
    ask_doc_p.add_argument("--mime-type")
    ask_doc_p.add_argument("--response-depth", default="balanced")

    ask_image_p = sub.add_parser("ask-image", help="Multimodal image query")
    ask_image_p.add_argument("layer")
    ask_image_p.add_argument("image")
    ask_image_p.add_argument("question")
    ask_image_p.add_argument("--api-key")
    ask_image_p.add_argument("--provider", default="gemini")
    ask_image_p.add_argument("--model", default="gemini-2.5-flash-lite")
    ask_image_p.add_argument("--mime-type")
    ask_image_p.add_argument("--response-depth", default="balanced")

    feedback_p = sub.add_parser("feedback", help="Send feedback")
    feedback_p.add_argument("layer")
    feedback_p.add_argument("question")
    feedback_p.add_argument("response")
    feedback_p.add_argument("--rating", type=float)
    feedback_p.add_argument("--suggestion", action="append", default=[])

    train_auto_p = sub.add_parser("train-auto", help="Auto-train from description")
    train_auto_p.add_argument("description")
    train_auto_p.add_argument("--out", required=True)
    train_auto_p.add_argument("--num-examples", type=int, default=5)
    train_auto_p.add_argument("--api-key")
    train_auto_p.add_argument("--provider", default="gemini")
    train_auto_p.add_argument("--model")

    train_data_p = sub.add_parser("train-from-data", help="Train from explicit examples and contexts")
    train_data_p.add_argument("--description", required=True)
    train_data_p.add_argument("--out", required=True)
    train_data_p.add_argument("--examples-jsonl", help="JSONL with input/output or question/answer keys")
    train_data_p.add_argument("--contexts-file", help="TXT (line-delimited), JSON array, or JSONL with {text}")
    train_data_p.add_argument("--epochs", type=int, default=10)
    train_data_p.add_argument("--domain")
    train_data_p.add_argument("--api-key")
    train_data_p.add_argument("--provider", default="gemini")
    train_data_p.add_argument("--model")

    extract_ctx_p = sub.add_parser("train-extract-contexts", help="Extract contexts from docs")
    extract_ctx_p.add_argument("--docs", nargs="+", required=True)
    extract_ctx_p.add_argument("--api-key")
    extract_ctx_p.add_argument("--recursive", action="store_true")
    extract_ctx_p.add_argument("--pattern")

    rag_index_p = sub.add_parser("rag-index", help="Build cloud RAG index")
    rag_index_p.add_argument("--index-id", required=True)
    rag_index_p.add_argument("--docs", nargs="+", required=True)
    rag_index_p.add_argument("--domain", default="general")
    rag_index_p.add_argument("--layer")
    rag_index_p.add_argument("--api-key")
    rag_index_p.add_argument("--recursive", action="store_true")
    rag_index_p.add_argument("--pattern")

    rag_query_p = sub.add_parser("rag-query", help="Query cloud RAG index")
    rag_query_p.add_argument("--layer", required=True)
    rag_query_p.add_argument("--index-id", required=True)
    rag_query_p.add_argument("--question", required=True)
    rag_query_p.add_argument("--api-key")
    rag_query_p.add_argument("--provider", default="gemini")
    rag_query_p.add_argument("--model")
    rag_query_p.add_argument("--top-k", type=int, default=5)
    rag_query_p.add_argument("--response-depth", default="balanced")

    agent_p = sub.add_parser("agent-run", help="Run cloud agent workflow")
    agent_p.add_argument("--config", required=True)
    agent_p.add_argument("--layer", action="append", default=[], help="agent_id=path/to/layer.toxo")
    agent_p.add_argument("--api-key")
    agent_p.add_argument("--provider", default="gemini")
    agent_p.add_argument("--model")
    agent_p.add_argument("--json-out")

    args = parser.parse_args()
    if not args.command:
        parser.print_help()
        return

    client = ToxoCloudClient(api_key=_api_key(getattr(args, "api_key", None)))

    if args.command == "root":
        print(json.dumps(client.root(), indent=2))
    elif args.command == "health":
        print(json.dumps(client.health(), indent=2))
    elif args.command == "query":
        print(client.query(args.layer, args.question, provider=args.provider, model=args.model, response_depth=args.response_depth))
    elif args.command == "ask-doc":
        print(
            client.query_multimodal(
                args.layer,
                args.question,
                document_path=args.document,
                mime_type=args.mime_type,
                provider=args.provider,
                model=args.model,
                response_depth=args.response_depth,
            )
        )
    elif args.command == "ask-image":
        print(
            client.query_multimodal(
                args.layer,
                args.question,
                image_path=args.image,
                mime_type=args.mime_type,
                provider=args.provider,
                model=args.model,
                response_depth=args.response_depth,
            )
        )
    elif args.command == "feedback":
        print(json.dumps(client.feedback(args.layer, args.question, args.response, rating=args.rating, suggestions=args.suggestion), indent=2))
    elif args.command == "train-auto":
        out = client.train_from_data(
            description=args.description,
            out_path=args.out,
            mode="auto_from_description",
            num_examples=args.num_examples,
            llm_provider=args.provider,
            llm_model=args.model,
        )
        print(json.dumps(out.get("summary", {}), indent=2))
    elif args.command == "train-from-data":
        examples = _read_jsonl_examples(args.examples_jsonl)
        contexts = _read_contexts(args.contexts_file)
        out = client.train_from_data(
            description=args.description,
            out_path=args.out,
            examples=examples or None,
            contexts=contexts or None,
            epochs=args.epochs,
            domain=args.domain,
            llm_provider=args.provider,
            llm_model=args.model,
        )
        print(json.dumps(out.get("summary", out), indent=2))
    elif args.command == "train-extract-contexts":
        docs = _collect_docs(args.docs, recursive=args.recursive, pattern=args.pattern)
        data = client.train_extract_contexts(docs=[str(d) for d in docs])
        print(json.dumps(data, indent=2))
    elif args.command == "rag-index":
        docs = _collect_docs(args.docs, recursive=args.recursive, pattern=args.pattern)
        data = client.rag_index(index_id=args.index_id, docs=[str(d) for d in docs], domain=args.domain, layer=args.layer)
        print(json.dumps(data, indent=2))
    elif args.command == "rag-query":
        data = client.rag_query(
            layer=args.layer,
            index_id=args.index_id,
            question=args.question,
            llm_provider=args.provider,
            llm_model=args.model,
            top_k=args.top_k,
            response_depth=args.response_depth,
        )
        print(json.dumps(data, indent=2))
    elif args.command == "agent-run":
        layer_map: Dict[str, str] = {}
        for entry in args.layer:
            if "=" not in entry:
                raise SystemExit("Invalid --layer format. Use agent_id=path/to/layer.toxo")
            k, v = entry.split("=", 1)
            layer_map[k.strip()] = v.strip()
        data = client.agent_run(
            config=args.config,
            layer_map=layer_map,
            llm_provider=args.provider,
            llm_model=args.model,
        )
        if args.json_out:
            p = Path(args.json_out)
            p.parent.mkdir(parents=True, exist_ok=True)
            p.write_text(json.dumps(data, indent=2), encoding="utf-8")
        print(data.get("final") or json.dumps(data, indent=2))


if __name__ == "__main__":
    main()
