
import os
import uuid
from typing import Any

from logos_adk.artifacts import ArtifactRecord, ArtifactSandbox
from logos_adk.context import get_run_context
from logos_adk.diffusion import build_diffusion_backend
from logos_adk.tools import Tool, tool
from logos_adk.tools.toolkit import Toolkit

# Global artifacts directory
ARTIFACTS_DIR = ".logos_adk/artifacts"

_IMAGE_EXTENSION_BY_MIME_TYPE = {
    "image/jpeg": ".jpg",
    "image/png": ".png",
    "image/svg+xml": ".svg",
    "image/webp": ".webp",
}

def _ensure_dir():
    if not os.path.exists(ARTIFACTS_DIR):
        os.makedirs(ARTIFACTS_DIR, exist_ok=True)

@tool
async def save_artifact(filename: str, content: Any, mime_type: str = "text/plain"):
    """
    Сохранить файл (артефакт) в рамках текущей сессии. 
    Используй это для создания отчетов, таблиц, текстов или изображений.
    Возвращает ссылку на сохраненный файл.
    """
    _ensure_dir()
    run_ctx = get_run_context()
    session_id = run_ctx.session_id if run_ctx else "global"
    
    # Generate unique ID for storage
    artifact_id = str(uuid.uuid4())
    ext = os.path.splitext(filename)[1]
    storage_name = f"{artifact_id}{ext}"
    storage_path = os.path.join(ARTIFACTS_DIR, storage_name)
    
    # Save content
    mode = "wb" if isinstance(content, bytes) else "w"
    with open(storage_path, mode) as f:
        f.write(content)
        
    # Store metadata in PipelineContext if available
    if run_ctx and hasattr(run_ctx, 'pipeline_ctx'):
        ctx = run_ctx.pipeline_ctx
        artifacts = ctx.metadata.setdefault("_artifacts", [])
        artifact_entry = {
            "id": artifact_id,
            "filename": filename,
            "storage_path": storage_path,
            "mime_type": mime_type,
            "session_id": session_id,
            "created_at": os.path.getctime(storage_path)
        }
        artifacts.append(artifact_entry)
        
    return f"Артефакт '{filename}' сохранен успешно (ID: {artifact_id}). Он появится в вашей галерее."


class ArtifactToolkit(Toolkit):
    """Shared workspace-scoped artifact tools."""

    def __init__(self, *, root: str | None = None) -> None:
        self._sandbox = ArtifactSandbox(root=root)

    def get_tools(self) -> list[Tool]:
        sandbox = self._sandbox

        async def create_text_artifact(
            *,
            name: str,
            content: str,
            mime_type: str = "text/plain; charset=utf-8",
            extension: str = ".txt",
        ) -> dict[str, Any]:
            """Persist a text artifact in the current workspace."""
            record = sandbox.create_text_artifact(
                workspace_id=None,
                name=name,
                content=content,
                mime_type=mime_type,
                extension=extension,
                created_by=_current_actor_name(),
            )
            return _record_to_payload(record)

        async def create_json_artifact(
            *,
            name: str,
            payload: Any,
        ) -> dict[str, Any]:
            """Persist structured JSON in the current workspace."""
            record = sandbox.create_json_artifact(
                workspace_id=None,
                name=name,
                payload=payload,
                created_by=_current_actor_name(),
            )
            return _record_to_payload(record)

        async def create_excel_artifact(
            *,
            name: str,
            rows: list[dict[str, Any]],
            columns: list[str] | None = None,
            sheet_name: str = "Sheet1",
        ) -> dict[str, Any]:
            """Persist tabular data as an Excel workbook."""
            record = sandbox.create_excel_artifact(
                workspace_id=None,
                name=name,
                rows=rows,
                columns=columns,
                sheet_name=sheet_name,
                created_by=_current_actor_name(),
            )
            return _record_to_payload(record)

        async def list_artifacts(
            *,
            kind: str | None = None,
            limit: int = 100,
        ) -> list[dict[str, Any]]:
            """List artifacts visible in the current workspace."""
            records = sandbox.list_artifacts(workspace_id=None, kind=kind, limit=limit)
            return [_record_to_payload(record) for record in records]

        async def read_artifact_text(
            *,
            artifact_id: str,
            encoding: str = "utf-8",
        ) -> str:
            """Read a text artifact from the current workspace."""
            _ensure_workspace_access(sandbox, artifact_id)
            return sandbox.read_text(artifact_id, encoding=encoding)

        async def generate_diffusion_image_artifact(
            *,
            name: str,
            prompt: str,
            negative_prompt: str | None = None,
            width: int = 1024,
            height: int = 1024,
            steps: int = 28,
            guidance_scale: float = 7.0,
            seed: int | None = None,
            model: str | None = None,
            backend: str | None = None,
            base_url: str | None = None,
            api_key: str | None = None,
            sampler_name: str | None = None,
        ) -> dict[str, Any]:
            """Generate an image via the configured diffusion backend and persist it as an artifact."""
            image_backend = build_diffusion_backend(
                backend=backend,
                base_url=base_url,
                api_key=api_key,
            )
            result = await image_backend.generate_image(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                steps=steps,
                guidance_scale=guidance_scale,
                seed=seed,
                model=model,
                sampler_name=sampler_name,
            )
            metadata = dict(result.metadata)
            metadata.setdefault("prompt", prompt)
            metadata.setdefault("negative_prompt", negative_prompt or "")
            metadata.setdefault("width", width)
            metadata.setdefault("height", height)
            metadata.setdefault("steps", steps)
            metadata.setdefault("guidance_scale", guidance_scale)
            metadata.setdefault("seed", seed)
            if model is not None:
                metadata.setdefault("model", model)
            record = sandbox.create_base64_artifact(
                workspace_id=None,
                name=name,
                base64_data=result.base64_data,
                mime_type=result.mime_type,
                kind="image",
                created_by=_current_actor_name(),
                metadata=metadata,
                extension=_image_extension(result.mime_type),
            )
            return _record_to_payload(record)

        return [
            Tool(
                fn=create_text_artifact,
                name="create_text_artifact",
                description="Create a text artifact in the current workspace.",
            ),
            Tool(
                fn=create_json_artifact,
                name="create_json_artifact",
                description="Create a JSON artifact in the current workspace.",
            ),
            Tool(
                fn=create_excel_artifact,
                name="create_excel_artifact",
                description="Create an Excel artifact in the current workspace.",
            ),
            Tool(
                fn=list_artifacts,
                name="list_artifacts",
                description="List artifacts available in the current workspace.",
            ),
            Tool(
                fn=read_artifact_text,
                name="read_artifact_text",
                description="Read a text artifact from the current workspace.",
            ),
            Tool(
                fn=generate_diffusion_image_artifact,
                name="generate_diffusion_image_artifact",
                description="Generate and persist an image artifact with the configured diffusion backend.",
            ),
        ]


def _current_actor_name() -> str | None:
    run_ctx = get_run_context()
    if run_ctx is None:
        return None
    agent_name = getattr(run_ctx, "agent_name", None)
    if isinstance(agent_name, str) and agent_name.strip():
        return agent_name.strip()
    return None


def _record_to_payload(record: ArtifactRecord) -> dict[str, Any]:
    return record.to_ref()


def _ensure_workspace_access(sandbox: ArtifactSandbox, artifact_id: str) -> ArtifactRecord:
    record = sandbox.get_artifact(artifact_id)
    if record is None:
        raise FileNotFoundError(f"Artifact '{artifact_id}' not found")
    current_workspace = sandbox.current_workspace_id()
    if record.workspace_id != current_workspace:
        raise FileNotFoundError(f"Artifact '{artifact_id}' not found")
    return record


def _image_extension(mime_type: str) -> str:
    return _IMAGE_EXTENSION_BY_MIME_TYPE.get(mime_type, ".bin")


__all__ = ["ArtifactToolkit", "save_artifact"]
