import uuid

from app.core.celery_app import celery_app
from app.core.config import get_settings
from app.db.session import SessionLocal
from app.models.dataset import DatasetExample
from app.models.knowledge import ExtractedText
from app.services.ollama_client import expand_instruction_variants, generate_instruction_examples
from app.ws.events import publish_job_event

settings = get_settings()

# A single Ollama call generating more than ~10 examples in one JSON response takes
# long enough (and produces enough output tokens) to risk the request timing out or
# the model truncating the JSON mid-array. Splitting into batches keeps each call
# short and reliable, and commits after every batch so partial progress survives if a
# later batch fails.
GENERATION_BATCH_SIZE = 10


@celery_app.task(name="datasets.generate_examples")
def generate_examples_task(
    dataset_id: str,
    extracted_text_id: str,
    count: int,
    difficulty_mix: list[str],
    model: str | None = None,
) -> None:
    db = SessionLocal()
    try:
        extracted = db.get(ExtractedText, uuid.UUID(extracted_text_id))
        if extracted is None or not extracted.cleaned_text:
            publish_job_event(dataset_id, {"status": "failed", "error": "No cleaned text available"})
            return

        total_batches = -(-count // GENERATION_BATCH_SIZE)  # ceil division
        created = []
        remaining = count
        batch_num = 0

        while remaining > 0:
            batch_num += 1
            batch_count = min(GENERATION_BATCH_SIZE, remaining)
            publish_job_event(
                dataset_id,
                {
                    "status": "generating",
                    "count": count,
                    "batch": batch_num,
                    "total_batches": total_batches,
                    "generated_so_far": len(created),
                },
            )

            raw_examples = generate_instruction_examples(
                extracted.cleaned_text, count=batch_count, difficulty_mix=difficulty_mix, model=model
            )
            for item in raw_examples:
                example = DatasetExample(
                    dataset_id=uuid.UUID(dataset_id),
                    extracted_text_id=extracted.id,
                    example_type=item["example_type"],
                    difficulty=item["difficulty"],
                    instruction=item["instruction"],
                    input=item.get("input"),
                    output=item["output"],
                    generated_by=model or settings.ollama_default_model,
                )
                db.add(example)
                created.append(example)
            db.commit()
            remaining -= batch_count

        publish_job_event(
            dataset_id, {"status": "completed", "generated": len(created)}
        )
    except Exception as exc:  # noqa: BLE001
        db.rollback()
        publish_job_event(dataset_id, {"status": "failed", "error": str(exc)})
        raise
    finally:
        db.close()


@celery_app.task(name="datasets.regenerate_example")
def regenerate_example_task(example_id: str) -> None:
    db = SessionLocal()
    try:
        example = db.get(DatasetExample, uuid.UUID(example_id))
        if example is None or example.extracted_text_id is None:
            return
        extracted = db.get(ExtractedText, example.extracted_text_id)
        if extracted is None or not extracted.cleaned_text:
            return

        publish_job_event(str(example.dataset_id), {"status": "regenerating", "example_id": example_id})

        raw_examples = generate_instruction_examples(
            extracted.cleaned_text, count=1, difficulty_mix=[example.difficulty.value]
        )
        if raw_examples:
            item = raw_examples[0]
            example.example_type = item["example_type"]
            example.difficulty = item["difficulty"]
            example.instruction = item["instruction"]
            example.input = item.get("input")
            example.output = item["output"]
            example.status = "pending"
            db.commit()

        publish_job_event(str(example.dataset_id), {"status": "regenerated", "example_id": example_id})
    except Exception as exc:  # noqa: BLE001
        db.rollback()
        publish_job_event(str(example.dataset_id) if example else "unknown", {"status": "failed", "error": str(exc)})
        raise
    finally:
        db.close()


@celery_app.task(name="datasets.expand_example")
def expand_example_task(example_id: str, variant_types: list[str], model: str | None = None) -> None:
    db = SessionLocal()
    dataset_id = "unknown"
    try:
        example = db.get(DatasetExample, uuid.UUID(example_id))
        if example is None:
            return
        dataset_id = str(example.dataset_id)

        grounding_text = None
        if example.extracted_text_id:
            extracted = db.get(ExtractedText, example.extracted_text_id)
            if extracted is not None and extracted.cleaned_text:
                grounding_text = extracted.cleaned_text

        publish_job_event(
            dataset_id, {"status": "expanding", "example_id": example_id, "count": len(variant_types)}
        )

        variants = expand_instruction_variants(
            instruction=example.instruction,
            output=example.output,
            input_text=example.input,
            variant_types=variant_types,
            grounding_text=grounding_text,
            model=model,
        )

        created = []
        for item in variants:
            variant = DatasetExample(
                dataset_id=example.dataset_id,
                extracted_text_id=example.extracted_text_id,
                parent_example_id=example.id,
                example_type=item["example_type"],
                difficulty=item["difficulty"],
                instruction=item["instruction"],
                input=item.get("input"),
                output=item["output"],
                generated_by=f"synthetic-expansion:{model or settings.ollama_default_model}",
            )
            db.add(variant)
            created.append(variant)
        db.commit()

        publish_job_event(
            dataset_id, {"status": "expanded", "example_id": example_id, "generated": len(created)}
        )
    except Exception as exc:  # noqa: BLE001
        db.rollback()
        publish_job_event(dataset_id, {"status": "failed", "error": str(exc)})
        raise
    finally:
        db.close()