import threading
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Dict, Optional

import dataiku

from backend.config import get_config
from backend.constants import PUBLISHED_ZONE
from backend.database.models import Agent, PublishingStatusEnum
from backend.schemas import schemas
from backend.services import agent_documents
from backend.services.agent_assets import compute_indexing_status, run_cloud_sync_scenarios
from backend.socket import socketio
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import (
    get_ua_project,
)
from backend.utils.utils import _agent_version_rows, get_store

logger = get_logger(__name__)


@dataclass
class PublishingConfig:
    monitor_interval_seconds: int = 5
    job_timeout_minutes: int = 30
    max_retry_attempts: int = 3
    retry_delay_seconds: int = 10
    batch_size: int = 10  # Max agents to check per cycle

    @classmethod
    def from_env(cls) -> "PublishingConfig":
        """Load config from environment/webapp config"""
        cfg = get_config()
        return cls(
            monitor_interval_seconds=cfg.get("publishing_monitor_interval", 5),
            job_timeout_minutes=cfg.get("publishing_job_timeout", 30),
            max_retry_attempts=cfg.get("publishing_max_retries", 3),
            retry_delay_seconds=cfg.get("publishing_retry_delay", 10),
            batch_size=cfg.get("publishing_batch_size", 10),
        )


class PublishingService:
    def __init__(self, app, config: Optional[PublishingConfig] = None):
        self.config = config or PublishingConfig.from_env()
        self._app = app
        self._monitor_thread: Optional[threading.Thread] = None
        self._monitor_event = threading.Event()
        self._shutdown_event = threading.Event()
        self._retry_counts: Dict[str, int] = {}  # agent_id -> retry count
        self._publishing_start_times: Dict[str, datetime] = {}  # agent_id -> start time
        self._lock = threading.Lock()
        self._start_monitor()

    def _start_monitor(self) -> None:
        with self._lock:
            if self._monitor_thread is None or not self._monitor_thread.is_alive():
                self._monitor_thread = threading.Thread(
                    target=self._monitor_loop, daemon=True, name="PublishingMonitor"
                )
                self._monitor_thread.start()
                logger.info("Publishing monitor thread started")

    def wake_monitor(self) -> None:
        self._monitor_event.set()

    def shutdown(self) -> None:
        logger.info("Shutting down publishing service")
        self._shutdown_event.set()
        self._monitor_event.set()  # Wake the monitor
        if self._monitor_thread:
            self._monitor_thread.join(timeout=10)

    def _monitor_loop(self) -> None:
        logger.info("Publishing monitor loop started")

        while not self._shutdown_event.is_set():
            try:
                # Wait indefinitely for wake signal
                logger.info("Publishing monitor waiting for wake signal...")
                self._monitor_event.wait()
                self._monitor_event.clear()

                if self._shutdown_event.is_set():
                    break

                logger.info("Publishing monitor woke up, starting job checks")

                # Poll while there are active jobs
                with self._app.app_context():
                    while not self._shutdown_event.is_set():
                        active_count = self._check_all_publishing_jobs()

                        if active_count == 0:
                            logger.info("No active publishing jobs, going back to sleep")
                            break

                        # Wait for interval or wake signal
                        if self._monitor_event.wait(self.config.monitor_interval_seconds):
                            self._monitor_event.clear()

            except Exception:
                logger.exception("Error in publishing monitor loop")
                # Continue running despite errors

        logger.info("Publishing monitor loop stopped")

    def _check_all_publishing_jobs(self) -> int:
        try:
            store = get_store()
            # Get agents with active publishing jobs
            all_agents = store.get_all_agents(all_users=True)
            publishing_agents = [
                a
                for a in all_agents
                if a.publishing_status == PublishingStatusEnum.PUBLISHING
                and (a.publishing_job_id or a.publishing_scenario_runs)
            ]
            if not publishing_agents:
                return 0

            logger.info(f"Checking {len(publishing_agents)} active publishing jobs")

            # Process in batches to avoid overloading DSS API
            for i in range(0, len(publishing_agents), self.config.batch_size):
                batch = publishing_agents[i : i + self.config.batch_size]
                for agent in batch:
                    if self._shutdown_event.is_set():
                        return 0
                    self._check_single_job(agent)

                # Small delay between batches
                if i + self.config.batch_size < len(publishing_agents):
                    time.sleep(0.5)

            return len(publishing_agents)

        except Exception:
            logger.exception("Error checking publishing jobs")
            return 0

    def _check_single_job(self, agent: Agent) -> None:
        agent_id = agent.id
        job_id = agent.publishing_job_id
        scenario_runs = agent.publishing_scenario_runs or []

        try:
            # Check for timeout
            start_time = self._publishing_start_times.get(agent_id)
            if start_time:
                elapsed = datetime.utcnow() - start_time
                if elapsed > timedelta(minutes=self.config.job_timeout_minutes):
                    logger.warning(f"Publishing for agent {agent_id} timed out after {elapsed}")
                    self._fail_publishing(agent, "Publishing timed out")
                    return

            run_status = compute_indexing_status(agent_id, job_id, scenario_runs)
            if not run_status:
                logger.warning(f"Publishing run for agent {agent_id} not found in DSS")
                self._handle_missing_job(agent)
                return

            current_status = run_status["status"]
            logger.info(f"Publishing run for agent {agent_id} status: {current_status}")
            # Handle based on status
            if current_status in ("pending", "running"):
                # Still in progress
                return
            elif current_status == "success":
                # Trigger the embedding job and complete
                self._complete_publishing(agent)
            else:
                # Job failed
                self._handle_failed_job(agent, f"DSS job failed: {current_status}")

        except Exception as e:
            logger.exception(f"Error checking job {job_id} for agent {agent_id}")
            self._handle_failed_job(agent, str(e))

    def _get_dss_job_status(self, job_id: str, agent_id: str) -> Optional[Dict]:
        try:
            client = dataiku.api_client()
            project = client.get_project(agent_id)
            job = project.get_job(job_id)

            base_status = job.get_status().get("baseStatus", {})
            state = base_status.get("state", "")

            # Map DSS states to our states
            status_map = {
                "NOT_STARTED": "pending",
                "RUNNING": "running",
                "DONE": "success",
                "FAILED": "failure",
                "ABORTED": "failure",
            }

            return {
                "status": status_map.get(state, "pending"),
                "raw_state": state,
                "updated_at": datetime.utcnow().isoformat(),
            }

        except Exception as e:
            logger.error(f"Failed to get DSS job status for {job_id}: {e}")
            return None

    def _handle_missing_job(self, agent: Agent) -> None:
        """Handle case where DSS job is not found."""
        agent_id = agent.id
        retry_count = self._retry_counts.get(agent_id, 0)

        if retry_count < self.config.max_retry_attempts:
            logger.info(f"Retrying publishing for agent {agent_id} (attempt {retry_count + 1})")
            self._retry_counts[agent_id] = retry_count + 1

            # Clear job ID and reset status
            update = schemas.AgentUpdate(
                publishing_status=PublishingStatusEnum.IDLE,
                publishing_job_id=None,
                publishing_scenario_runs=None,
            )
            get_store().update_agent(agent_id, update, bypass_user=True)

            # Wait before retry
            time.sleep(self.config.retry_delay_seconds)

            # Restart publishing
            try:
                self.start_publish(agent_id, agent.owner)
            except Exception:
                logger.exception(f"Failed to retry publishing for agent {agent_id}")
                self._fail_publishing(agent, "Failed to retry publishing")
        else:
            self._fail_publishing(agent, "Publishing job lost - max retries exceeded")

    def _handle_failed_job(self, agent: Agent, error: str) -> None:
        """Handle a failed publishing job with retry logic."""
        agent_id = agent.id
        retry_count = self._retry_counts.get(agent_id, 0)

        if retry_count < self.config.max_retry_attempts:
            logger.warning(f"Publishing failed for {agent_id}, will retry: {error}")
            self._retry_counts[agent_id] = retry_count + 1

            # Reset status for retry
            update = schemas.AgentUpdate(
                publishing_status=PublishingStatusEnum.IDLE,
                publishing_job_id=None,
                publishing_scenario_runs=None,
            )
            get_store().update_agent(agent_id, update, bypass_user=True)

            # Schedule retry
            threading.Timer(self.config.retry_delay_seconds, lambda: self.start_publish(agent_id, agent.owner)).start()
        else:
            # Max retries exceeded
            self._fail_publishing(agent, error)

    def start_publish(self, agent_id: str, user_id: str) -> dict:
        logger.info(f"Starting publish for agent {agent_id} by user {user_id}")

        agent = get_store().get_agent(agent_id)
        if not agent:
            raise ValueError(f"Agent {agent_id} not found")

        if agent.owner != user_id:
            raise PermissionError("Only the owner can publish an agent")

        # Check if already publishing
        if agent.publishing_status == PublishingStatusEnum.PUBLISHING and (
            agent.publishing_job_id or agent.publishing_scenario_runs
        ):
            job_id = agent.publishing_job_id
            scenario_runs = agent.publishing_scenario_runs or []
            logger.info(f"Agent {agent_id} already publishing with job {job_id} and scenarios {scenario_runs}")
            return {
                "status": "already_publishing",
                "agent_id": agent_id,
                "job_id": job_id,
                "scenario_runs": scenario_runs,
            }

        store = get_store()
        try:
            # Update agent status
            store.update_agent(agent_id, {"publishing_status": PublishingStatusEnum.PUBLISHING})
            # Start the sync job
            job_id, scenario_runs = self._start_sync_job(agent)
            if not job_id and not scenario_runs:
                self._complete_publishing(agent)
                return {"status": PublishingStatusEnum.PUBLISHED.value, "agent_id": agent_id}
            # Track publishing start time
            self._publishing_start_times[agent_id] = datetime.utcnow()

            # Update agent status
            store.update_agent(
                agent_id,
                {
                    "publishing_job_id": job_id,
                    "publishing_scenario_runs": scenario_runs or None,
                },
            )

            # Clear retry count
            self._retry_counts.pop(agent_id, None)

            # Emit start event
            self._emit_publish_event(
                agent_id,
                user_id,
                "publish_started",
                {"job_id": job_id, "scenario_runs": scenario_runs or []},
            )

            # Wake the monitor
            self.wake_monitor()

            logger.info(f"Started publishing for agent {agent_id} with job {job_id} and scenarios {scenario_runs}")

            return {"status": "started", "agent_id": agent_id, "job_id": job_id, "scenario_runs": scenario_runs or []}

        except Exception:
            logger.exception(f"Failed to start publishing for agent {agent_id}")
            store.update_agent(
                agent_id,
                {
                    "publishing_status": PublishingStatusEnum.FAILED,
                    "publishing_job_id": None,
                    "publishing_scenario_runs": None,
                },
                bypass_user=True,
            )
            raise

    def _start_sync_job(self, agent: Agent) -> tuple[Optional[str], list[dict]]:
        logger.info(f"Starting publish sync for agent {agent.id}")

        try:
            from backend.services.agent_assets import sync_agent_to_published, sync_folders_to_published

            agent_dict = schemas.AgentRead.model_validate(agent).model_dump()
            sync_folders_to_published(agent_dict)
            sync_agent_to_published(agent_dict)
            project = get_ua_project(agent.id)

            docs = agent.documents or []
            local_docs = [
                doc
                for doc in docs
                if isinstance(doc, str)
                or (
                    isinstance(doc, dict)
                    and doc.get("source") != "connection_folder"
                    and not doc.get("connectionType")
                    and not doc.get("deletePending")
                )
            ]
            cloud_docs = [
                doc
                for doc in docs
                if isinstance(doc, dict)
                and doc.get("source") == "connection_folder"
                and doc.get("connectionType")
                and not doc.get("deletePending")
            ]

            scenario_runs = []
            if cloud_docs:
                scenario_runs = run_cloud_sync_scenarios(project, cloud_docs, PUBLISHED_ZONE)
                if not scenario_runs:
                    raise RuntimeError("Failed to start cloud sync scenario")

            indexing_blob = agent_documents.build_indexing_blob(
                agent.id,
                zone=PUBLISHED_ZONE,
                include_local=bool(local_docs) and not bool(scenario_runs),
                include_cloud=False,
                start_job=not bool(scenario_runs),
            )

            job_id = indexing_blob.get("jobId")
            if job_id:
                logger.info(f"Started published indexing job {job_id} for agent {agent.id}")
            elif scenario_runs:
                logger.info(f"Started published cloud sync scenarios for agent {agent.id}: {scenario_runs}")
            else:
                logger.info(f"Refreshed published merge outputs for agent {agent.id} without starting an index job")
            return job_id, scenario_runs
        except Exception as e:
            logger.error(f"Failed to start publish sync: {e}")
            raise

    def _complete_publishing(self, agent: Agent) -> None:
        agent_id = agent.id
        logger.info(f"Completing publishing for agent {agent_id}")

        try:
            # Update agent name in published zone
            from backend.utils.project_utils import get_ua_project, update_agent_name_in_zone

            store = get_store()
            agent_dict = schemas.AgentRead.model_validate(agent).model_dump()
            project = get_ua_project(agent_dict.get("id"))
            update_agent_name_in_zone(project, PUBLISHED_ZONE, agent.name)
            # Create published version snapshot
            published_version = self._create_published_snapshot(agent)
            # Update agent with published state
            now = datetime.utcnow().isoformat()
            updates = {
                "published_version": published_version,
                "published_at": now,
                "publishing_status": PublishingStatusEnum.PUBLISHED,
                "publishing_job_id": None,  # Clear job ID
                "publishing_scenario_runs": None,
            }
            store.update_agent(agent_id, updates, bypass_user=True)

            # Clean up tracking
            self._publishing_start_times.pop(agent_id, None)
            self._retry_counts.pop(agent_id, None)

            # Emit success event
            self._emit_publish_event(
                agent_id, agent.owner, "publish_completed", {"published_at": now, "agent_name": agent.name}
            )

            # Notify shared users
            self._notify_shared_users(agent_id)

            logger.info(f"Successfully published agent {agent_id}")

        except Exception as e:
            logger.exception(f"Error completing publishing for agent {agent_id}")
            self._fail_publishing(agent, str(e))

    def _create_published_snapshot(self, agent: Agent) -> dict:
        """Create a snapshot of the agent configuration for the published version."""
        from backend.config import get_ui_config

        documents = agent.documents or []
        # Filter to only active, non-deleted documents
        active_docs = [d for d in documents if isinstance(d, dict) and d.get("active") and not d.get("deletePending")]

        # TODO SQLAlchemy/Pydantic: improve and use pydantic models
        return {
            # "name": agent.name,
            "description": agent.description,
            # "system_prompt": agent.system_prompt,
            # "kb_description": agent.kb_description,
            "sample_questions": agent.sample_questions,
            # "llmid": agent.llmid,
            # "tools": agent.tools,
            "documents": active_docs,
            "indexing": {"status": "success"},
            "datasources": agent.datasources,
            "sem_models_llm": get_ui_config()["sem_models_llm"] if agent.datasources else "",
        }

    def _fail_publishing(self, agent: Agent, error: str) -> None:
        agent_id = agent.id
        logger.error(f"Publishing failed for agent {agent_id}: {error}")

        # Update status
        store = get_store()
        update = {
            "publishing_status": PublishingStatusEnum.FAILED,
            "publishing_job_id": None,
            "publishing_scenario_runs": None,
        }
        store.update_agent(agent_id, update, bypass_user=True)

        # Clean up tracking
        self._publishing_start_times.pop(agent_id, None)
        self._retry_counts.pop(agent_id, None)

        # Emit failure event
        self._emit_publish_event(agent_id, agent.owner, "publish_failed", {"error": error})

    def _emit_publish_event(self, agent_id: str, user_id: str, event_type: str, data: Optional[dict] = None) -> None:
        """Emit publishing events via WebSocket."""
        event_data = {"agent_id": agent_id, "event_type": event_type, "timestamp": datetime.utcnow().isoformat()}
        if data:
            event_data.update(data)

        try:
            socketio.emit("agent:publish_event", event_data, room=f"user:{user_id}")
        except Exception:
            logger.exception(f"Failed to emit publish event {event_type}")

    def _notify_shared_users(self, agent_id: str) -> None:
        try:
            store = get_store()
            agent = store.get_agent(agent_id)
            if not agent:
                return

            # Fetch all share records for this agent
            shares = store.get_agent_shares(agent_id)
            share_count = len([s for s in shares if s.principal_type == schemas.PrincipalTypeEnum.USER])

            agent_dict = schemas.AgentRead.model_validate(agent).model_dump(mode="json")
            version_rows = _agent_version_rows(agent_dict, share_count, agent.owner)

            # owner
            for row in version_rows:
                socketio.emit("agent:updated", row, room=f"user:{agent.owner}")

            # shared users
            for share in shares:
                if share.principal_type == schemas.PrincipalTypeEnum.USER:
                    for row in version_rows:
                        socketio.emit(
                            "agent:updated",
                            {**row, "isShared": True},
                            room=f"user:{share.principal}",
                        )
        except Exception:
            logger.exception(f"Failed to emit agent:updated for {agent_id}")

    def resume_publishing_jobs(self) -> None:
        logger.info("Resuming publishing job monitoring")
        with self._app.app_context():
            try:
                store = get_store()
                # Get agents with active publishing jobs
                all_agents = store.get_all_agents(all_users=True)
                publishing_agents = [
                    a
                    for a in all_agents
                    if a.publishing_status == PublishingStatusEnum.PUBLISHING
                    and (a.publishing_job_id or a.publishing_scenario_runs)
                ]
                # Restore publishing start times (assume they started recently)
                for agent in publishing_agents:
                    # Assume publishing started 5 minutes ago if we don't know
                    self._publishing_start_times[agent.id] = datetime.utcnow() - timedelta(minutes=5)

                if publishing_agents:
                    logger.info(f"Found {len(publishing_agents)} agents with active publishing jobs")
                    self.wake_monitor()
                else:
                    logger.info("No active publishing jobs found")

            except Exception:
                logger.exception("Error resuming publishing jobs")

    def get_stats(self) -> dict:
        return {
            "monitor_running": self._monitor_thread and self._monitor_thread.is_alive(),
            "active_jobs": len(self._publishing_start_times),
            "retry_queue": len(self._retry_counts),
            "config": {
                "monitor_interval": self.config.monitor_interval_seconds,
                "job_timeout_minutes": self.config.job_timeout_minutes,
                "max_retries": self.config.max_retry_attempts,
            },
        }
