"""
Evaluation Script: Multi-Agent Invoice Classification (Minh - Nam - Thủy)
Comparing Gemini 3.8 Flash vs Gemini 3.7 Flash on 100% Real DSCons Invoices.
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from datetime import datetime, timezone
from decimal import Decimal
from typing import Any

from google.antigravity import Agent, LocalAgentConfig
from app.core.postgres.erp_client import ErpDatabaseClient
from app.core.settings import get_settings


def load_real_invoices() -> list[dict[str, Any]]:
    """Trích xuất danh sách hóa đơn thực tế từ CSDL PostgreSQL của DSCons."""
    db = ErpDatabaseClient()
    target_numbers = ["00000067", "00000054", "00000495", "00000160"]
    invoices = []

    with db.get_connection() as conn:
        with conn.cursor() as cur:
            for num in target_numbers:
                cur.execute(
                    """
                    SELECT id, invoice_number, invoice_series, issue_date, seller_name, seller_tax_code,
                           buyer_name, buyer_tax_code, subtotal_amount_vnd, vat_amount_vnd, total_amount_vnd,
                           vat_rate_percent, direction, notes, tax_risk_level, status
                    FROM erp_invoices
                    WHERE invoice_number = %s
                    LIMIT 1;
                    """,
                    (num,),
                )
                r = cur.fetchone()
                if not r:
                    continue
                inv = dict(r)
                inv["id"] = str(inv["id"])
                inv["issue_date"] = str(inv["issue_date"])
                inv["subtotal_amount_vnd"] = float(inv["subtotal_amount_vnd"] or 0)
                inv["vat_amount_vnd"] = float(inv["vat_amount_vnd"] or 0)
                inv["total_amount_vnd"] = float(inv["total_amount_vnd"] or 0)
                inv["vat_rate_percent"] = float(inv["vat_rate_percent"] or 0)

                cur.execute(
                    """
                    SELECT item_name, quantity, unit, unit_price_vnd, total_item_amount_vnd
                    FROM erp_invoice_items
                    WHERE invoice_id = %s
                    ORDER BY item_order ASC;
                    """,
                    (inv["id"],),
                )
                inv["items"] = [
                    {
                        "item_name": it["item_name"],
                        "quantity": float(it["quantity"] or 0),
                        "unit": it["unit"],
                        "unit_price_vnd": float(it["unit_price_vnd"] or 0),
                        "total_item_amount_vnd": float(it["total_item_amount_vnd"] or 0),
                    }
                    for it in cur.fetchall()
                ]
                invoices.append(inv)
    return invoices


MINH_SYSTEM_PROMPT = """Bạn là Minh - Điều phối Tổng / Pháp lý Luật Xây dựng 2026 của DSCons (MST: 0202111150).
Nhiệm vụ của bạn là tiếp nhận dữ liệu hóa đơn, rà soát tính hợp pháp và phân định vào 4 TRỤ CỘT KINH DOANH ĐỊNH SƠN:
- Trụ Cột 1: Thi công Xây lắp & Thủy lợi (Đa Độ, Cống Bến Kem, Đại Thắng, Rạng Đông, Bảo Long...).
- Trụ Cột 2: Dịch vụ Vận tải Xe Ben & Logistics Cơ Giới (Vận chuyển bùn đất, cát đá bằng đoàn xe Howo/xe ben, hợp đồng Cty Thoát Nước Hải Phòng 0200149705, vận chuyển cho Vinhomes...). TUYỆT ĐỐI KHÔNG gộp thành Dự án Thi công Xây dựng.
- Trụ Cột 3: Cho thuê Máy móc & Ca máy Cơ giới (Cho thuê máy đào bánh xích, máy ủi, cừ larsen...).
- Trụ Cột 4: Mua bán Vật tư & San lấp (Cát, đá thương mại...).

Hãy đưa ra:
1. Xác định Chiều hóa đơn (Đầu vào Input hay Đầu ra Output đối với DSCons MST 0202111150).
2. Phân loại Trụ cột chính xác (1, 2, 3 hay 4) kèm lý giải nghiệp vụ chặt chẽ.
3. Rà soát căn cứ pháp lý (Luật Xây dựng 2025, Nghị định 123/2020/NĐ-CP).
4. Các điểm lưu ý bàn giao cho Kế toán trưởng Nam hạch toán.
"""

NAM_SYSTEM_PROMPT = """Bạn là Nam - Kế toán Trưởng & Quản trị Dòng tiền Dự án của DSCons (MST: 0202111150).
Nhiệm vụ của bạn là tiếp nhận hóa đơn và ý kiến của Minh để thực hiện kiểm soát tài chính kế toán:
- Kiểm tra tính hợp lệ thuế VAT (Nghị định 123/2020/NĐ-CP, thuế suất 8% vs 10%, khớp toán học: Tổng = Trước thuế + VAT).
- Hạch toán định khoản kế toán xây lắp / vận tải:
  + Đầu vào: TK 621 (NVL trực tiếp), TK 623 (Chi phí ca máy), TK 627 (Chi phí chung/vận chuyển), TK 1331 (Thuế VAT được khấu trừ), TK 331 (Phải trả NCC).
  + Đầu ra: TK 5111/5112/5113 (Doanh thu hàng hóa/xây lắp/dịch vụ vận tải), TK 33311 (Thuế VAT đầu ra), TK 131 (Phải thu khách hàng).
- Quy tắc kiểm soát chống khống giá & ĐVT (AEC Định Sơn): cấm áp giá thép tấn cho ĐVT m/m2; cừ Larsen thuê thì VL=0; kiểm soát tỷ trọng MR/T.
- Đánh giá rủi ro thuế (Thấp/Trung bình/Cao) và đề xuất phương án thanh toán/ghi nhận sổ sách.
"""

THUY_SYSTEM_PROMPT = """Bạn là Thủy - Trợ lý Ban Giám Đốc & Thư ký Điều hành của DSCons.
Nhiệm vụ của bạn là tiếp nhận phân tích phối hợp từ Minh (Điều phối/Pháp lý) và Nam (Kế toán trưởng):
- Tổng hợp thành Báo cáo Điều hành Ngắn gọn, Chuẩn xác trình Giám đốc Nguyễn Sĩ Sơn.
- Đưa ra Khuyến nghị Quyết định cuối cùng: [DUYỆT CHI / GHI NHẬN DOANH THU] hoặc [TẠM GIỮ CHỨNG TỪ - BỔ SUNG NHẬT TRÌNH] hoặc [TỪ CHỐI].
- Phân công đầu việc luân chuyển chứng từ và lưu trữ vào hồ sơ dự án.
- Giữ phong cách lịch thiệp, chỉn chu, bảo mật và chuẩn xác số liệu.
"""


async def run_single_agent(model_name: str, system_prompt: str, user_prompt: str, api_key: str) -> dict[str, Any]:
    start = time.time()
    config = LocalAgentConfig(
        model=model_name,
        api_key=api_key,
        system_instruction=system_prompt,
    )
    async with Agent(config) as agent:
        resp = await agent.chat(user_prompt)
        text = await resp.text()
    elapsed = time.time() - start
    word_count = len(text.split())
    approx_tokens = int(word_count * 1.3)
    tps = round(approx_tokens / elapsed, 1) if elapsed > 0 else 0
    return {
        "text": text,
        "elapsed_seconds": round(elapsed, 3),
        "word_count": word_count,
        "approx_tokens": approx_tokens,
        "tps": tps,
    }


async def run_collaborative_workflow(model_name: str, invoice: dict[str, Any], api_key: str) -> dict[str, Any]:
    flow_start = time.time()
    inv_summary = (
        f"HÓA ĐƠN SỐ: {invoice['invoice_number']} (Ký hiệu: {invoice['invoice_series']})\n"
        f"Ngày lập: {invoice['issue_date']}\n"
        f"Bên bán: {invoice['seller_name']} (MST: {invoice['seller_tax_code']})\n"
        f"Bên mua: {invoice['buyer_name']} (MST: {invoice['buyer_tax_code']})\n"
        f"Tiền hàng trước thuế: {invoice['subtotal_amount_vnd']:,.0f} VNĐ\n"
        f"Thuế suất: {invoice['vat_rate_percent']}% | Tiền thuế VAT: {invoice['vat_amount_vnd']:,.0f} VNĐ\n"
        f"Tổng thanh toán: {invoice['total_amount_vnd']:,.0f} VNĐ\n"
        f"Các dòng hàng hóa/dịch vụ:\n"
    )
    for idx, item in enumerate(invoice["items"], 1):
        inv_summary += (
            f"  {idx}. {item['item_name']} | SL: {item['quantity']} {item['unit']} | "
            f"Đơn giá: {item['unit_price_vnd']:,.0f} đ | Thành tiền: {item['total_item_amount_vnd']:,.0f} đ\n"
        )

    # 1. Minh phân loại pháp lý & 4 Trụ Cột
    minh_prompt = f"Phân tích hóa đơn sau và xác định 4 Trụ Cột cùng pháp lý:\n\n{inv_summary}"
    minh_res = await run_single_agent(model_name, MINH_SYSTEM_PROMPT, minh_prompt, api_key)

    # 2. Nam hạch toán chi phí & kiểm soát tài chính
    nam_prompt = (
        f"Dữ liệu hóa đơn:\n{inv_summary}\n\n"
        f"Ý kiến phân loại từ Minh (Điều phối & Pháp lý):\n{minh_res['text']}\n\n"
        f"Hãy kiểm tra thuế, hạch toán định khoản kế toán và kiểm soát đơn vị tính."
    )
    nam_res = await run_single_agent(model_name, NAM_SYSTEM_PROMPT, nam_prompt, api_key)

    # 3. Thủy tổng hợp báo cáo điều hành trình Ban Giám Đốc
    thuy_prompt = (
        f"Dữ liệu hóa đơn:\n{inv_summary}\n\n"
        f"Phân tích pháp lý của Minh:\n{minh_res['text']}\n\n"
        f"Hạch toán kế toán và kiểm soát rủi ro của Nam:\n{nam_res['text']}\n\n"
        f"Hãy tổng hợp Báo cáo Điều hành trình Giám đốc Nguyễn Sĩ Sơn và đưa ra khuyến nghị phê duyệt."
    )
    thuy_res = await run_single_agent(model_name, THUY_SYSTEM_PROMPT, thuy_prompt, api_key)

    total_time = time.time() - flow_start
    total_tokens = minh_res["approx_tokens"] + nam_res["approx_tokens"] + thuy_res["approx_tokens"]
    avg_tps = round(total_tokens / total_time, 1) if total_time > 0 else 0

    return {
        "invoice_number": invoice["invoice_number"],
        "model": model_name,
        "total_duration_seconds": round(total_time, 3),
        "total_tokens": total_tokens,
        "average_tps": avg_tps,
        "steps": {
            "minh": minh_res,
            "nam": nam_res,
            "thuy": thuy_res,
        },
    }


async def main():
    settings = get_settings()
    api_key = settings.gemini_api_key
    if not api_key:
        print("ERROR: GEMINI_API_KEY is not set!")
        sys.exit(1)

    invoices = load_real_invoices()
    print(f"Loaded {len(invoices)} real invoices from PostgreSQL database.")

    results_38 = []
    results_37 = []

    print("\n========================================================")
    print(" BẮT ĐẦU ĐÁNH GIÁ: GEMINI 3.8 FLASH (BẢN MỚI)")
    print("========================================================")
    for inv in invoices:
        print(f"--> Đang chạy hóa đơn {inv['invoice_number']} với Gemini 3.8 Flash...")
        res = await run_collaborative_workflow("gemini-3.8-flash", inv, api_key)
        results_38.append(res)
        print(f"    Hoàn tất: {res['total_duration_seconds']}s | {res['total_tokens']} tokens | {res['average_tps']} TPS")

    print("\n========================================================")
    print(" BẮT ĐẦU ĐÁNH GIÁ: GEMINI 3.7 FLASH (BẢN TIỀN NHIỆM)")
    print("========================================================")
    for inv in invoices:
        print(f"--> Đang chạy hóa đơn {inv['invoice_number']} với Gemini 3.7 Flash...")
        res = await run_collaborative_workflow("gemini-3.7-flash", inv, api_key)
        results_37.append(res)
        print(f"    Hoàn tất: {res['total_duration_seconds']}s | {res['total_tokens']} tokens | {res['average_tps']} TPS")

    output_payload = {
        "evaluated_at": datetime.now(timezone.utc).isoformat(),
        "total_invoices": len(invoices),
        "gemini_3_8_flash": results_38,
        "gemini_3_7_flash": results_37,
    }

    os.makedirs("artifacts", exist_ok=True)
    out_file = "artifacts/invoice_classification_benchmark_38_vs_37.json"
    with open(out_file, "w", encoding="utf-8") as f:
        json.dump(output_payload, f, ensure_ascii=False, indent=2)
    print(f"\nSaved benchmark results to {out_file}")


if __name__ == "__main__":
    asyncio.run(main())
