from __future__ import annotations

from datetime import datetime
from typing import Optional, Tuple

import dataiku
from dataikuapi.dss.managedfolder import DSSManagedFolder
from dataikuapi.dss.project import DSSProject
from dataikuapi.utils import DataikuException

from backend.config import get_config, get_default_embedding_llm
from backend.constants import DRAFT_ZONE, PUBLISHED_ZONE
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import (
    create_aug_llm,
    create_embed_doc_recipe,
    create_kb_tool,
    create_visual_agent,
    get_kb_in_zone,
    get_kb_tool_in_zone,
    get_managed_folder_in_zone,
    get_recipe_in_zone,
    get_ua_project,
    get_visual_agent_in_zone,
    update_visual_agent,
)

logger = get_logger(__name__)


def prepare_agent_tools(project: DSSProject, agent: dict, zone_name: str):
    docs = agent.get("documents", [])
    tools_ids = list(agent.get("tools", []))  # make sure keep it immutable
    kb_id = get_kb_in_zone(project, zone_name)
    knowledge_description = (agent.get("kb_description") or "").strip()
    kb_tool = get_kb_tool_in_zone(project, zone_name)

    if docs:
        if not kb_tool:
            kb_tool = create_kb_tool(
                project=project,
                kb_id=kb_id,
                kb_desc=knowledge_description,
                maxDocuments=get_config().get("number_of_documents_to_retrieve", 4),
                zone_name=zone_name,
            )
            logger.info(
                f"create kb_tool {kb_tool.id} - project {project.project_key} - in zone {zone_name} with description '{knowledge_description}'"
            )
        else:
            # Update description if needed
            tool_settings = kb_tool.get_settings()
            tool_settings_raw = tool_settings.get_raw()
            logger.info(f"Updating kb_tool description, project {project.project_key}, version {zone_name}")
            tool_settings_raw["additionalDescriptionForLLM"] = knowledge_description
            tool_settings.save()
            logger.info(f"Updating kb_tool description, project {project.project_key}, version {zone_name} -- done")
        kb_tool_id = f"{project.project_key}:{kb_tool.id}"
        tools_ids.append(kb_tool_id)
    return tools_ids


def create_quick_agent(project: DSSProject, agent: dict, project_name: str, zone_name: str):
    logger.info(f"Bundling user agent {agent.get('id')} in zone {zone_name}")
    try:
        tools_ids = prepare_agent_tools(project, agent, zone_name)
        llm_id = agent.get("llmid")
        prompt = agent.get("system_prompt", "")
        # Create visual agent in case we have tools
        return create_visual_agent(
            project=project,
            agent_name=f"{project_name} {zone_name}",
            tools_ids=tools_ids,
            llm_id=llm_id,
            prompt=prompt,
        )
    except Exception as e:
        logger.exception(f"Failed to create visual agent in zone {zone_name}: {e}")

def create_agent_project(project_key, project_name, project_folder_id, agent, owner):
    client = dataiku.api_client()
    connection = get_config().get("default_fs_connection")

    project: DSSProject = client.create_project(
        project_key=project_key, owner=owner, name=project_name, project_folder_id=project_folder_id
    )
    llm_id = agent.get("llmid")
    embedding_llm = get_default_embedding_llm()

    logger.info(
        f"Created project {project_key} for user agent, owner: {owner}, llm: {llm_id}, embedding llm: {embedding_llm}"
    )
    try:
        # Prepare flow zones
        for zone_name in [DRAFT_ZONE, PUBLISHED_ZONE]:
            logger.info(f"Creating managed folder with connection {connection} in zone {zone_name}")
            # Prepare folders, recipes & ra models
            folder = project.create_managed_folder(f"documents_{zone_name}", connection_name=connection)
            r = create_embed_doc_recipe(
                project=project, zone=zone_name, folder_id=folder.id, llm_id=llm_id, embedding_llm=embedding_llm
            )
            # Finding the KB ID
            new_kb_id = r.get_settings().data["recipe"]["outputs"]["knowledge_bank"]["items"][0]["ref"]
            ra_model = create_aug_llm(project=project, name=f"ra_model_{zone_name}", kb_id=new_kb_id, llm_id=llm_id)
            # We need to retrieve the DSSSavedModel for the ra_model
            sm = project.get_saved_model(ra_model.id)
            zone = project.get_flow().create_zone(zone_name)
            zone.add_items([folder, r, sm])
            va = create_quick_agent(project, agent, project_name, zone_name)
            if va:
                # TODO uncomment when fixed on DSS side
                zone.add_item(va)
        settings = project.get_settings()
        psettings = settings.get_raw()
        psettings["projectStatus"] = "Archived"
        settings.save()
        return project
    except Exception as e:
        logger.exception("Failed to create all elements in the project {e}")
        logger.info("Deleting project")
        project.delete()
        return None


def get_ua_project_folder(agent: dict) -> Tuple[DSSProject, DSSManagedFolder]:
    project = get_ua_project(agent)
    documents_folder = get_managed_folder_in_zone(project, f"documents_{DRAFT_ZONE}", DRAFT_ZONE)
    if not documents_folder:
        raise Exception("Draft documents folder not found in template")
    return project, documents_folder


def ensure_ua_project(agent: dict, owner) -> DSSProject:
    if "ac_user_agent_" not in agent.get("id"):
        project_key = f"ac_user_agent_{agent['id']}"
        logger.info(f"Creating new project {project_key} owner: {agent.get('owner')}")
        logger.debug(f"ensure_ua_project, Agent details: {agent}")
        agent_hub_folder_id = get_config().get("agents_folder")
        logger.info(f"Using agent hub folder id: {agent_hub_folder_id}")
        return create_agent_project(
            project_key=project_key,
            project_name=agent.get("name", project_key),
            project_folder_id=agent_hub_folder_id,
            agent=agent,
            owner=owner,
        )
    return get_ua_project(agent)


def upload_documents(agent: dict, files) -> None:
    """Upload documents to the draft folder."""
    project, folder = get_ua_project_folder(agent)
    vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
    if not vis_agent:
        # TODO should we throw an exception instead?
        raise Exception("Visual agent not found in draft zone")
        # vis_agent = create_quick_agent(project=project, agent=agent, project_name=project.name, zone_name=DRAFT_ZONE)
    tools_ids = prepare_agent_tools(project, agent, DRAFT_ZONE)
    update_visual_agent(
        agent=vis_agent, tools_ids=tools_ids, llm_id=agent.get("llmid"), prompt=agent.get("system_prompt", "")
    )
    for f in files:
        try:
            f.stream.seek(0)
            folder.put_file(f.filename, f.stream)
            logger.info(f"Uploaded document: {f.filename}")
        except Exception as e:
            logger.error(f"Failed to upload {f.filename}: {e}")
            raise


def remove_document(agent: dict, filename: str) -> None:
    """Remove a document from the draft folder."""
    project, folder = get_ua_project_folder(agent)
    if agent.get("documents") and len(agent.get("documents")) == 1:
        vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
        if not vis_agent:
            raise Exception("Visual agent not found in draft zone")
        logger.info("No more documents, removing kb tool")
        copy_agent = dict(agent)
        copy_agent["documents"] = []
        # only update when no more documents to remove kb tool
        tools_ids = prepare_agent_tools(project, copy_agent, DRAFT_ZONE)
        update_visual_agent(
            agent=vis_agent, tools_ids=tools_ids, llm_id=agent.get("llmid"), prompt=agent.get("system_prompt", "")
        )
    try:
        folder.delete_file(filename)
        logger.info(f"Deleted document: {filename}")
    except Exception as e:
        logger.error(f"Failed to delete document {filename}: {e}")
        # Don't raise - file might already be gone


def delete_agent_project(agent_id: str) -> None:
    """Delete the entire agent project."""
    client = dataiku.api_client()
    try:
        project = client.get_project(agent_id)
        project.delete(clear_managed_datasets=True, clear_output_managed_folders=True, clear_job_and_scenario_logs=True)
        logger.info(f"Deleted project {agent_id}")
    except DataikuException as e:
        logger.warning(f"Project {agent_id} not found or already deleted: {e}")


def launch_index_job(agent: dict, zone: str = DRAFT_ZONE) -> str:
    project = get_ua_project(agent)

    # Find the embed_documents recipe in the specified zone
    recipe = get_recipe_in_zone(project, "embed_documents", zone)
    if not recipe:
        raise Exception(f"No embed_documents recipe found in zone '{zone}'")

    # Start the job
    job = recipe.run(wait=False)
    job_id = job.id

    logger.info(f"Started indexing job {job_id} in zone '{zone}' for agent {agent['id']}")

    return job_id


def get_job_status(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",
        }

        if state == "FAILED":
            activities = base_status.get("activities")
            if not activities or not isinstance(activities, dict) or len(activities) == 0:
                logger.error("No activities found")
            first_key = next(iter(activities))
            first_activity = activities[first_key]
            if first_activity and first_activity.get("firstFailure"):
                error_message = first_activity.get("firstFailure").get("detailedMessage")
                logger.error(f"Job {job_id} failed: {error_message}")

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

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


def sync_agent_to_published(agent: dict) -> None:
    project = get_ua_project(agent)
    vis_agent = get_visual_agent_in_zone(project, PUBLISHED_ZONE)
    if vis_agent:
        logger.info(f"Updating existing Agent {vis_agent.id} in zone {PUBLISHED_ZONE}")
        tools_ids = prepare_agent_tools(project, agent, PUBLISHED_ZONE)
        update_visual_agent(
            agent=vis_agent,
            tools_ids=tools_ids,
            llm_id=agent.get("llmid"),
            prompt=agent.get("system_prompt", ""),
        )
    else:
        logger.info(f"Creating new Agent {agent.get('name', 'Agent')} in zone {PUBLISHED_ZONE}")
        return create_quick_agent(project=project, agent=agent, project_name=project.name, zone_name=PUBLISHED_ZONE)


def sync_folders_to_published(agent: dict) -> None:
    project, draft_folder = get_ua_project_folder(agent)

    published_folder = get_managed_folder_in_zone(project, f"documents_{PUBLISHED_ZONE}", PUBLISHED_ZONE)
    if not published_folder:
        raise Exception("Published documents folder not found")

    logger.info(f"Syncing documents from draft to published for agent {agent['id']}")

    active_docs = {
        d["name"]
        for d in agent.get("documents", [])
        if isinstance(d, dict) and d.get("active") and not d.get("deletePending")
    }

    try:
        for item in published_folder.list_contents()["items"]:
            if item["type"] == "FILE":
                published_folder.delete_file(item["path"])
                logger.info(f"Deleted old file from published: {item['path']}")
    except Exception as e:
        logger.error(f"Error clearing published folder: {e}")

    copied_count = 0
    for doc_name in active_docs:
        try:
            # Read from draft
            response = draft_folder.get_file(doc_name)
            content = response.content
            # Write to published
            published_folder.put_file(doc_name, content)
            logger.info(f"document {doc_name} Copied to published folder")
            copied_count += 1

        except Exception as e:
            logger.error(f"Failed to copy document {doc_name}: {e}")

    logger.info(f"Synced {copied_count} documents to published folder")
