from typing import List, Optional

from backend import schemas
from backend.database.base import db
from backend.database.models import Agent
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)


def _load_agent_identity_and_publication_details(project, has_published_version: bool) -> tuple[dict, Optional[dict]]:
    from backend.constants import DRAFT_ZONE, PUBLISHED_ZONE
    from backend.utils.project_utils import get_agent_details

    draft_agent_details = get_agent_details(project, DRAFT_ZONE)
    pub_details = get_agent_details(project, PUBLISHED_ZONE) if has_published_version else None
    return draft_agent_details, pub_details


def _apply_agent_identity_and_publication_details(
    agent_obj: schemas.AgentRead, draft_agent_details: dict, pub_details: Optional[dict]
) -> None:
    agent_obj.llmid = draft_agent_details["llm_id"]
    agent_obj.name = draft_agent_details["name"]

    published_version = agent_obj.published_version
    if published_version and pub_details:
        published_version.llmid = pub_details["llm_id"]
        published_version.name = pub_details["name"]
        agent_obj.published_version = published_version


def _get_agent_identity_and_publication_details(db_agent: Agent) -> Optional[schemas.AgentRead]:
    """Load the minimal DSS-enriched draft and published identity fields for an agent."""
    from backend.utils.project_utils import get_ua_project

    if not db_agent:
        return None
    agent_obj = None
    try:
        agent_obj = schemas.AgentRead.model_validate(db_agent)
        project = get_ua_project(db_agent.id)
        draft_agent_details, pub_details = _load_agent_identity_and_publication_details(
            project, bool(agent_obj.published_version)
        )
        _apply_agent_identity_and_publication_details(agent_obj, draft_agent_details, pub_details)
        return agent_obj
    except Exception as e:
        logger.exception("Error getting identity/publication details for agent_id=%s: %s", db_agent.id, e)
        return agent_obj


def _get_full_agent_details(db_agent: Agent) -> Optional[schemas.AgentRead]:
    from backend.constants import DRAFT_ZONE
    from backend.utils.project_utils import get_cloud_docs_folder_details, get_ua_project

    if not db_agent:
        return None
    agent_obj = None
    try:
        agent_obj = schemas.AgentRead.model_validate(db_agent)
        project = get_ua_project(db_agent.id)
        draft_agent_details, pub_details = _load_agent_identity_and_publication_details(
            project, bool(agent_obj.published_version)
        )
        _apply_agent_identity_and_publication_details(agent_obj, draft_agent_details, pub_details)
        agent_obj.tools = draft_agent_details["tools"]
        agent_obj.kb_description = draft_agent_details["kb_description"]
        agent_obj.system_prompt = draft_agent_details["system_prompt"]
        if agent_obj.documents:
            for doc in agent_obj.documents:
                if isinstance(doc, dict) and doc.get("folderId") and doc.get("connectionType"):
                    doc_details = get_cloud_docs_folder_details(
                        project, DRAFT_ZONE, doc["folderId"], doc["connectionType"]
                    )
                    doc = {
                        **doc,
                        **doc_details,
                    }
                    if agent_obj.cloud_folders is None:
                        agent_obj.cloud_folders = []
                    agent_obj.cloud_folders.append(doc)
        agent_obj.semantic_datasource = draft_agent_details.get("semantic_datasource") or None
        agent_obj.sem_models_llm = draft_agent_details.get("sem_models_llm", "")
        agent_obj.sem_models_interaction_mode = draft_agent_details.get("sem_models_interaction_mode", "simple")
        agent_obj.datasources = db_agent.datasources or (
            [agent_obj.semantic_datasource] if agent_obj.semantic_datasource else []
        )

        logger.debug(f"Agent[{agent_obj.name}] details: {draft_agent_details}")
        published_version = agent_obj.published_version
        if published_version and pub_details:
            published_version.tools = pub_details["tools"]
            published_version.kb_description = pub_details["kb_description"]
            published_version.system_prompt = pub_details["system_prompt"]
            published_version.semantic_datasource = pub_details.get("semantic_datasource") or None
            published_version.sem_models_llm = pub_details.get("sem_models_llm", "")
            published_version.sem_models_interaction_mode = pub_details.get("sem_models_interaction_mode", "simple")
            published_version.datasources = pub_details.get("datasources", []) or (
                [published_version.semantic_datasource] if published_version.semantic_datasource else []
            )
            agent_obj.published_version = published_version
            logger.debug(f"Published version details: {pub_details}")
        return agent_obj
    except Exception as e:
        logger.exception(f"Error getting full agent details for agent_id={db_agent.id}: {e}")
        return agent_obj


def get_agent(agent_id: str) -> Optional[schemas.AgentRead]:
    with db.session.begin():
        agent = db.session.query(Agent).filter(Agent.id == agent_id).first()
        if not agent:
            return None
        return _get_full_agent_details(agent)


def get_agent_identity_and_publication_details(agent_id: str) -> Optional[schemas.AgentRead]:
    with db.session.begin():
        agent = db.session.query(Agent).filter(Agent.id == agent_id).first()
        if not agent:
            return None
        return _get_agent_identity_and_publication_details(agent)


def get_agent_summaries_by_ids(agent_ids: list[str]) -> List[schemas.AgentRead]:
    """Load DB-backed agent rows without the DSS-heavy enrichment performed by full details."""
    if not agent_ids:
        return []

    with db.session.begin():
        agents = db.session.query(Agent).filter(Agent.id.in_(agent_ids)).all()
        return [schemas.AgentRead.model_validate(agent) for agent in agents]


def get_agents_by_ids(agent_ids: list[str]) -> List[schemas.AgentRead]:
    """Load full agent payloads for a targeted subset of ids."""
    if not agent_ids:
        return []

    with db.session.begin():
        agents = db.session.query(Agent).filter(Agent.id.in_(agent_ids)).all()
        return [details for agent in agents if (details := _get_full_agent_details(agent)) is not None]


def get_agents_identity_and_publication_details_by_ids(agent_ids: list[str]) -> List[schemas.AgentRead]:
    """Load the DSS-enriched draft/published identity fields for a targeted subset of ids."""
    if not agent_ids:
        return []

    with db.session.begin():
        agents = db.session.query(Agent).filter(Agent.id.in_(agent_ids)).all()
        return [
            details for agent in agents if (details := _get_agent_identity_and_publication_details(agent)) is not None
        ]


def get_all_agents() -> List[schemas.AgentRead]:
    with db.session.begin():
        agents = db.session.query(Agent).all()
        return [details for agent in agents if (details := _get_full_agent_details(agent)) is not None]


def _get_agents_by_owner_no_ctx_mgr(user_id: str) -> List[schemas.AgentRead]:
    agents = db.session.query(Agent).filter(Agent.owner == user_id).all()
    return [details for agent in agents if (details := _get_full_agent_details(agent)) is not None]


def get_agents_by_owner(user_id: str) -> List[schemas.AgentRead]:
    with db.session.begin():
        return _get_agents_by_owner_no_ctx_mgr(user_id)


def create_agent(agent: schemas.AgentCreate) -> schemas.AgentRead:
    with db.session.begin():
        new_agent = Agent(
            id=agent.id,
            owner=agent.owner,
            description=agent.description,
            sample_questions=agent.sample_questions,
            documents=agent.documents,
            datasources=agent.datasources,
            indexing=agent.indexing,
            published_version=agent.published_version,
            published_at=agent.published_at,
            publishing_status=agent.publishing_status,
            publishing_job_id=agent.publishing_job_id,
            publishing_scenario_runs=agent.publishing_scenario_runs,
        )
        # if get_agent(agent.id):
        #     raise ValueError(f"Agent '{agent.id}' already exists")
        db.session.add(new_agent)

        # flush here replace the commit/refresh pattern to get the new agent's ID populated without committing the transaction yet, the newer fields will be available
        # after flush, but the transaction can still be rolled back.
        # Note : Flush only matter for values that come from the database, such as the generated id. We don't face this issue here but it is a good practice to flush after adding a new object to the session.
        db.session.flush()
        # TODO handle success and failure
        return schemas.AgentRead.model_validate(new_agent)


def update_agent(agent_id: str, agent_update: schemas.AgentUpdate) -> Optional[schemas.AgentRead]:
    with db.session.begin():
        db_agent = db.session.query(Agent).filter(Agent.id == agent_id).first()
        if not db_agent:
            return None
        update_data = agent_update.model_dump(exclude_unset=True)
        for key, value in update_data.items():
            setattr(db_agent, key, value)
        db.session.flush()
        return schemas.AgentRead.model_validate(db_agent)


def delete_agent(agent_id: str):
    with db.session.begin():
        db_agent = db.session.query(Agent).filter(Agent.id == agent_id).first()
        if db_agent:
            db.session.delete(db_agent)
            return True
        return False


def list_user_agent_owners() -> list[dict]:
    """
    Return a list of agent owners and their agent counts.
    """
    with db.session.begin():
        results = (
            db.session.query(Agent.owner, db.func.count(Agent.id)).group_by(Agent.owner).order_by(Agent.owner).all()
        )
        out = []
        for owner, count in results:
            display_name = ""
            email = ""
            try:
                from backend.utils.user_utils import get_user_info

                ui = get_user_info(owner)
                display_name = ui.display_name or ""
                email = ui.email or ""
            except Exception:
                pass
            out.append({"id": owner, "displayName": display_name, "email": email, "agentCount": int(count)})
        return out
