import argparse
from dataclasses import asdict
from pathlib import Path
import json


def main():
    p = argparse.ArgumentParser(description="NEXORA research and local runtime")
    sub = p.add_subparsers(dest="command", required=True)
    t = sub.add_parser("train")
    t.add_argument("--config", default="configs/tiny.json")
    t.add_argument("--data", default="artifacts/data")
    t.add_argument("--output", default="artifacts/tiny")
    t.add_argument("--resume", action="store_true")
    t.add_argument("--stop-after", type=int)
    d = sub.add_parser("prepare")
    d.add_argument("input")
    d.add_argument("--output", default="artifacts/data")
    g = sub.add_parser("generate")
    g.add_argument("prompt")
    g.add_argument("--checkpoint", default="artifacts/tiny")
    g.add_argument("--tokens", type=int, default=64)
    c = sub.add_parser("chat")
    c.add_argument("prompt")
    c.add_argument("--model", default=".cache/Qwen3.5-0.8B")
    c.add_argument("--tokens", type=int, default=128)
    s = sub.add_parser("serve")
    s.add_argument("--model", default=".cache/Qwen3.5-0.8B")
    s.add_argument("--port", type=int, default=8765)
    i = sub.add_parser("index")
    i.add_argument("root", nargs="?", default=".")
    i.add_argument("--query", default="")
    a = sub.add_parser("agent")
    a.add_argument("task")
    a.add_argument("--policy", required=True)
    a.add_argument("--model", default=".cache/Qwen3.5-0.8B")
    a.add_argument("--url")
    a.add_argument("--allow-network", action="store_true")
    a.add_argument("--verify-command", help="Owner-configured command key; should refer to trusted external tests")
    a.add_argument("--steps", type=int, default=8)
    e = sub.add_parser("estimate")
    e.add_argument("--parameters", type=float, default=120e9)
    e.add_argument("--active", type=float)
    e.add_argument("--tokens", type=float, default=2.4e12)
    e.add_argument("--gpus", type=int, default=1024)
    e.add_argument("--mfu", type=float, default=.35)
    args = p.parse_args()
    if args.command == "train":
        from .training import train
        result = train(args.config, args.data, args.output, resume=args.resume, stop_after=args.stop_after)
    elif args.command == "prepare":
        from .data import prepare
        result = prepare([json.loads(l) for l in Path(args.input).read_text(encoding="utf-8").splitlines() if l.strip()], args.output)
    elif args.command == "generate":
        from .inference import tiny_generate
        result = tiny_generate(args.checkpoint, args.prompt, args.tokens)
    elif args.command == "chat":
        from .inference import HFBackend
        backend = HFBackend(args.model, max_new_tokens=args.tokens)
        result = {"answer": backend.complete([{"role": "system", "content": "You are NEXORA, a concise assistant. Be accurate about your limitations."}, {"role": "user", "content": args.prompt}]), "metrics": backend.last_metrics}
    elif args.command == "serve":
        import os
        from .inference import HFBackend
        from .server import create_server
        token = os.environ.get("NEXORA_LOCAL_TOKEN", "")
        if len(token) < 16:
            raise ValueError("Set NEXORA_LOCAL_TOKEN to at least 16 random characters")
        server = create_server(HFBackend(args.model), token, args.port)
        print(f"Local inference listening at http://127.0.0.1:{args.port}; Ctrl+C stops it", flush=True)
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            pass
        finally:
            server.server_close()
        return
    elif args.command == "estimate":
        from .compute import Estimate
        result = Estimate(args.parameters, args.active or args.parameters, args.tokens, args.gpus, mfu=args.mfu).calculate()
    elif args.command == "index":
        from .tools import Policy, Executor
        from .coding import index_repository, retrieve
        result = index_repository(Executor(Policy(args.root)))
        if args.query:
            result = retrieve(result, args.query)
    else:
        from .tools import Policy, Executor
        from .inference import HFBackend, HTTPBackend
        from .agent import Agent
        executor = Executor(Policy(**json.loads(Path(args.policy).read_text(encoding="utf-8"))))
        backend = HTTPBackend(args.url, args.model, allow_network=args.allow_network) if args.url else HFBackend(args.model)
        verifier = None
        if args.verify_command:
            def verifier():
                r = executor.execute("shell.exec", {"command": args.verify_command})
                return r.ok and r.exit_code == 0 and not r.truncated
        result = Agent(backend, executor, verifier, max_steps=args.steps).run(args.task)
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()