from __future__ import annotations

import json
import logging
import re
from pathlib import Path
from typing import Any

import sentry_sdk
from ninja import NinjaAPI, Schema
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam

from aiservice.analytics.posthog import ph
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.llm import llm_client
from aiservice.llm_models import EXECUTE_MODEL
from authapp.auth import AuthenticatedRequest

_steps_section_re = re.compile(r"steps:\s*\n((?:[^\n]+\n)*)", re.MULTILINE)

_array_match_re = re.compile(r"(- name:.*?)(?=\n\n|\Z)", re.DOTALL)

_yaml_like_re = re.compile(r"(- name:.*)", re.DOTALL)

logger = logging.getLogger(__name__)

workflow_gen_api = NinjaAPI(urls_namespace="workflow_gen")

# Get the directory of the current file
current_dir = Path(__file__).parent
SYSTEM_PROMPT = (current_dir / "system_prompt.md").read_text()
USER_PROMPT = (current_dir / "user_prompt.md").read_text()


def _extract_yaml_steps(text: str) -> str | None:
    """Extract YAML steps section from LLM response."""
    # Try to find steps section
    steps_match = _steps_section_re.search(text)
    if steps_match:
        steps_content = steps_match.group(1)
        formatted_lines = []
        for line in steps_content.split("\n"):
            stripped_line = line.strip()
            if not stripped_line:
                continue
            if line.startswith("-"):
                formatted_lines.append(line)
            elif stripped_line.startswith("-"):
                # Prepend two spaces only once per line
                formatted_lines.append("  " + stripped_line)
            else:
                formatted_lines.append(line)
        return "steps:\n" + "\n".join(formatted_lines)

    # Try to find just the steps array without "steps:" prefix
    array_match = _array_match_re.search(text)
    if array_match:
        return "steps:\n" + array_match.group(1)

    # If no clear pattern, try to extract any YAML-like content
    yaml_match = _yaml_like_re.search(text)
    if yaml_match:
        content = yaml_match.group(1)
        return "steps:\n" + content

    return None


async def generate_workflow_steps_llm(
    repo_files: dict[str, str],
    directory_structure: dict[str, Any],
    user_id: str,
    trace_id: str = "",
    codeflash_config: dict[str, Any] | None = None,
) -> str | None:
    """Generate workflow steps using LLM."""
    # Format repo files for prompt
    files_text = "\n".join([f"**{path}:**\n{content[:8000]}\n" for path, content in repo_files.items()])

    # Format directory structure
    structure_text = json.dumps(directory_structure, indent=2)

    # Format codeflash config
    config_text = json.dumps(codeflash_config, indent=2) if codeflash_config else "Not available"

    user_prompt = USER_PROMPT.format(
        repo_files=files_text, directory_structure=structure_text, codeflash_config=config_text
    )

    system_message = ChatCompletionSystemMessageParam(role="system", content=SYSTEM_PROMPT)
    user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)

    debug_log_sensitive_data(f"Generating workflow steps with prompt length: {len(user_prompt)}")

    try:
        response = await llm_client.call(
            llm=EXECUTE_MODEL,
            messages=[system_message, user_message],
            call_type="workflow_generation",
            trace_id=trace_id or "",
            user_id=user_id,
            context={"num_files": len(repo_files)},
        )

        if not response.content:
            logger.warning("LLM returned empty response for workflow generation")
            return None

        response_text = response.content.strip()

        # Extract YAML steps
        steps_yaml = _extract_yaml_steps(response_text)
        if steps_yaml:
            logger.info("Successfully generated workflow steps (%d chars)", len(steps_yaml))
            return steps_yaml

        logger.warning("Could not extract valid YAML steps from LLM response: %s", response_text[:200])
        return None

    except Exception:
        return None


class WorkflowGenInputSchema(Schema):
    repo_files: dict[str, str]  # path -> content
    directory_structure: dict[str, Any]  # 2-level nested structure
    codeflash_config: dict[str, Any] | None = None
    trace_id: str = ""  # Optional trace ID for observability


class WorkflowGenResponseSchema(Schema):
    workflow_steps: str  # YAML string for steps section


class WorkflowGenErrorResponseSchema(Schema):
    error: str


@workflow_gen_api.post(
    "/",
    response={200: WorkflowGenResponseSchema, 400: WorkflowGenErrorResponseSchema, 500: WorkflowGenErrorResponseSchema},
)
async def generate_workflow_steps(
    request: AuthenticatedRequest, data: WorkflowGenInputSchema
) -> tuple[int, WorkflowGenResponseSchema | WorkflowGenErrorResponseSchema]:
    """Generate GitHub Actions workflow steps based on repository analysis."""
    ph(request.user, "aiservice-workflow-gen-called")

    try:
        # Validate input
        if not data.repo_files and not data.directory_structure:
            return 400, WorkflowGenErrorResponseSchema(
                error="At least one of repo_files or directory_structure must be provided"
            )

        # Generate workflow steps using LLM
        workflow_steps = await generate_workflow_steps_llm(
            repo_files=data.repo_files,
            directory_structure=data.directory_structure,
            user_id=request.user,
            trace_id=data.trace_id,
            codeflash_config=data.codeflash_config,
        )

        if not workflow_steps:
            logger.warning("Failed to generate workflow steps from LLM")
            return 500, WorkflowGenErrorResponseSchema(
                error="Failed to generate workflow steps. Please try again or use the static template."
            )

        ph(request.user, "aiservice-workflow-gen-success", properties={"steps_length": len(workflow_steps)})
        return 200, WorkflowGenResponseSchema(workflow_steps=workflow_steps)

    except Exception:
        logger.exception("Error in generate_workflow_steps endpoint")
        sentry_sdk.capture_exception()
        return 500, WorkflowGenErrorResponseSchema(error="Internal server error while generating workflow steps")
