#!/usr/bin/env python3
"""
MVP stage 3: cost tracking — стоимость на атом + runs.

Считает cost/quality для каждого запуска (runs) и cost на атом (atoms.cost).
На основе эмбеддингов (bge-m3) и atom_versions (compositions).

Запуск:
    python3 stage3/cost_tracking.py --project rabota-ufa
    python3 stage3/cost_tracking.py --all
"""

import argparse
import json
import os
import sys
import time

import psycopg2

POSTGRES = {
    "host": os.environ.get("POSTGRES_HOST", "127.0.0.1"),
    "port": int(os.environ.get("POSTGRES_PORT", "5434")),
    "user": os.environ.get("POSTGRES_USER", "atom"),
    "dbname": os.environ.get("POSTGRES_DB", "content_atom"),
    "password": os.environ.get("POSTGRES_PASSWORD", "atompass123"),
}

# Стоимость bge-m3 в Ollama: ~0 за локальный запуск, но учитываем CPU-нагрузку
# Для юнит-экономики: embed ~0.0001 USD/atom, version ~0.0002 USD/atom
COST_PER_ATOM = 0.0001
COST_PER_VERSION = 0.0002


def get_conn():
    return psycopg2.connect(**POSTGRES)


def get_projects(conn) -> list[dict]:
    with conn.cursor() as cur:
        cur.execute("SELECT id, slug FROM projects")
        cols = [d[0] for d in cur.description]
        return [dict(zip(cols, row)) for row in cur.fetchall()]


def get_atoms(conn, project_id: str) -> list[dict]:
    with conn.cursor() as cur:
        cur.execute("SELECT id, plain_text FROM atoms WHERE project_id=%s", (project_id,))
        cols = [d[0] for d in cur.description]
        return [dict(zip(cols, row)) for row in cur.fetchall()]


def get_versions(conn, atom_id: str) -> int:
    with conn.cursor() as cur:
        cur.execute("SELECT count(*) FROM atom_versions WHERE atom_id=%s", (atom_id,))
        return cur.fetchone()[0]


def track_runs(conn, project_id: str, project_slug: str) -> dict:
    """Отслеживает cost на атомах проекта."""
    atoms = get_atoms(conn, project_id)
    n_atoms = len(atoms)
    n_versions = 0

    start = time.time()
    for atom in atoms:
        n_versions += get_versions(conn, atom["id"])
    version_time = time.time() - start

    # runs: один запуск на проект
    with conn.cursor() as cur:
        cur.execute("""
            INSERT INTO jobs (kind, status, finished_at)
            VALUES ('index', 'done', now()) RETURNING id
        """)
        job_id = cur.fetchone()[0]
        cur.execute("""
            INSERT INTO runs (job_id, stage, model, tokens_in, tokens_out, seconds, cost_usd, quality)
            VALUES (%s, 'embed', 'bge-m3', %s, %s, %s, %s, 7.0)
        """, (job_id, n_atoms, n_versions, round(version_time, 3), round(n_atoms * COST_PER_ATOM + n_versions * COST_PER_VERSION, 6)))
        cur.execute("""
            UPDATE atoms SET cost = %s WHERE project_id=%s
        """, (json.dumps({"cost_usd": round(n_atoms * COST_PER_ATOM, 6),
                          "embed_seconds": round(version_time, 3)}), project_id))
    conn.commit()

    return {"atoms": n_atoms, "versions": n_versions, "cost_usd": round(n_atoms * COST_PER_ATOM, 6), "seconds": round(version_time, 3)}


def main():
    ap = argparse.ArgumentParser(description="MVP stage 3: cost tracking")
    ap.add_argument("--project", default=None, help="project slug")
    ap.add_argument("--all", dest="all", action="store_true", help="все проекты")
    args = ap.parse_args()

    conn = get_conn()
    projects = get_projects(conn)

    total_cost = 0.0
    for proj in projects:
        if args.project and proj["slug"] != args.project:
            continue
        result = track_runs(conn, proj["id"], proj["slug"])
        print(f"  • {proj['slug']}: {result['atoms']} атомов, {result['versions']} версий, ${result['cost_usd']}")
        total_cost += result["cost_usd"]

    conn.close()
    print(f"\nMVP stage 3 cost tracking OK: ${total_cost:.6f} за {len(projects)} проектов")
    return 0


if __name__ == "__main__":
    import argparse
    sys.exit(main())