# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import os
import pathlib
import shutil
import tempfile
from dataclasses import dataclass
from typing import Any

import yaml
from cookiecutter.main import cookiecutter


@dataclass
class TemplateConfig:
    name: str
    description: str
    settings: dict[str, bool | list[str]]

    @classmethod
    def from_file(cls, config_path: pathlib.Path) -> "TemplateConfig":
        """Load template config from file with validation"""
        try:
            with open(config_path) as f:
                data = yaml.safe_load(f)

            if not isinstance(data, dict):
                raise ValueError(f"Invalid template config format in {config_path}")

            required_fields = ["name", "description", "settings"]
            missing_fields = [f for f in required_fields if f not in data]
            if missing_fields:
                raise ValueError(
                    f"Missing required fields in template config: {missing_fields}"
                )

            return cls(
                name=data["name"],
                description=data["description"],
                settings=data["settings"],
            )
        except yaml.YAMLError as err:
            raise ValueError(f"Invalid YAML in template config: {err}") from err
        except Exception as err:
            raise ValueError(f"Error loading template config: {err}") from err


OVERWRITE_FOLDERS = ["app", "frontend", "tests", "notebooks"]
TEMPLATE_CONFIG_FILE = ".templateconfig.yaml"
DEPLOYMENT_FOLDERS = ["cloud_run", "agent_engine"]
DEFAULT_FRONTEND = "streamlit"


def get_available_agents(deployment_target: str | None = None) -> dict:
    """Dynamically load available agents from the agents directory.

    Args:
        deployment_target: Optional deployment target to filter agents
    """
    # Define priority agents that should appear first
    PRIORITY_AGENTS = [
        "langgraph_base_react"  # Add other priority agents here as needed
    ]

    agents_list = []
    priority_agents = []
    agents_dir = pathlib.Path(__file__).parent.parent.parent.parent / "agents"

    for agent_dir in agents_dir.iterdir():
        if agent_dir.is_dir() and not agent_dir.name.startswith("__"):
            template_config_path = agent_dir / "template" / ".templateconfig.yaml"
            if template_config_path.exists():
                try:
                    with open(template_config_path) as f:
                        config = yaml.safe_load(f)
                    agent_name = agent_dir.name

                    # Skip if deployment target specified and agent doesn't support it
                    if deployment_target:
                        targets = config.get("settings", {}).get(
                            "deployment_targets", []
                        )
                        if isinstance(targets, str):
                            targets = [targets]
                        if deployment_target not in targets:
                            continue

                    description = config.get("description", "No description available")
                    agent_info = {"name": agent_name, "description": description}

                    # Add to priority list or regular list based on agent name
                    if agent_name in PRIORITY_AGENTS:
                        priority_agents.append(agent_info)
                    else:
                        agents_list.append(agent_info)
                except Exception as e:
                    logging.warning(f"Could not load agent from {agent_dir}: {e}")

    # Only sort the non-priority agents
    agents_list.sort(key=lambda x: x["name"])

    # Combine priority agents with regular agents (no sorting of priority_agents)
    combined_agents = priority_agents + agents_list

    # Convert to numbered dictionary starting from 1
    agents = {i + 1: agent for i, agent in enumerate(combined_agents)}

    return agents


def load_template_config(template_dir: pathlib.Path) -> dict[str, Any]:
    """Read .templateconfig.yaml file to get agent configuration."""
    config_file = template_dir / TEMPLATE_CONFIG_FILE
    if not config_file.exists():
        return {}

    try:
        with open(config_file) as f:
            config = yaml.safe_load(f)
            return config if config else {}
    except Exception as e:
        logging.error(f"Error loading template config: {e}")
        return {}


def get_deployment_targets(agent_name: str) -> list:
    """Get available deployment targets for the selected agent."""
    template_path = (
        pathlib.Path(__file__).parent.parent.parent.parent
        / "agents"
        / agent_name
        / "template"
    )
    config = load_template_config(template_path)

    if not config:
        return []

    targets = config.get("settings", {}).get("deployment_targets", [])
    return targets if isinstance(targets, list) else [targets]


def prompt_deployment_target(agent_name: str) -> str:
    """Ask user to select a deployment target for the agent."""
    targets = get_deployment_targets(agent_name)

    # Define deployment target friendly names and descriptions
    TARGET_INFO = {
        "agent_engine": {
            "display_name": "Vertex AI Agent Engine",
            "description": "Vertex AI Managed platform for scalable agent deployments",
        },
        "cloud_run": {
            "display_name": "Cloud Run",
            "description": "GCP Serverless container execution",
        },
    }

    if not targets:
        return ""

    from rich.console import Console

    console = Console()
    console.print("\n> Please select a deployment target:")
    for idx, target in enumerate(targets, 1):
        info = TARGET_INFO.get(target, {})
        display_name = info.get("display_name", target)
        description = info.get("description", "")
        console.print(f"{idx}. {display_name} - {description}")

    from rich.prompt import IntPrompt

    choice = IntPrompt.ask(
        "\nEnter the number of your deployment target choice",
        default=1,
        show_default=True,
    )
    return targets[choice - 1]


def prompt_data_ingestion(agent_name: str) -> bool:
    """Ask user if they want to include data pipeline if the agent supports it."""
    template_path = (
        pathlib.Path(__file__).parent.parent.parent.parent
        / "agents"
        / agent_name
        / "template"
    )
    config = load_template_config(template_path)

    if config:
        # If requires_data_ingestion is true, return True without prompting
        if config.get("settings", {}).get("requires_data_ingestion"):
            return True

        # Only prompt if the agent has optional data processing support
        if "data_ingestion" in config.get("settings", {}):
            from rich.prompt import Prompt

            return (
                Prompt.ask(
                    "\n> This agent supports a data pipeline. Would you like to include it?",
                    choices=["y", "n"],
                    default="n",
                ).lower()
                == "y"
            )
    return False


def get_template_path(agent_name: str, debug: bool = False) -> str:
    """Get the absolute path to the agent template directory."""
    current_dir = pathlib.Path(__file__).parent.parent.parent.parent
    template_path = current_dir / "agents" / agent_name / "template"
    if debug:
        logging.debug(f"Looking for template in: {template_path}")
        logging.debug(f"Template exists: {template_path.exists()}")
        if template_path.exists():
            logging.debug(f"Template contents: {list(template_path.iterdir())}")

    if not template_path.exists():
        raise ValueError(f"Template directory not found at {template_path}")

    return str(template_path)


def copy_data_ingestion_files(project_template: pathlib.Path) -> None:
    """Copy data processing files to the project template.

    Args:
        project_template: Path to the project template directory
    """
    data_ingestion_src = pathlib.Path(__file__).parent.parent.parent / "data_ingestion"
    data_ingestion_dst = project_template / "data_ingestion"

    if data_ingestion_src.exists():
        logging.debug(
            f"Copying data processing files from {data_ingestion_src} to {data_ingestion_dst}"
        )
        copy_files(data_ingestion_src, data_ingestion_dst, overwrite=True)
    else:
        logging.warning(
            f"Data processing source directory not found at {data_ingestion_src}"
        )


def process_template(
    agent_name: str,
    template_dir: str,
    project_name: str,
    deployment_target: str | None = None,
    include_data_ingestion: bool = False,
    output_dir: pathlib.Path | None = None,
) -> None:
    """Process the template directory and create a new project.

    Args:
        agent_name: Name of the agent template to use
        template_dir: Directory containing the template files
        project_name: Name of the project to create
        deployment_target: Optional deployment target (agent_engine or cloud_run)
        include_data_ingestion: Whether to include data pipeline components
        output_dir: Optional output directory path, defaults to current directory
    """
    logging.debug(f"Processing template from {template_dir}")
    logging.debug(f"Project name: {project_name}")
    logging.debug(f"Include pipeline: {include_data_ingestion}")
    logging.debug(f"Output directory: {output_dir}")

    # Get paths
    agent_path = pathlib.Path(template_dir).parent  # Get parent of template dir
    logging.debug(f"agent path: {agent_path}")
    logging.debug(f"agent path exists: {agent_path.exists()}")
    logging.debug(
        f"agent path contents: {list(agent_path.iterdir()) if agent_path.exists() else 'N/A'}"
    )

    base_template_path = pathlib.Path(__file__).parent.parent.parent / "base_template"

    # Use provided output_dir or current directory
    destination_dir = output_dir if output_dir else pathlib.Path.cwd()

    # Create output directory if it doesn't exist
    if not destination_dir.exists():
        destination_dir.mkdir(parents=True)

    # Create a new temporary directory and use it as our working directory
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = pathlib.Path(temp_dir)

        # Important: Store the original working directory
        original_dir = pathlib.Path.cwd()

        try:
            os.chdir(temp_path)  # Change to temp directory

            # Create the cookiecutter template structure
            cookiecutter_template = temp_path / "template"
            cookiecutter_template.mkdir(parents=True)
            project_template = cookiecutter_template / "{{cookiecutter.project_name}}"
            project_template.mkdir(parents=True)

            # 1. First copy base template files
            base_template_path = (
                pathlib.Path(__file__).parent.parent.parent / "base_template"
            )
            copy_files(base_template_path, project_template, agent_name, overwrite=True)
            logging.debug(f"1. Copied base template from {base_template_path}")

            # 2. Process deployment target if specified
            if deployment_target and deployment_target in DEPLOYMENT_FOLDERS:
                deployment_path = (
                    pathlib.Path(__file__).parent.parent.parent
                    / "deployment_targets"
                    / deployment_target
                )
                if deployment_path.exists():
                    copy_files(
                        deployment_path,
                        project_template,
                        agent_name=agent_name,
                        overwrite=True,
                    )
                    logging.debug(
                        f"2. Processed deployment files for target: {deployment_target}"
                    )

            # 3. Copy data ingestion files if needed
            template_config = load_template_config(pathlib.Path(template_dir))
            requires_data_ingestion = template_config.get("settings", {}).get(
                "requires_data_ingestion", False
            )
            should_include_data_ingestion = (
                include_data_ingestion or requires_data_ingestion
            )

            if should_include_data_ingestion:
                logging.debug("3. Including data processing files")
                copy_data_ingestion_files(project_template)

            # 4. Process frontend files
            frontend_type = template_config.get("settings", {}).get(
                "frontend_type", DEFAULT_FRONTEND
            )
            copy_frontend_files(frontend_type, project_template)
            logging.debug(f"4. Processed frontend files for type: {frontend_type}")

            # 5. Finally, copy agent-specific files to override everything else
            if agent_path.exists():
                for folder in OVERWRITE_FOLDERS:
                    agent_folder = agent_path / folder
                    project_folder = project_template / folder
                    if agent_folder.exists():
                        logging.debug(f"5. Copying agent folder {folder} with override")
                        copy_files(
                            agent_folder, project_folder, agent_name, overwrite=True
                        )

            # Copy agent README.md if it exists
            agent_readme = agent_path / "README.md"
            if agent_readme.exists():
                agent_readme_dest = project_template / "agent_README.md"
                shutil.copy2(agent_readme, agent_readme_dest)
                logging.debug(
                    f"Copied agent README from {agent_readme} to {agent_readme_dest}"
                )

            # Load and validate template config first
            template_path = pathlib.Path(template_dir)
            config = load_template_config(template_path)
            if not config:
                raise ValueError(f"Could not load template config from {template_path}")

            # Validate deployment target
            available_targets = config.get("settings", {}).get("deployment_targets", [])
            if isinstance(available_targets, str):
                available_targets = [available_targets]

            if deployment_target and deployment_target not in available_targets:
                raise ValueError(
                    f"Invalid deployment target '{deployment_target}'. Available targets: {available_targets}"
                )

            # Load template config
            template_config = load_template_config(pathlib.Path(template_dir))

            # Check if data processing should be included
            requires_data_ingestion = template_config.get("settings", {}).get(
                "requires_data_ingestion", False
            )
            should_include_data_ingestion = (
                include_data_ingestion or requires_data_ingestion
            )

            if should_include_data_ingestion:
                logging.debug(
                    "Including data processing files based on template config or user request"
                )
                copy_data_ingestion_files(project_template)

            # Create cookiecutter.json in the template root
            # Process extra dependencies
            extra_deps = template_config.get("settings", {}).get(
                "extra_dependencies", []
            )
            otel_instrumentations = get_otel_instrumentations(dependencies=extra_deps)

            # Get frontend type from template config
            frontend_type = template_config.get("settings", {}).get(
                "frontend_type", DEFAULT_FRONTEND
            )

            cookiecutter_config = {
                "project_name": "my-project",
                "agent_name": agent_name,
                "agent_description": template_config.get("description", ""),
                "deployment_target": deployment_target or "",
                "frontend_type": frontend_type,
                "extra_dependencies": [extra_deps],
                "otel_instrumentations": otel_instrumentations,
                "data_ingestion": should_include_data_ingestion,
                "_copy_without_render": [
                    "*.ipynb",  # Don't render notebooks
                    "*.json",  # Don't render JSON files
                    "frontend/*",  # Don't render frontend directory
                    "tests/*",  # Don't render tests directory
                    "notebooks/*",  # Don't render notebooks directory
                    ".git/*",  # Don't render git directory
                    "__pycache__/*",  # Don't render cache
                    "**/__pycache__/*",
                    ".pytest_cache/*",
                    ".venv/*",
                    "*templates.py",  # Don't render templates files
                    "!*.py",  # render Python files
                    "!Makefile",  # DO render Makefile
                    "!README.md",  # DO render README.md
                ],
            }

            with open(cookiecutter_template / "cookiecutter.json", "w") as f:
                import json

                json.dump(cookiecutter_config, f, indent=4)

            logging.debug(f"Template structure created at {cookiecutter_template}")
            logging.debug(
                f"Directory contents: {list(cookiecutter_template.iterdir())}"
            )

            # Process the template
            cookiecutter(
                str(cookiecutter_template),
                no_input=True,
                extra_context={
                    "project_name": project_name,
                    "agent_name": agent_name,
                },
            )
            logging.debug("Template processing completed successfully")

            # Move the generated project to the final destination
            output_dir = temp_path / project_name
            final_destination = destination_dir / project_name

            logging.debug(f"Moving project from {output_dir} to {final_destination}")

            if output_dir.exists():
                if final_destination.exists():
                    shutil.rmtree(final_destination)
                shutil.copytree(output_dir, final_destination, dirs_exist_ok=True)
                logging.debug(f"Project successfully created at {final_destination}")

                # After copying template files, handle the lock file
                if deployment_target:
                    # Get the source lock file path
                    lock_path = (
                        pathlib.Path(__file__).parent.parent.parent.parent
                        / "src"
                        / "resources"
                        / "locks"
                        / f"uv-{agent_name}-{deployment_target}.lock"
                    )
                    logging.debug(f"Looking for lock file at: {lock_path}")
                    logging.debug(f"Lock file exists: {lock_path.exists()}")
                    if not lock_path.exists():
                        raise FileNotFoundError(f"Lock file not found: {lock_path}")
                    # Copy and rename to uv.lock in the project directory
                    shutil.copy2(lock_path, final_destination / "uv.lock")
                    logging.debug(
                        f"Copied lock file from {lock_path} to {final_destination}/uv.lock"
                    )

                    # Replace cookiecutter project name with actual project name in lock file
                    lock_file_path = final_destination / "uv.lock"
                    with open(lock_file_path, "r+", encoding="utf-8") as f:
                        content = f.read()
                        f.seek(0)
                        f.write(
                            content.replace(
                                "{{cookiecutter.project_name}}", project_name
                            )
                        )
                        f.truncate()
                    logging.debug(
                        f"Updated project name in lock file at {lock_file_path}"
                    )
            else:
                logging.error(f"Generated project directory not found at {output_dir}")
                raise FileNotFoundError(
                    f"Generated project directory not found at {output_dir}"
                )

        except Exception as e:
            logging.error(f"Failed to process template: {e!s}")
            raise

        finally:
            # Always restore the original working directory
            os.chdir(original_dir)


def should_exclude_path(path: pathlib.Path, agent_name: str) -> bool:
    """Determine if a path should be excluded based on the agent type."""
    if agent_name == "multimodal_live_api":
        # Exclude the unit test utils folder and app/utils folder for multimodal_live_api
        if "tests/unit/test_utils" in str(path) or "app/utils" in str(path):
            logging.debug(f"Excluding path for multimodal_live_api: {path}")
            return True
    return False


def copy_files(
    src: pathlib.Path,
    dst: pathlib.Path,
    agent_name: str | None = None,
    overwrite: bool = False,
) -> None:
    """
    Copy files with configurable behavior for exclusions and overwrites.

    Args:
        src: Source path
        dst: Destination path
        agent_name: Name of the agent (for agent-specific exclusions)
        overwrite: Whether to overwrite existing files (True) or skip them (False)
    """

    def should_skip(path: pathlib.Path) -> bool:
        """Determine if a file/directory should be skipped during copying."""
        if path.suffix in [".pyc"]:
            return True
        if "__pycache__" in str(path) or path.name == "__pycache__":
            return True
        if agent_name is not None and should_exclude_path(path, agent_name):
            return True
        return False

    if src.is_dir():
        if not dst.exists():
            dst.mkdir(parents=True)
        for item in src.iterdir():
            if should_skip(item):
                logging.debug(f"Skipping file/directory: {item}")
                continue
            d = dst / item.name
            if item.is_dir():
                copy_files(item, d, agent_name, overwrite)
            else:
                if overwrite or not d.exists():
                    logging.debug(f"Copying file: {item} -> {d}")
                    shutil.copy2(item, d)
                else:
                    logging.debug(f"Skipping existing file: {d}")
    else:
        if not should_skip(src):
            if overwrite or not dst.exists():
                shutil.copy2(src, dst)


def copy_frontend_files(frontend_type: str, project_template: pathlib.Path) -> None:
    """Copy files from the specified frontend folder directly to project root."""
    # Use default frontend if none specified
    frontend_type = frontend_type or DEFAULT_FRONTEND

    # Get the frontends directory path
    frontends_path = (
        pathlib.Path(__file__).parent.parent.parent / "frontends" / frontend_type
    )

    if frontends_path.exists():
        logging.debug(f"Copying frontend files from {frontends_path}")
        # Copy frontend files directly to project root instead of a nested frontend directory
        copy_files(frontends_path, project_template, overwrite=True)
    else:
        logging.warning(f"Frontend type directory not found: {frontends_path}")
        if frontend_type != DEFAULT_FRONTEND:
            logging.info(f"Falling back to default frontend: {DEFAULT_FRONTEND}")
            copy_frontend_files(DEFAULT_FRONTEND, project_template)


def copy_deployment_files(
    deployment_target: str, agent_name: str, project_template: pathlib.Path
) -> None:
    """Copy files from the specified deployment target folder."""
    if not deployment_target:
        return

    deployment_path = (
        pathlib.Path(__file__).parent.parent.parent
        / "deployment_targets"
        / deployment_target
    )

    if deployment_path.exists():
        logging.debug(f"Copying deployment files from {deployment_path}")
        # Pass agent_name to respect agent-specific exclusions
        copy_files(
            deployment_path, project_template, agent_name=agent_name, overwrite=True
        )
    else:
        logging.warning(f"Deployment target directory not found: {deployment_path}")


def get_otel_instrumentations(dependencies: list) -> list[list[str]]:
    """Returns OpenTelemetry instrumentation statements for enabled dependencies."""
    otel_deps = {
        "langgraph": "Instruments.LANGCHAIN",
        "crewai": "Instruments.CREW",
    }
    imports = []
    for dep in dependencies:
        if any(otel_dep in dep for otel_dep in otel_deps):
            imports.append(otel_deps[next(key for key in otel_deps if key in dep)])
    return [imports]
