from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
from datetime import datetime
import json

from pygeai import logger
from pygeai.lab.managers import AILabManager
from pygeai.lab.models import FilterSettings
from pygeai.assistant.rag.clients import RAGAssistantClient
from pygeai.assistant.rag.mappers import RAGAssistantMapper
from pygeai.core.files.managers import FileManager
from pygeai.core.secrets.clients import SecretClient
from pygeai.organization.limits.managers import UsageLimitManager
from pygeai.core.utils.console import Console
from pygeai.admin.clients import AdminClient


@dataclass
class ResourceDiff:
    resource_type: str
    only_in_source: List[Dict[str, Any]] = field(default_factory=list)
    only_in_destination: List[Dict[str, Any]] = field(default_factory=list)
    in_both: List[Dict[str, Any]] = field(default_factory=list)
    total_source: int = 0
    total_destination: int = 0


@dataclass
class ProjectDiff:
    source_project_id: str
    destination_project_id: str
    source_instance: str
    destination_instance: str
    resource_diffs: Dict[str, ResourceDiff] = field(default_factory=dict)
    timestamp: str = ""
    summary: Dict[str, int] = field(default_factory=dict)


class ProjectComparator:
    def __init__(
        self,
        from_api_key: str,
        from_instance: str,
        to_api_key: str,
        to_instance: str,
        from_project_id: Optional[str] = None,
        to_project_id: Optional[str] = None,
        from_organization_id: Optional[str] = None,
        to_organization_id: Optional[str] = None,
        from_organization_api_key: Optional[str] = None,
        to_organization_api_key: Optional[str] = None,
        resource_types: Optional[List[str]] = None
    ):
        self.from_api_key = from_api_key
        self.from_instance = from_instance
        self.to_api_key = to_api_key
        self.to_instance = to_instance
        self.from_organization_api_key = from_organization_api_key
        self.to_organization_api_key = to_organization_api_key
        self.resource_types = resource_types or [
            "agents", "tools", "processes", "tasks", "artifact_types",
            "rag_assistants", "files", "usage_limits", "secrets"
        ]

        if from_project_id is None or from_organization_id is None:
            from_admin = AdminClient(api_key=self.from_api_key, base_url=self.from_instance)
            from_token_info = from_admin.validate_api_token()
            self.from_project_id = from_project_id if from_project_id is not None else from_token_info.get("projectId")
            self.from_organization_id = from_organization_id if from_organization_id is not None else from_token_info.get("organizationId")
            logger.info(f"Auto-detected source - project: {self.from_project_id}, org: {self.from_organization_id}")
        else:
            self.from_project_id = from_project_id
            self.from_organization_id = from_organization_id

        if to_project_id is None or to_organization_id is None:
            to_admin = AdminClient(api_key=self.to_api_key, base_url=self.to_instance)
            to_token_info = to_admin.validate_api_token()
            self.to_project_id = to_project_id if to_project_id is not None else to_token_info.get("projectId")
            self.to_organization_id = to_organization_id if to_organization_id is not None else to_token_info.get("organizationId")
            logger.info(f"Auto-detected destination - project: {self.to_project_id}, org: {self.to_organization_id}")
        else:
            self.to_project_id = to_project_id
            self.to_organization_id = to_organization_id

        self._source_lab_manager = AILabManager(
            api_key=self.from_api_key,
            base_url=self.from_instance
        )
        self._dest_lab_manager = AILabManager(
            api_key=self.to_api_key,
            base_url=self.to_instance
        )

    def compare(self) -> ProjectDiff:
        logger.info(f"Starting comparison between projects {self.from_project_id} and {self.to_project_id}")
        
        diff = ProjectDiff(
            source_project_id=self.from_project_id,
            destination_project_id=self.to_project_id,
            source_instance=self.from_instance,
            destination_instance=self.to_instance,
            timestamp=datetime.utcnow().isoformat()
        )

        comparison_methods = {
            "agents": self._compare_agents,
            "tools": self._compare_tools,
            "processes": self._compare_processes,
            "tasks": self._compare_tasks,
            "artifact_types": self._compare_artifact_types,
            "rag_assistants": self._compare_rag_assistants,
            "files": self._compare_files,
            "usage_limits": self._compare_usage_limits,
            "secrets": self._compare_secrets
        }

        total_differences = 0
        for resource_type in self.resource_types:
            if resource_type in comparison_methods:
                try:
                    logger.debug(f"Comparing {resource_type}...")
                    resource_diff = comparison_methods[resource_type]()
                    diff.resource_diffs[resource_type] = resource_diff
                    total_differences += len(resource_diff.only_in_source) + len(resource_diff.only_in_destination)
                except Exception as e:
                    logger.error(f"Failed to compare {resource_type}: {e}")
                    Console.write_stderr(f"Warning: Could not compare {resource_type}: {e}")

        diff.summary = {
            "resource_types_compared": len(diff.resource_diffs),
            "total_differences": total_differences
        }

        logger.info(f"Comparison completed: {len(diff.resource_diffs)} resource types, {total_differences} differences")
        return diff

    def _compare_agents(self) -> ResourceDiff:
        try:
            source_list = self._source_lab_manager.list_agents(FilterSettings(count=1000))
            source_agents = [{"id": a.id, "name": a.name} for a in source_list.agents if a.id]
        except Exception as e:
            logger.warning(f"Could not retrieve source agents: {e}")
            source_agents = []

        try:
            dest_list = self._dest_lab_manager.list_agents(FilterSettings(count=1000))
            dest_agents = [{"id": a.id, "name": a.name} for a in dest_list.agents if a.id]
        except Exception as e:
            logger.warning(f"Could not retrieve destination agents: {e}")
            dest_agents = []

        return self._compare_resources(source_agents, dest_agents, "id", "name", "agents")

    def _compare_tools(self) -> ResourceDiff:
        try:
            source_list = self._source_lab_manager.list_tools(FilterSettings(count=1000))
            source_tools = [{"id": t.id, "name": t.name} for t in source_list.tools if t.id]
        except Exception as e:
            logger.warning(f"Could not retrieve source tools: {e}")
            source_tools = []

        try:
            dest_list = self._dest_lab_manager.list_tools(FilterSettings(count=1000))
            dest_tools = [{"id": t.id, "name": t.name} for t in dest_list.tools if t.id]
        except Exception as e:
            logger.warning(f"Could not retrieve destination tools: {e}")
            dest_tools = []

        return self._compare_resources(source_tools, dest_tools, "id", "name", "tools")

    def _compare_processes(self) -> ResourceDiff:
        try:
            source_list = self._source_lab_manager.list_processes(FilterSettings(count=1000))
            source_processes = [{"id": p.id, "name": p.name} for p in source_list.processes if p.id]
        except Exception as e:
            logger.warning(f"Could not retrieve source processes: {e}")
            source_processes = []

        try:
            dest_list = self._dest_lab_manager.list_processes(FilterSettings(count=1000))
            dest_processes = [{"id": p.id, "name": p.name} for p in dest_list.processes if p.id]
        except Exception as e:
            logger.warning(f"Could not retrieve destination processes: {e}")
            dest_processes = []

        return self._compare_resources(source_processes, dest_processes, "id", "name", "processes")

    def _compare_tasks(self) -> ResourceDiff:
        try:
            source_list = self._source_lab_manager.list_tasks(FilterSettings(count=1000))
            source_tasks = [{"id": t.id, "name": t.name} for t in source_list.tasks if t.id]
        except Exception as e:
            logger.warning(f"Could not retrieve source tasks: {e}")
            source_tasks = []

        try:
            dest_list = self._dest_lab_manager.list_tasks(FilterSettings(count=1000))
            dest_tasks = [{"id": t.id, "name": t.name} for t in dest_list.tasks if t.id]
        except Exception as e:
            logger.warning(f"Could not retrieve destination tasks: {e}")
            dest_tasks = []

        return self._compare_resources(source_tasks, dest_tasks, "id", "name", "tasks")

    def _compare_artifact_types(self) -> ResourceDiff:
        try:
            source_list = self._source_lab_manager.list_artifact_types(FilterSettings(count=1000))
            source_artifacts = [{"id": at.id, "name": at.name} for at in source_list.artifact_types if at.id]
        except Exception as e:
            logger.warning(f"Could not retrieve source artifact types: {e}")
            source_artifacts = []

        try:
            dest_list = self._dest_lab_manager.list_artifact_types(FilterSettings(count=1000))
            dest_artifacts = [{"id": at.id, "name": at.name} for at in dest_list.artifact_types if at.id]
        except Exception as e:
            logger.warning(f"Could not retrieve destination artifact types: {e}")
            dest_artifacts = []

        return self._compare_resources(source_artifacts, dest_artifacts, "id", "name", "artifact_types")

    def _compare_rag_assistants(self) -> ResourceDiff:
        try:
            source_client = RAGAssistantClient(api_key=self.from_api_key, base_url=self.from_instance)
            source_data = source_client.get_assistants_from_project()
            source_raw = source_data.get("assistants", [])
            source_assistants = [
                {"name": RAGAssistantMapper.map_to_rag_assistant(a).name}
                for a in source_raw if a
            ]
        except Exception as e:
            logger.warning(f"Could not retrieve source RAG assistants: {e}")
            source_assistants = []

        try:
            dest_client = RAGAssistantClient(api_key=self.to_api_key, base_url=self.to_instance)
            dest_data = dest_client.get_assistants_from_project()
            dest_raw = dest_data.get("assistants", [])
            dest_assistants = [
                {"name": RAGAssistantMapper.map_to_rag_assistant(a).name}
                for a in dest_raw if a
            ]
        except Exception as e:
            logger.warning(f"Could not retrieve destination RAG assistants: {e}")
            dest_assistants = []

        return self._compare_resources(source_assistants, dest_assistants, "name", "name", "rag_assistants")

    def _compare_files(self) -> ResourceDiff:
        if not self.from_organization_id or not self.to_organization_id:
            logger.warning("Organization IDs required for file comparison, skipping")
            return ResourceDiff(resource_type="files")

        try:
            source_manager = FileManager(
                api_key=self.from_api_key,
                base_url=self.from_instance,
                organization_id=self.from_organization_id,
                project_id=self.from_project_id
            )
            source_list = source_manager.get_file_list()
            source_files = [{"id": f.id, "name": f.filename} for f in source_list.files if f.id]
        except Exception as e:
            logger.warning(f"Could not retrieve source files: {e}")
            source_files = []

        try:
            dest_manager = FileManager(
                api_key=self.to_api_key,
                base_url=self.to_instance,
                organization_id=self.to_organization_id,
                project_id=self.to_project_id
            )
            dest_list = dest_manager.get_file_list()
            dest_files = [{"id": f.id, "name": f.filename} for f in dest_list.files if f.id]
        except Exception as e:
            logger.warning(f"Could not retrieve destination files: {e}")
            dest_files = []

        return self._compare_resources(source_files, dest_files, "id", "name", "files")

    def _compare_usage_limits(self) -> ResourceDiff:
        if not self.from_organization_api_key or not self.to_organization_api_key:
            logger.warning("Organization API keys required for usage limits comparison, skipping")
            return ResourceDiff(resource_type="usage_limits")

        if not self.from_organization_id or not self.to_organization_id:
            logger.warning("Organization IDs required for usage limits comparison, skipping")
            return ResourceDiff(resource_type="usage_limits")

        try:
            source_manager = UsageLimitManager(
                api_key=self.from_organization_api_key,
                base_url=self.from_instance
            )
            source_limits_response = source_manager.get_all_usage_limits_from_project(self.from_organization_id)
            source_limits = [
                {"model": ul.model, "limit": ul.limit}
                for ul in source_limits_response.usage_limits if ul.model
            ] if hasattr(source_limits_response, 'usage_limits') else []
        except Exception as e:
            logger.warning(f"Could not retrieve source usage limits: {e}")
            source_limits = []

        try:
            dest_manager = UsageLimitManager(
                api_key=self.to_organization_api_key,
                base_url=self.to_instance
            )
            dest_limits_response = dest_manager.get_latest_usage_limit_from_project(self.to_organization_id)
            dest_limits = [
                {"model": ul.model, "limit": ul.limit}
                for ul in dest_limits_response.usage_limits if ul.model
            ] if hasattr(dest_limits_response, 'usage_limits') else []
        except Exception as e:
            logger.warning(f"Could not retrieve destination usage limits: {e}")
            dest_limits = []

        return self._compare_resources(source_limits, dest_limits, "model", "model", "usage_limits")

    def _compare_secrets(self) -> ResourceDiff:
        try:
            source_client = SecretClient(api_key=self.from_api_key, base_url=self.from_instance)
            source_data = source_client.list_secrets(count=1000)
            source_secrets_list = source_data.get("secrets", []) if isinstance(source_data, dict) else []
            source_secrets = [
                {"id": s.get("id"), "name": s.get("name")}
                for s in source_secrets_list if s.get("id")
            ]
        except Exception as e:
            logger.warning(f"Could not retrieve source secrets: {e}")
            source_secrets = []

        try:
            dest_client = SecretClient(api_key=self.to_api_key, base_url=self.to_instance)
            dest_data = dest_client.list_secrets(count=1000)
            dest_secrets_list = dest_data.get("secrets", []) if isinstance(dest_data, dict) else []
            dest_secrets = [
                {"id": s.get("id"), "name": s.get("name")}
                for s in dest_secrets_list if s.get("id")
            ]
        except Exception as e:
            logger.warning(f"Could not retrieve destination secrets: {e}")
            dest_secrets = []

        return self._compare_resources(source_secrets, dest_secrets, "id", "name", "secrets")

    def _compare_resources(
        self,
        source_items: List[Dict[str, Any]],
        dest_items: List[Dict[str, Any]],
        id_field: str,
        name_field: str,
        resource_type: str
    ) -> ResourceDiff:
        source_ids = {item[id_field] for item in source_items if item.get(id_field)}
        dest_ids = {item[id_field] for item in dest_items if item.get(id_field)}

        only_in_source = [
            item for item in source_items
            if item.get(id_field) and item[id_field] not in dest_ids
        ]
        only_in_dest = [
            item for item in dest_items
            if item.get(id_field) and item[id_field] not in source_ids
        ]
        in_both = [
            item for item in source_items
            if item.get(id_field) and item[id_field] in dest_ids
        ]

        return ResourceDiff(
            resource_type=resource_type,
            only_in_source=only_in_source,
            only_in_destination=only_in_dest,
            in_both=in_both,
            total_source=len(source_items),
            total_destination=len(dest_items)
        )


class DiffFormatter:
    def format_as_text(self, diff: ProjectDiff) -> str:
        lines = []
        lines.append("=" * 80)
        lines.append("PROJECT COMPARISON REPORT")
        lines.append("=" * 80)
        lines.append(f"Source:      {diff.source_instance} / Project: {diff.source_project_id}")
        lines.append(f"Destination: {diff.destination_instance} / Project: {diff.destination_project_id}")
        lines.append(f"Timestamp:   {diff.timestamp}")
        lines.append("")
        lines.append("SUMMARY")
        lines.append("-" * 80)
        lines.append(f"Total resource types compared: {diff.summary.get('resource_types_compared', 0)}")
        lines.append(f"Total differences found: {diff.summary.get('total_differences', 0)}")
        lines.append("")

        for resource_type, resource_diff in diff.resource_diffs.items():
            lines.append("")
            lines.append(f"{resource_type.upper().replace('_', ' ')}")
            lines.append("-" * 80)
            lines.append(f"  Source: {resource_diff.total_source} {resource_type}")
            lines.append(f"  Destination: {resource_diff.total_destination} {resource_type}")
            lines.append(f"  Only in source: {len(resource_diff.only_in_source)}")
            lines.append(f"  Only in destination: {len(resource_diff.only_in_destination)}")
            lines.append(f"  In both: {len(resource_diff.in_both)}")
            lines.append("")

            if resource_diff.only_in_source:
                lines.append("  Only in source:")
                for item in resource_diff.only_in_source:
                    name = item.get("name", "N/A")
                    id_val = item.get("id", item.get("model", "N/A"))
                    if name != "N/A" and name != id_val:
                        lines.append(f"    - {name} (ID: {id_val})")
                    else:
                        lines.append(f"    - {id_val}")
                lines.append("")

            if resource_diff.only_in_destination:
                lines.append("  Only in destination:")
                for item in resource_diff.only_in_destination:
                    name = item.get("name", "N/A")
                    id_val = item.get("id", item.get("model", "N/A"))
                    if name != "N/A" and name != id_val:
                        lines.append(f"    - {name} (ID: {id_val})")
                    else:
                        lines.append(f"    - {id_val}")
                lines.append("")

        lines.append("=" * 80)
        return "\n".join(lines)

    def format_as_json(self, diff: ProjectDiff) -> str:
        data = {
            "source_project_id": diff.source_project_id,
            "destination_project_id": diff.destination_project_id,
            "source_instance": diff.source_instance,
            "destination_instance": diff.destination_instance,
            "timestamp": diff.timestamp,
            "summary": diff.summary,
            "resource_diffs": {}
        }

        for resource_type, resource_diff in diff.resource_diffs.items():
            data["resource_diffs"][resource_type] = {
                "resource_type": resource_diff.resource_type,
                "total_source": resource_diff.total_source,
                "total_destination": resource_diff.total_destination,
                "only_in_source": resource_diff.only_in_source,
                "only_in_destination": resource_diff.only_in_destination,
                "in_both": resource_diff.in_both
            }

        return json.dumps(data, indent=2)

    def format_as_markdown(self, diff: ProjectDiff) -> str:
        lines = []
        lines.append("# Project Comparison Report")
        lines.append("")
        lines.append(f"**Source:** {diff.source_instance} / Project: `{diff.source_project_id}`")
        lines.append(f"**Destination:** {diff.destination_instance} / Project: `{diff.destination_project_id}`")
        lines.append(f"**Timestamp:** {diff.timestamp}")
        lines.append("")
        lines.append("## Summary")
        lines.append("")
        lines.append(f"- **Resource types compared:** {diff.summary.get('resource_types_compared', 0)}")
        lines.append(f"- **Total differences found:** {diff.summary.get('total_differences', 0)}")
        lines.append("")

        for resource_type, resource_diff in diff.resource_diffs.items():
            lines.append(f"## {resource_type.replace('_', ' ').title()}")
            lines.append("")
            lines.append("| Metric | Count |")
            lines.append("|--------|-------|")
            lines.append(f"| Source | {resource_diff.total_source} |")
            lines.append(f"| Destination | {resource_diff.total_destination} |")
            lines.append(f"| Only in source | {len(resource_diff.only_in_source)} |")
            lines.append(f"| Only in destination | {len(resource_diff.only_in_destination)} |")
            lines.append(f"| In both | {len(resource_diff.in_both)} |")
            lines.append("")

            if resource_diff.only_in_source:
                lines.append("### Only in Source")
                lines.append("")
                for item in resource_diff.only_in_source:
                    name = item.get("name", "N/A")
                    id_val = item.get("id", item.get("model", "N/A"))
                    if name != "N/A" and name != id_val:
                        lines.append(f"- **{name}** (`{id_val}`)")
                    else:
                        lines.append(f"- `{id_val}`")
                lines.append("")

            if resource_diff.only_in_destination:
                lines.append("### Only in Destination")
                lines.append("")
                for item in resource_diff.only_in_destination:
                    name = item.get("name", "N/A")
                    id_val = item.get("id", item.get("model", "N/A"))
                    if name != "N/A" and name != id_val:
                        lines.append(f"- **{name}** (`{id_val}`)")
                    else:
                        lines.append(f"- `{id_val}`")
                lines.append("")

        return "\n".join(lines)

    def write_to_file(self, content: str, filepath: str) -> None:
        import os
        
        directory = os.path.dirname(filepath)
        if directory and not os.path.exists(directory):
            os.makedirs(directory)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)
        
        logger.info(f"Diff report written to {filepath}")
