from __future__ import annotations

import asyncio
import inspect
import json
import queue
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any

from backend import config
from backend.logging_audit import audit_event
from backend.run_report import write_upload_reports


class JobStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class PipelineJob:
    job_id: str
    user_n: str
    ip: str
    status: JobStatus = JobStatus.PENDING
    stage: str = "pending"
    created_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc)
        .isoformat()
        .replace("+00:00", "Z")
    )
    updated_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc)
        .isoformat()
        .replace("+00:00", "Z")
    )
    error: str | None = None
    uploaded_files: list[str] = field(default_factory=list)
    artifacts: dict[str, str] = field(default_factory=dict)
    stage_counts: dict[str, Any] = field(default_factory=dict)
    _thread: threading.Thread | None = None
    _queue: queue.Queue[str | None] = field(default_factory=queue.Queue)

    @property
    def root_dir(self) -> Path:
        return config.job_dir(self.job_id)

    @property
    def uploads_dir(self) -> Path:
        return config.job_uploads_dir(self.job_id)

    @property
    def ingestion_dir(self) -> Path:
        return config.job_ingestion_dir(self.job_id)

    @property
    def chunking_dir(self) -> Path:
        return config.job_chunking_dir(self.job_id)

    @property
    def embedding_dir(self) -> Path:
        return config.job_embedding_dir(self.job_id)


_JOBS: dict[str, PipelineJob] = {}
_JOBS_LOCK = threading.Lock()


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _sse(event_type: str, payload: dict[str, Any]) -> str:
    return f"event: {event_type}\ndata: {json.dumps(payload, ensure_ascii=True)}\n\n"


def _emit(job: PipelineJob, event_type: str, payload: dict[str, Any]) -> None:
    body = {
        "job_id": job.job_id,
        "stage": job.stage,
        "status": job.status.value,
        **payload,
    }
    job._queue.put(_sse(event_type, body))


def _ensure_stage_supports_output_dir(callable_obj: Any, stage_name: str) -> None:
    params = inspect.signature(callable_obj).parameters
    if "output_dir" not in params:
        raise RuntimeError(
            f"{stage_name} stage cannot be constrained to backend output path using existing API."
        )


def create_job(user_n: str, ip: str) -> PipelineJob:
    config.bootstrap()
    job_id = uuid.uuid4().hex[:12]
    job = PipelineJob(job_id=job_id, user_n=user_n, ip=ip or "unknown")
    for path in (
        job.root_dir,
        job.uploads_dir,
        job.ingestion_dir,
        job.chunking_dir,
        job.embedding_dir,
    ):
        path.mkdir(parents=True, exist_ok=True)
    with _JOBS_LOCK:
        _JOBS[job_id] = job
    _emit(job, "job_created", {"message": "job accepted"})
    return job


def get_job(job_id: str) -> PipelineJob:
    with _JOBS_LOCK:
        job = _JOBS.get(job_id)
    if job is None:
        raise KeyError(f"Job not found: {job_id}")
    return job


def record_uploads(job_id: str, filenames: list[str]) -> None:
    job = get_job(job_id)
    job.uploaded_files = filenames
    job.updated_at = _utc_now()
    _emit(job, "upload_received", {"files": filenames, "file_count": len(filenames)})
    audit_event(
        "kb_upload",
        {"files": filenames, "file_count": len(filenames)},
        user_n=job.user_n,
        ip=job.ip,
        job_id=job.job_id,
    )


def start_job(job_id: str) -> None:
    job = get_job(job_id)
    if job._thread and job._thread.is_alive():
        return
    thread = threading.Thread(target=_run_job, args=(job,), daemon=True)
    job._thread = thread
    thread.start()


def _stage_start(job: PipelineJob, stage: str) -> None:
    job.stage = stage
    job.status = JobStatus.RUNNING
    job.updated_at = _utc_now()
    _emit(job, "stage_started", {"stage_name": stage})
    audit_event(
        "stage_started",
        {"stage_name": stage},
        user_n=job.user_n,
        ip=job.ip,
        job_id=job.job_id,
    )


def _stage_done(job: PipelineJob, stage: str, payload: dict[str, Any]) -> None:
    job.updated_at = _utc_now()
    _emit(job, "stage_completed", {"stage_name": stage, **payload})
    audit_event(
        "stage_completed",
        {"stage_name": stage, **payload},
        user_n=job.user_n,
        ip=job.ip,
        job_id=job.job_id,
    )


def _run_job(job: PipelineJob) -> None:
    try:
        modules = config.get_stage_modules()
        ingestion_engine = modules["ingestion_engine"]
        ingestion_events = modules["ingestion_events"]
        chunking = modules["chunking"]
        embedding_chunks = modules["embedding_chunks"]
        vectordb_upsert = modules["vectordb_upsert"]

        if not any(path.is_file() for path in job.uploads_dir.rglob("*")):
            raise RuntimeError("No uploaded files were provided for ingestion.")

        # Ingestion
        _stage_start(job, "ingestion")
        callback_emitter = ingestion_events.CallbackEmitter(
            lambda event: _on_ingestion_event(job, ingestion_events, event)
        )
        ingestion_cfg = ingestion_engine.IngestionConfig(
            input_dir=job.uploads_dir,
            output_dir=job.ingestion_dir,
            workers=config.env_int("INGEST_WORKERS", 4),
        )
        engine = ingestion_engine.IngestionEngine(ingestion_cfg, emitter=callback_emitter)
        ingestion_stats = engine.run()
        cleaned_path = job.ingestion_dir / "cleaned_documents.jsonl"
        if not cleaned_path.exists():
            raise RuntimeError("Ingestion did not produce cleaned_documents.jsonl.")
        job.artifacts["cleaned_documents_jsonl"] = str(cleaned_path)
        job.stage_counts["ingestion"] = ingestion_stats.get("counts", {})
        _stage_done(job, "ingestion", {"counts": ingestion_stats.get("counts", {})})

        # Chunking
        _stage_start(job, "chunking")
        _ensure_stage_supports_output_dir(chunking.chunk_file, "chunking")
        chunk_stats = chunking.chunk_file(
            input_path=cleaned_path,
            output_dir=job.chunking_dir,
            chunk_size=config.env_int("CHUNK_SIZE", 2000),
            chunk_overlap=config.env_int("CHUNK_OVERLAP", 200),
        )
        chunked_path = job.chunking_dir / "chunked_documents.jsonl"
        if not chunked_path.exists():
            raise RuntimeError("Chunking did not produce chunked_documents.jsonl.")
        job.artifacts["chunked_documents_jsonl"] = str(chunked_path)
        job.stage_counts["chunking"] = chunk_stats.get("counts", {})
        _stage_done(job, "chunking", {"counts": chunk_stats.get("counts", {})})

        # Embedding
        _stage_start(job, "embedding")
        _ensure_stage_supports_output_dir(embedding_chunks.embed_chunks, "embedding")
        embed_stats = embedding_chunks.embed_chunks(
            input_path=str(chunked_path),
            output_dir=str(job.embedding_dir),
            model_name=config.env_value(
                "EMBEDDING_MODEL", "nvidia/llama-3.2-nv-embedqa-1b-v2"
            ),
            api_key=config.env_value("NVIDIA_API_KEY", ""),
            truncate="NONE",
            batch_size=config.env_int("EMBED_BATCH_SIZE", 32),
            normalize=False,
        )
        embedding_path = job.embedding_dir / "embeddings_mtsamples.jsonl"
        if not embedding_path.exists():
            raise RuntimeError("Embedding stage did not produce embeddings_mtsamples.jsonl.")
        job.artifacts["embeddings_jsonl"] = str(embedding_path)
        job.stage_counts["embedding"] = embed_stats.get("counts", {})
        _stage_done(job, "embedding", {"counts": embed_stats.get("counts", {})})

        # Upload reports (per file stage artifacts copied under backend/run_report/upload)
        _stage_start(job, "upload_report")
        report_summary = write_upload_reports(
            job_id=job.job_id,
            uploaded_files=job.uploaded_files,
            ingestion_jsonl=cleaned_path,
            chunking_jsonl=chunked_path,
            embedding_jsonl=embedding_path,
        )
        job.artifacts["run_report_upload"] = str(config.RUN_REPORT_UPLOAD_DIR)
        _stage_done(job, "upload_report", {"summary": report_summary})

        # Dense upsert
        _stage_start(job, "upsert_dense")
        dense_result = vectordb_upsert.upsert_dense(
            embedding_path,
            namespace=config.env_value("DOC_NAMESPACE", ""),
        )
        job.stage_counts["upsert_dense"] = {
            "stored_count": dense_result.get("stored_count", 0)
        }
        _stage_done(job, "upsert_dense", {"result": dense_result})

        # Sparse upsert
        _stage_start(job, "upsert_sparse")
        sparse_result = vectordb_upsert.upsert_sparse(
            embedding_path,
            namespace=config.env_value("DOC_NAMESPACE", ""),
        )
        job.stage_counts["upsert_sparse"] = {
            "stored_count": sparse_result.get("stored_count", 0)
        }
        _stage_done(job, "upsert_sparse", {"result": sparse_result})

        job.status = JobStatus.COMPLETED
        job.stage = "completed"
        job.updated_at = _utc_now()
        _emit(job, "pipeline_completed", {"artifacts": job.artifacts, "counts": job.stage_counts})
        audit_event(
            "pipeline_completed",
            {"artifacts": job.artifacts, "counts": job.stage_counts},
            user_n=job.user_n,
            ip=job.ip,
            job_id=job.job_id,
        )
    except Exception as exc:
        job.status = JobStatus.FAILED
        job.error = str(exc)
        job.updated_at = _utc_now()
        _emit(job, "pipeline_failed", {"error": str(exc)})
        audit_event(
            "pipeline_failed",
            {"error": str(exc)},
            user_n=job.user_n,
            ip=job.ip,
            job_id=job.job_id,
        )
    finally:
        job._queue.put(None)


def _on_ingestion_event(job: PipelineJob, ingestion_events_module: Any, event: Any) -> None:
    payload = ingestion_events_module.event_to_dict(event)
    event_type = payload.get("event_type", "ingestion_event")
    _emit(job, event_type, payload)


async def stream_events(job_id: str):
    job = get_job(job_id)
    while True:
        try:
            value = job._queue.get_nowait()
        except queue.Empty:
            if job.status in {JobStatus.COMPLETED, JobStatus.FAILED}:
                return
            await asyncio.sleep(0.05)
            continue
        if value is None:
            return
        yield value

def job_status(job_id: str) -> dict[str, Any]:
    job = get_job(job_id)
    return {
        "job_id": job.job_id,
        "status": job.status.value,
        "stage": job.stage,
        "created_at": job.created_at,
        "updated_at": job.updated_at,
        "uploaded_files": job.uploaded_files,
        "artifacts": job.artifacts,
        "counts": job.stage_counts,
        "error": job.error,
    }