"""
CodeHandler: Utilities for code build, extraction, and manipulation in agent generation.
"""

import logging
import re

logger = logging.getLogger(__name__)


class CodeHandler:
    @staticmethod
    def clean_generated_code(code: str) -> str:
        if not code:
            return ""

        # Remove markdown code blocks
        if "```dana" in code:
            start = code.find("```dana") + 7
            end = code.find("```", start)
            if end != -1:
                code = code[start:end].strip()
        elif "```python" in code:
            start = code.find("```python") + 9
            end = code.find("```", start)
            if end != -1:
                code = code[start:end].strip()
        elif "```na" in code:
            start = code.find("```na") + 5
            end = code.find("```", start)
            if end != -1:
                code = code[start:end].strip()
        elif "```" in code:
            start = code.find("```") + 3
            newline_pos = code.find("\n", start)
            if newline_pos != -1:
                start = newline_pos + 1
            end = code.find("```", start)
            if end != -1:
                code = code[start:end].strip()

        # Remove any remaining markdown artifacts
        lines = code.split("\n")
        cleaned_lines = []
        for line in lines:
            stripped = line.strip()
            # Skip lines that are just markdown code block markers
            if stripped in ["```", "```python", "```dana", "```na"]:
                continue
            cleaned_lines.append(line)

        code = "\n".join(cleaned_lines).strip()

        # Ensure proper spacing between code blocks
        # Add newline after agent definitions if missing, but don't add extra newlines
        code = re.sub(r"(agent\s+\w+:)\n(?!\n)", r"\1\n\n", code)
        # Remove excessive newlines (more than 2 consecutive)
        code = re.sub(r"\n{3,}", r"\n\n", code)
        # Ensure exactly one newline after agent definitions, but not extra newlines
        code = re.sub(r"(agent\s+\w+:\n)\n+(\s+)", r"\1\n\2", code)
        # Remove extra newlines between agent definition and first field
        code = re.sub(r"(agent\s+\w+:\n)\n+(\s+\w+:)", r"\1\n\2", code)
        # Ensure exactly one newline after agent definition (not two)
        code = re.sub(r"(agent\s+\w+:\n)\n(\s+\w+:)", r"\1\2", code)
        # Ensure newline between agent fields and function definitions
        code = re.sub(r"(\s+\w+:\s+[^\n]+\n)(def\s+)", r"\1\n\2", code)
        # Ensure no trailing newline for empty content
        if not code.strip():
            code = ""
        # Ensure exactly one newline at the end, but not for empty input
        if code.strip():
            code = code.rstrip()
        else:
            code = ""

        return code

    @staticmethod
    def parse_multi_file_response(response: str) -> dict:
        files = []
        project_name = "Georgia Training Project"
        project_description = "Georgia training project generated by Dana from user requirements"

        # Check if response is empty or whitespace only
        if not response or not response.strip():
            logger.warning("Empty response received, using fallback template")
            # Return a basic project structure
            return {
                "files": [
                    {
                        "filename": "main.na",
                        "content": CodeHandler.get_file_template("main.na"),
                        "file_type": "main",
                        "description": "Main agent file",
                        "dependencies": [],
                    }
                ],
                "main_file": "main.na",
                "name": project_name,
                "description": project_description,
                "structure_type": "simple",
            }

        try:
            # Try parsing as JSON first
            import json

            json_data = json.loads(response.strip())

            # Extract file contents from JSON
            for filename, content in json_data.items():
                if filename.endswith(".na"):
                    files.append(
                        {
                            "filename": filename,
                            "content": content.strip(),
                            "file_type": CodeHandler.determine_file_type(filename),
                            "description": CodeHandler.get_file_description(filename),
                            "dependencies": CodeHandler.extract_dependencies(content),
                        }
                    )

        except (json.JSONDecodeError, ValueError) as e:
            # Fallback to FILE_START/FILE_END format if JSON parsing fails
            logger.warning(f"Failed to parse JSON response, falling back to FILE_START/FILE_END format: {e}")

            lines = response.split("\n")
            current_file = None
            current_content = []
            for line in lines:
                if line.startswith("FILE_START:"):
                    if current_file:
                        files.append(
                            {
                                "filename": current_file,
                                "content": "\n".join(current_content).strip(),
                                "file_type": CodeHandler.determine_file_type(current_file),
                                "description": CodeHandler.get_file_description(current_file),
                                "dependencies": CodeHandler.extract_dependencies("\n".join(current_content)),
                            }
                        )
                    current_file = line.split(":", 1)[1].strip()
                    current_content = []
                elif line.startswith("FILE_END:"):
                    if current_file:
                        files.append(
                            {
                                "filename": current_file,
                                "content": "\n".join(current_content).strip(),
                                "file_type": CodeHandler.determine_file_type(current_file),
                                "description": CodeHandler.get_file_description(current_file),
                                "dependencies": CodeHandler.extract_dependencies("\n".join(current_content)),
                            }
                        )
                        current_file = None
                        current_content = []
                elif current_file:
                    # Skip markdown code block markers
                    if line.strip().startswith("```"):
                        continue
                    current_content.append(line)
            if current_file and current_content:
                files.append(
                    {
                        "filename": current_file,
                        "content": "\n".join(current_content).strip(),
                        "file_type": CodeHandler.determine_file_type(current_file),
                        "description": CodeHandler.get_file_description(current_file),
                        "dependencies": CodeHandler.extract_dependencies("\n".join(current_content)),
                    }
                )

        # Ensure all required files are present
        project = {
            "files": files,
            "main_file": "main.na",
            "name": project_name,
            "description": project_description,
            "structure_type": "complex",  # Default structure type
        }

        project = CodeHandler.ensure_all_files_present(project)

        # --- BEGIN PATCH: Ensure all imported tools are defined in tools.na ---
        # Find all symbols imported from tools in methods.na
        methods_file = next((f for f in project["files"] if f["filename"] == "methods.na"), None)
        tools_file = next((f for f in project["files"] if f["filename"] == "tools.na"), None)
        if methods_file and tools_file:
            import_re = re.compile(r"from tools import ([a-zA-Z0-9_,\s]+)")
            import_matches = import_re.findall(methods_file["content"])
            imported_tools = set()
            for match in import_matches:
                # Split by comma and handle multiple imports
                # Also split by newlines to handle cases where the import statement spans multiple lines
                tools = []
                for part in match.split(","):
                    part = part.strip()
                    # Split by newlines and take only the first part (before any newline)
                    if "\n" in part:
                        part = part.split("\n")[0].strip()
                    if part:
                        tools.append(part)
                imported_tools.update(tools)

            # Find all defined symbols in tools.na (simple assignment or def)
            defined_tools = set()
            for line in tools_file["content"].split("\n"):
                line = line.strip()
                if line.startswith("def "):
                    defined_tools.add(line.split()[1].split("(")[0])
                elif "=" in line:
                    defined_tools.add(line.split("=")[0].strip())
            # Add stubs for missing tools
            missing_tools = imported_tools - defined_tools
            if missing_tools:
                stub_lines = []
                for tool in missing_tools:
                    stub_lines.append(f"def {tool}(query: str) -> str:")
                    stub_lines.append("    return reason(f'Reasoning for: {query}')  # Auto-generated stub")
                if tools_file["content"] and not tools_file["content"].endswith("\n"):
                    tools_file["content"] += "\n"
                tools_file["content"] += "\n" + "\n".join(stub_lines)
        # --- END PATCH ---

        # Clean up any remaining markdown artifacts from all files
        for file_info in project["files"]:
            file_info["content"] = CodeHandler.clean_generated_code(file_info["content"])

        return project

    @staticmethod
    def ensure_all_files_present(project: dict) -> dict:
        """Ensure all required files are present in the project."""
        required_files = ["main.na", "workflows.na", "methods.na", "common.na", "knowledge.na", "tools.na"]
        existing_files = {f["filename"] for f in project["files"]}

        for required_file in required_files:
            if required_file not in existing_files:
                # Create missing file with appropriate template
                template_content = CodeHandler.get_file_template(required_file)
                project["files"].append(
                    {
                        "filename": required_file,
                        "content": template_content,
                        "file_type": CodeHandler.determine_file_type(required_file),
                        "description": CodeHandler.get_file_description(required_file),
                        "dependencies": CodeHandler.extract_dependencies(template_content),
                    }
                )

        # Clean up any remaining markdown artifacts from all files
        for file_info in project["files"]:
            file_info["content"] = CodeHandler.clean_generated_code(file_info["content"])

        return project

    @staticmethod
    def get_file_template(filename: str) -> str:
        """Get template content for missing files."""
        templates = {
            "main.na": """from workflows import workflow
from common import AgentPackage

agent Georgia:
    name: str = "Georgia"
    description: str = "A specialized agent trained by Dana to help with various tasks"

def solve(self : Georgia, query: str) -> str:
    package = AgentPackage(query=query)
    return workflow(package)

this_agent = Georgia()

# Use this_agent.solve() in your application
# response = this_agent.solve("Your query here")""",
            "workflows.na": """from methods import process_request
from methods import generate_response

workflow = process_request | generate_response""",
            "methods.na": """from common import AgentPackage

def process_request(package: AgentPackage) -> AgentPackage:
    # Process the user request
    package.processed = True
    return package

def generate_response(package: AgentPackage) -> str:
    # Generate response using reasoning
    return reason(f"Help me with: {package.query}")""",
            "common.na": '''# Data structures for the agent
struct AgentPackage:
    query: str
    processed: bool = False
    result: str = ""''',
            "knowledge.na": '''"""Knowledge base/resource configurations.

Knowledge Description:
- No external knowledge sources required for this agent
- Uses only built-in knowledge and reasoning capabilities
"""

# No external knowledge sources needed for this agent''',
            "tools.na": '''"""Tool definitions and integrations.

Tools Description:
- No external tools required for this agent
- Uses only built-in reasoning capabilities
"""

# No external tools required - this agent uses only built-in reasoning capabilities''',
        }

        return templates.get(filename, f"# {filename} - Generated template file")

    @staticmethod
    def determine_file_type(filename: str) -> str:
        if filename == "main.na":
            return "agent"
        elif filename == "workflows.na":
            return "workflow"
        elif filename == "knowledge.na":
            return "resources"
        elif filename == "methods.na":
            return "methods"
        elif filename == "tools.na":
            return "tools"
        elif filename == "common.na":
            return "common"
        return "other"

    @staticmethod
    def get_file_description(filename: str) -> str:
        descriptions = {
            "main.na": "Main agent definition and orchestration",
            "workflows.na": "Workflow orchestration and pipelines",
            "knowledge.na": "Resource configurations (RAG, databases, APIs)",
            "methods.na": "Core processing methods and utilities",
            "tools.na": "Tool definitions and integrations",
            "common.na": "Shared data structures and utilities",
        }
        return descriptions.get(filename, "Dana agent file")

    @staticmethod
    def extract_dependencies(content: str) -> list[str]:
        dependencies = []
        lines = content.split("\n")
        for line in lines:
            line = line.strip()
            if line.startswith("import ") and not line.endswith(".py"):
                dep = line.replace("import ", "").strip()
                if dep not in dependencies:
                    dependencies.append(dep)
            elif line.startswith("from ") and " import " in line:
                dep = line.split(" import ")[0].replace("from ", "").strip()
                if not dep.endswith(".py") and dep not in dependencies:
                    dependencies.append(dep)
        return dependencies

    @staticmethod
    def get_fallback_template() -> str:
        return '''"""Basic Agent Template."""

# Agent Card declaration
agent BasicAgent:
    name : str = "Basic Agent"
    description : str = "A basic agent that can handle general queries."

# Agent's problem solver
def solve(basic_agent : BasicAgent, problem : str):
    """Solve a problem using reasoning."""
    return reason(f"Help me to answer the question: {problem}")

# Example usage
example_input = "Hello, how can you help me?"
print(solve(BasicAgent(), example_input))'''
