cli.py

#!/usr/bin/env python3
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from core.database import AnalystDB
from core.models import DataSource, Session, Query
from core.loader import load_data, infer_schema, get_preview
from core.executor import analyze


def cmd_load(args):
    db = AnalystDB()
    df, meta = load_data(
        path=Path(args.file) if args.file else None,
        source_type=args.type,
        raw_text=args.text or "",
        table_name=args.table or "",
    )
    schema = infer_schema(df)
    preview = get_preview(df, 5)

    src = DataSource(
        name=meta.get("filename", args.file or "pasted"),
        source_type=meta["source_type"],
        filename=meta.get("filename", ""),
        table_name=meta.get("table", ""),
        row_count=schema["row_count"],
        column_count=schema["column_count"],
        size_bytes=int(schema.get("memory_bytes", 0)),
        schema_json=json.dumps(schema),
        preview_json=json.dumps(preview),
    )
    src_id = db.add_source(src)
    print(f"  Source ID: {src_id}")
    print(f"  {schema['row_count']} rows x {schema['column_count']} columns")
    print(f"  Columns: {', '.join(c['name'] for c in schema['columns'])}")
    db.close()


def cmd_ask(args):
    api_key = args.api_key or os.getenv("GROQ_API_KEY", "")
    if not api_key:
        print("  Error: GROQ_API_KEY not set")
        return

    db = AnalystDB()
    src = db.get_source(args.source)
    if not src:
        print(f"  Source {args.source} not found")
        return

    # Rebuild DataFrame from preview (for small analysis) or from file
    schema = json.loads(src["schema_json"])
    df, _ = load_data(
        path=Path(src["filename"]) if src["filename"] and Path(src["filename"]).exists() else None,
        source_type=src["source_type"],
    )

    # Create/get session
    ses = Session(source_id=args.source, title=args.question[:100])
    ses_id = db.add_session(ses)

    # Get conversation history
    history = []
    if args.session:
        for q in db.get_queries(args.session):
            history.append({"question": q["question"], "code": q["code"]})
    elif args.continuous:
        existing = db.get_queries(db.get_sessions(args.source)[0]["id"] if db.get_sessions(args.source) else 0)
        for q in existing:
            history.append({"question": q["question"], "code": q["code"]})

    print(f"  Analyzing: {args.question}")
    result = analyze(args.question, df, schema, api_key, history if history else None)

    query = Query(
        session_id=ses_id,
        question=args.question,
        code=result.get("code", ""),
        result_text=result.get("result_text", ""),
        error=result.get("error", ""),
        chart_b64=result.get("chart_b64", ""),
        execution_time_ms=result.get("execution_time_ms", 0),
        tokens_used=result.get("tokens_used", 0),
    )
    db.add_query(query)
    db.close()

    if result.get("error"):
        print(f"  Error: {result['error'][:200]}")
    if result.get("result_text"):
        print(f"\n{result['result_text']}\n")
    if result.get("chart_b64"):
        print(f"  [Chart generated: {len(result['chart_b64'])} bytes base64]")
    print(f"  [{result.get('execution_time_ms', 0)}ms, {result.get('tokens_used', 0)} tokens, {result.get('retries', 0)} retries]")


def cmd_list(args):
    db = AnalystDB()
    if args.type == "sources":
        for r in db.get_sources():
            print(f"  [{r['id']}] {r['name'][:40]:40s} | {r['source_type']:8s} | {r['row_count']}r x {r['column_count']}c | {r['created_at'][:16]}")
    elif args.type == "sessions":
        for r in db.get_sessions():
            print(f"  [{r['id']}] src#{r['source_id']} | {r['title'][:50]:50s} | {r['query_count']}q | {r['created_at'][:16]}")
    db.close()


def cmd_interactive(args):
    import readline
    api_key = args.api_key or os.getenv("GROQ_API_KEY", "")
    if not api_key:
        print("  Error: GROQ_API_KEY not set")
        return

    db = AnalystDB()
    src = db.get_source(args.source)
    if not src:
        print(f"  Source {args.source} not found")
        return

    schema = json.loads(src["schema_json"])
    df, _ = load_data(
        path=Path(src["filename"]) if src["filename"] and Path(src["filename"]).exists() else None,
        source_type=src["source_type"],
    )

    ses = Session(source_id=args.source, title="Interactive session")
    ses_id = db.add_session(ses)
    history = []

    print(f"  Data: {schema['row_count']} rows x {schema['column_count']} columns")
    print(f"  Columns: {', '.join(c['name'] for c in schema['columns'])}")
    print("  Type 'exit' to quit, 'schema' to see schema\n")

    while True:
        try:
            q = input("  >>> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not q:
            continue
        if q.lower() in ("exit", "quit", "q"):
            break
        if q.lower() == "schema":
            for c in schema["columns"]:
                print(f"    {c['name']:25s} {c['dtype']:12s} {c['non_null']} non-null, {c.get('unique','?')} unique")
            continue

        result = analyze(q, df, schema, api_key, history)

        q_obj = Query(session_id=ses_id, question=q, code=result.get("code",""),
                      result_text=result.get("result_text",""), error=result.get("error",""),
                      chart_b64=result.get("chart_b64",""),
                      execution_time_ms=result.get("execution_time_ms",0),
                      tokens_used=result.get("tokens_used",0))
        db.add_query(q_obj)
        history.append({"question": q, "code": result.get("code","")})

        if result.get("error"):
            print(f"  Error: {result['error'][:200]}")
        if result.get("result_text"):
            print(f"\n{result['result_text']}")
        if result.get("chart_b64"):
            print(f"  [Chart: {len(result['chart_b64'])} bytes]")
        print(f"  [{result.get('execution_time_ms',0)}ms]")

    db.close()


def main():
    import argparse
    parser = argparse.ArgumentParser(prog="analyst", description="AI Data Analyst - Chat with your data")
    sub = parser.add_subparsers(dest="command")

    p_load = sub.add_parser("load", help="Load a data source")
    p_load.add_argument("--file", help="File path")
    p_load.add_argument("--type", default="csv", choices=["csv", "excel", "json", "sqlite"])
    p_load.add_argument("--text", help="Raw CSV/JSON text")
    p_load.add_argument("--table", help="Table name for SQLite")

    p_ask = sub.add_parser("ask", help="Ask a question about loaded data")
    p_ask.add_argument("source", type=int, help="Source ID")
    p_ask.add_argument("question", help="Your question")
    p_ask.add_argument("--session", type=int, help="Session ID for context")
    p_ask.add_argument("--continuous", action="store_true", help="Continue last session")
    p_ask.add_argument("--api-key", help="Groq API key")

    p_interactive = sub.add_parser("chat", help="Interactive chat mode")
    p_interactive.add_argument("source", type=int, help="Source ID")
    p_interactive.add_argument("--api-key", help="Groq API key")

    p_list = sub.add_parser("list", help="List sources or sessions")
    p_list.add_argument("type", choices=["sources", "sessions"])

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

    cmds = {
        "load": cmd_load,
        "ask": cmd_ask,
        "chat": cmd_interactive,
        "list": cmd_list,
    }
    cmds[args.command](args)


if __name__ == "__main__":
    main()