from __future__ import annotations
import asyncio, json, uuid, time
from sqlalchemy import select
from app.models.batch_job import BatchJob
from app.services.storage import async_session_maker
from app.services.connection_pool import get_client

MAX_BATCH_SIZE = 50
MAX_CONCURRENT = 5


async def create_batch(session, requests: list, model: str) -> dict:
    if len(requests) > MAX_BATCH_SIZE:
        requests = requests[:MAX_BATCH_SIZE]

    batch_id = str(uuid.uuid4())
    total = len(requests)
    results = []
    completed = 0
    failed = 0

    session.add(BatchJob(
        id=batch_id, status="processing", total=total,
        completed=0, failed=0, results="[]",
    ))
    await session.commit()

    semaphore = asyncio.Semaphore(MAX_CONCURRENT)
    client = get_client()

    async def _process(req: dict, idx: int):
        nonlocal completed, failed
        async with semaphore:
            try:
                url = req.get("url", "")
                method = req.get("method", "POST").upper()
                headers = req.get("headers", {})
                body = req.get("body", {})

                if method == "GET":
                    resp = await client.get(url, headers=headers, timeout=120.0)
                else:
                    resp = await client.post(url, headers=headers, json=body, timeout=120.0)

                result = {
                    "index": idx, "status_code": resp.status_code,
                    "body": resp.text[:10000],
                    "error": "",
                }
                results.append(result)
                if resp.is_success:
                    completed += 1
                else:
                    failed += 1
            except Exception as e:
                results.append({"index": idx, "status_code": 0, "body": "", "error": str(e)})
                failed += 1

    tasks = [_process(req, i) for i, req in enumerate(requests)]
    await asyncio.gather(*tasks)

    status = "completed" if failed == 0 else ("partial" if completed > 0 else "failed")
    await session.execute(
        select(BatchJob).where(BatchJob.id == batch_id)
    )
    row = (await session.execute(select(BatchJob).where(BatchJob.id == batch_id))).scalar_one()
    row.status = status
    row.completed = completed
    row.failed = failed
    row.results = json.dumps(results)
    row.completed_at = time.time()
    await session.commit()

    return {
        "batch_id": batch_id,
        "status": status,
        "total": total,
        "completed": completed,
        "failed": failed,
    }


async def get_batch_history(session) -> list:
    result = await session.execute(
        select(BatchJob).order_by(BatchJob.created_at.desc()).limit(50)
    )
    rows = result.scalars().all()
    return [
        {
            "batch_id": r.id, "status": r.status,
            "total": r.total, "completed": r.completed, "failed": r.failed,
            "created_at": r.created_at, "completed_at": r.completed_at,
        }
        for r in rows
    ]