import ast
from pathlib import Path
from typing import cast

from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai.models import Model
from pydantic_ai.tools import ToolDefinition

from .persist import persist_skill, update_skill
from .render import render_script_boilerplate, render_skill_draft, render_skill_script
from .types import (
    CreationStage,
    NonEmptyString,
    SkillCreationDeps,
    SkillScript,
    SkillScriptArgument,
    SkillScriptDraft,
    SkillScriptFileName,
    SkillSpec,
)

PROMPT_PATH = Path(__file__).parent / "prompt.md"

STAGE_TOOLS: dict[CreationStage, set[str]] = {
    "PLANNING": {"create_script", "remove_script", "complete_skill", "cancel_skill"},
    "SCRIPTING": {"complete_script", "cancel_script"},
    "CONCLUDED": set(),
}

STAGE_RETRY_MESSAGES: dict[CreationStage, str] = {
    "PLANNING": (
        "A skill ainda não foi finalizada nem cancelada. "
        "Use complete_skill para salvar ou cancel_skill para descartar."
    ),
    "SCRIPTING": (
        "Há um rascunho de script em andamento. "
        "Use complete_script para finalizar ou cancel_script para descartar."
    ),
    "CONCLUDED": "",
}


def get_skill_creation_stage(deps: SkillCreationDeps) -> CreationStage:
    if deps.concluded:
        return "CONCLUDED"
    if deps.current_script is not None:
        return "SCRIPTING"
    return "PLANNING"


async def prepare_skill_creation_tool(
    ctx: RunContext[SkillCreationDeps],
    tool: ToolDefinition,
) -> ToolDefinition | None:
    stage = get_skill_creation_stage(ctx.deps)
    allowed_tools = STAGE_TOOLS[stage]

    if tool.name in allowed_tools:
        return tool

    return None


def build_skill_creation_agent(model: Model) -> Agent[SkillCreationDeps, str]:
    agent: Agent[SkillCreationDeps, str] = Agent(
        model,
        deps_type=SkillCreationDeps,
        system_prompt=PROMPT_PATH.read_text(encoding="utf-8"),
        retries=3,
    )

    @agent.tool(prepare=prepare_skill_creation_tool)
    async def create_script(
        ctx: RunContext[SkillCreationDeps],
        file_name: SkillScriptFileName,
        when_to_use: NonEmptyString,
        dependencies: set[NonEmptyString],
        args: list[SkillScriptArgument],
    ) -> str:
        """Inicia a criação de um script Python para automatizar tarefas da skill.

        Use esta ferramenta para definir o nome do arquivo, sua finalidade, dependências e
        os argumentos que ele aceitará via linha de comando.

        Args:
            file_name: Nome do arquivo (ex.: "buscar_dados.py"). Deve terminar em .py.
            when_to_use: Descrição sucinta de quando este script deve ser executado.
            dependencies: Lista de bibliotecas externas (ex.: ["httpx", "beautifulsoup4"])
            args: Definição dos argumentos que o script receberá.
        """

        skill_draft = ctx.deps.skill_draft
        skill_draft.scripts = [
            s for s in skill_draft.scripts if s.file_name.casefold() != file_name.casefold()
        ]

        ctx.deps.current_script = SkillScriptDraft(
            file_name=file_name,
            when_to_use=when_to_use,
            dependencies=dependencies,
            args=args,
        )

        boilerplate = render_script_boilerplate(
            file_name=file_name,
            when_to_use=when_to_use,
            dependencies=dependencies,
            args=args,
        )

        return (
            f"Boilerplate de {file_name}:\n\n```python\n{boilerplate}\n```\n\n"
            "Preencha a lógica da função principal para finalizar o script."
        )

    @agent.tool(prepare=prepare_skill_creation_tool)
    async def complete_script(
        ctx: RunContext[SkillCreationDeps],
        body: NonEmptyString,
    ) -> str:
        """Finaliza o script Python fornecendo a lógica da função principal.

        Regras para o body:
        - Escreva apenas o que vai dentro da função main().
        - Não inclua 'def main(...)', 'import argparse' ou blocos 'if __name__'.
        - Os argumentos previamente definidos já estão disponíveis como variáveis locais.
        - O código deve imprimir o resultado final no console (stdout).

        Args:
            body: O código Python com a lógica interna da função.
        """
        skill_draft = ctx.deps.skill_draft
        current = cast(SkillScriptDraft, ctx.deps.current_script)

        script = SkillScript(**current.model_dump(), body=body)

        skill_script_code = render_skill_script(script)

        try:
            ast.parse(skill_script_code)
        except SyntaxError as e:
            raise ModelRetry(f"Script com erro na linha {e.lineno}: {e.msg}.")

        skill_draft.scripts.append(script)
        ctx.deps.current_script = None
        skill_markdown_draft = render_skill_draft(
            name=skill_draft.name,
            description=skill_draft.description,
            scripts=skill_draft.scripts,
        )

        return (
            f"Script {script.file_name}:\n\n```python\n{skill_script_code}\n```\n\n"
            f"Rascunho do SKILL.md:\n\n```markdown\n{skill_markdown_draft}\n```\n\n"
            "Você pode criar mais scripts ou finalizar a skill."
        )

    @agent.tool(prepare=prepare_skill_creation_tool)
    async def cancel_script(ctx: RunContext[SkillCreationDeps]) -> str:
        """Descarta o rascunho do script em construção."""
        ctx.deps.current_script = None
        return "Script cancelado. Você pode criar um novo script ou finalizar a skill."

    @agent.tool(prepare=prepare_skill_creation_tool)
    async def remove_script(
        ctx: RunContext[SkillCreationDeps],
        file_name: SkillScriptFileName,
    ) -> str:
        """Remove um script do rascunho da skill.

        Use quando quiser excluir um script existente da versão final.
        Scripts não removidos explicitamente são mantidos automaticamente.

        Args:
            file_name: Nome do arquivo do script a remover (ex.: "buscar_dados.py").
        """
        draft = ctx.deps.skill_draft
        before = len(draft.scripts)
        draft.scripts = [s for s in draft.scripts if s.file_name.casefold() != file_name.casefold()]
        if len(draft.scripts) == before:
            raise ModelRetry(
                f"Script '{file_name}' não encontrado no rascunho. "
                f"Scripts disponíveis: {[s.file_name for s in draft.scripts]}"
            )
        remaining = [s.file_name for s in draft.scripts]
        return f"Script '{file_name}' removido. Scripts restantes: {remaining}"

    @agent.tool(prepare=prepare_skill_creation_tool)
    async def cancel_skill(ctx: RunContext[SkillCreationDeps]) -> str:
        """Cancela todo o processo de criação da skill e descarta o rascunho."""
        ctx.deps.concluded = True
        return "Criação da skill cancelada. Nenhum arquivo foi gravado."

    @agent.tool(prepare=prepare_skill_creation_tool)
    async def complete_skill(
        ctx: RunContext[SkillCreationDeps],
        instructions: NonEmptyString,
    ) -> str:
        """Conclui a criação da skill, definindo as instruções de uso e salvando-a.

        As instruções devem ser escritas em Markdown e explicar claramente que problemas
        a skill resolve, como ela faz isso e casos de uso para os scripts criados (se houver).
        Isso substituirá o placeholder de instruções no rascunho do SKILL.md.

        Args:
            instructions: Texto detalhado de orientação sobre a skill.
        """
        skill_draft = ctx.deps.skill_draft

        spec = SkillSpec(
            name=skill_draft.name,
            description=skill_draft.description,
            instructions=instructions,
            scripts=skill_draft.scripts,
        )

        if ctx.deps.is_update:
            result = update_skill(ctx.deps.skills_directory, spec)
        else:
            result = persist_skill(ctx.deps.skills_directory, spec)

        ctx.deps.concluded = True
        action = "atualizada" if ctx.deps.is_update else "criada e está pronta para uso"
        return f"Sucesso! A skill '{result.name}' foi {action}."

    @agent.output_validator
    def validate_stage_allows_output(ctx: RunContext[SkillCreationDeps], output: str) -> str:
        stage = get_skill_creation_stage(ctx.deps)
        if stage != "CONCLUDED":
            raise ModelRetry(STAGE_RETRY_MESSAGES[stage])
        return output

    return agent
