#!/usr/bin/env python
"""
Seed or rollback the Multica organization collaboration workflow and agents.

Usage:
  python scripts/seed_multica_org_agents.py --dry-run
  python scripts/seed_multica_org_agents.py --apply
  python scripts/seed_multica_org_agents.py --rollback
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))


PROJECT_ROOT = Path(__file__).resolve().parent.parent
FIXTURE_PATH = PROJECT_ROOT / "ai_platform" / "fixtures" / "multica_org_agents.json"


def load_fixture() -> Dict[str, Any]:
    with FIXTURE_PATH.open("r", encoding="utf-8") as fixture_file:
        return json.load(fixture_file)


def build_workflow_payload(fixture: Dict[str, Any]) -> Dict[str, Any]:
    workflow = deepcopy(fixture["workflow"])
    definition = workflow["definition"]
    workflow["published_definition"] = deepcopy(definition)
    workflow["published_at"] = datetime.utcnow()
    return workflow


def build_agent_payload(agent_fixture: Dict[str, Any], workflow_id: str | None) -> Dict[str, Any]:
    payload = deepcopy(agent_fixture)
    workflow_code = payload.pop("workflow_code", None)
    if workflow_code:
        payload["workflow_id"] = workflow_id
    return payload


def _find_parallel_merge(
    node_by_id: Dict[str, Dict[str, Any]],
    outgoing_edges: Dict[str, list[str]],
    branch_starts: list[str],
) -> str | None:
    branch_merges: list[str] = []
    for start in branch_starts:
        current = start
        seen: set[str] = set()
        while current and current not in seen:
            seen.add(current)
            next_nodes = outgoing_edges.get(current, [])
            if not next_nodes or len(next_nodes) > 1:
                break
            current = next_nodes[0]
            if (node_by_id.get(current) or {}).get("type") == "merge":
                branch_merges.append(current)
                break
    unique_merges = set(branch_merges)
    return branch_merges[0] if len(unique_merges) == 1 and len(branch_merges) == len(branch_starts) else None


def _walk_workflow_nodes(
    node_by_id: Dict[str, Dict[str, Any]],
    outgoing_edges: Dict[str, list[str]],
    start_id: str,
) -> set[str]:
    visited: set[str] = set()
    current = start_id
    while current and current not in visited:
        visited.add(current)
        node_type = (node_by_id.get(current) or {}).get("type")
        if node_type == "end":
            break
        if node_type == "parallel":
            branch_starts = outgoing_edges.get(current, [])
            for branch_start in branch_starts:
                branch_current = branch_start
                branch_seen: set[str] = set()
                while branch_current and branch_current not in branch_seen:
                    if (node_by_id.get(branch_current) or {}).get("type") == "merge":
                        break
                    visited.add(branch_current)
                    branch_seen.add(branch_current)
                    branch_next_nodes = outgoing_edges.get(branch_current, [])
                    branch_current = branch_next_nodes[0] if branch_next_nodes else None
            current = _find_parallel_merge(node_by_id, outgoing_edges, branch_starts)
            continue
        next_nodes = outgoing_edges.get(current, [])
        current = next_nodes[0] if next_nodes else None
    return visited


def validate_fixture(fixture: Dict[str, Any]) -> Dict[str, int]:
    agents = fixture.get("agents", [])
    definition = fixture.get("workflow", {}).get("definition", {})
    nodes = definition.get("nodes", [])
    edges = definition.get("edges", [])
    agent_codes = {agent.get("code") for agent in agents}
    workflow_agent_codes = {
        (node.get("data") or {}).get("agent_code")
        for node in nodes
        if isinstance(node, dict) and (node.get("data") or {}).get("agent_code")
    }
    node_by_id = {node.get("id"): node for node in nodes if isinstance(node, dict)}
    node_types = {node.get("type") for node in nodes}
    outgoing_edges: Dict[str, list[str]] = {}
    for edge in edges:
        if not isinstance(edge, dict):
            continue
        source = edge.get("source")
        target = edge.get("target")
        if source and target:
            outgoing_edges.setdefault(source, []).append(target)

    required_persona_fields = {
        "role",
        "skills",
        "constraints",
        "background",
        "examples",
    }
    required_node_types = {"start", "end", "template", "parallel", "merge"}
    multi_outgoing_node_types = {"parallel", "condition", "intent", "choice"}
    start_node = next((node for node in nodes if node.get("type") == "start"), None)

    missing_workflow_agents = sorted(workflow_agent_codes - agent_codes)
    missing_persona_fields = {
        agent.get("code"): sorted(required_persona_fields - set((agent.get("persona") or {}).keys()))
        for agent in agents
        if required_persona_fields - set((agent.get("persona") or {}).keys())
    }
    missing_node_types = sorted(required_node_types - node_types)
    invalid_multi_outgoing = {
        source: targets
        for source, targets in outgoing_edges.items()
        if len(targets) > 1
        and (node_by_id.get(source) or {}).get("type") not in multi_outgoing_node_types
    }
    reachable_nodes = _walk_workflow_nodes(node_by_id, outgoing_edges, start_node.get("id")) if start_node else set()
    unreachable_nodes = sorted(set(node_by_id) - reachable_nodes)
    invalid_parallel_merges = {
        source: targets
        for source, targets in outgoing_edges.items()
        if (node_by_id.get(source) or {}).get("type") == "parallel"
        and not _find_parallel_merge(node_by_id, outgoing_edges, targets)
    }
    project_manager = next(
        (agent for agent in agents if agent.get("code") == "project_manager"),
        None,
    )
    project_manager_workflow = (project_manager or {}).get("workflow_code")
    if (
        missing_workflow_agents
        or missing_persona_fields
        or missing_node_types
        or invalid_multi_outgoing
        or unreachable_nodes
        or invalid_parallel_merges
        or project_manager_workflow != "multica_org_collaboration_flow"
    ):
        raise ValueError(
            f"Fixture validation failed: missing_workflow_agents={missing_workflow_agents}, "
            f"missing_persona_fields={missing_persona_fields}, "
            f"missing_node_types={missing_node_types}, "
            f"invalid_multi_outgoing={invalid_multi_outgoing}, "
            f"unreachable_nodes={unreachable_nodes}, "
            f"invalid_parallel_merges={invalid_parallel_merges}, "
            f"project_manager_workflow={project_manager_workflow}"
        )

    return {
        "agents": len(agents),
        "workflow_nodes": len(nodes),
        "workflow_edges": len(edges),
        "workflow_agent_refs": len(workflow_agent_codes),
    }


async def resolve_default_chat_model_id(session) -> str | None:
    from sqlalchemy import select

    from ai_platform.models.model import LLMModel

    result = await session.execute(
        select(LLMModel)
        .where(
            LLMModel.is_deleted == False,
            LLMModel.is_active == True,
            LLMModel.model_type == "chat",
        )
        .order_by(LLMModel.sort.desc(), LLMModel.sys_create_datetime.desc())
    )
    model = result.scalars().first()
    return str(model.id) if model else None


async def upsert_workflow(session, workflow_payload: Dict[str, Any]) -> AIWorkflow:
    from sqlalchemy import select

    from ai_platform.models.workflow import AIWorkflow, AIWorkflowVersion

    code = workflow_payload["code"]
    result = await session.execute(select(AIWorkflow).where(AIWorkflow.code == code))
    workflow = result.scalar_one_or_none()

    if workflow:
        for key, value in workflow_payload.items():
            setattr(workflow, key, value)
        workflow.is_deleted = False
    else:
        workflow = AIWorkflow(**workflow_payload)
        session.add(workflow)

    await session.flush()

    result = await session.execute(
        select(AIWorkflowVersion).where(
            AIWorkflowVersion.workflow_id == workflow.id,
            AIWorkflowVersion.version == workflow.version,
        )
    )
    version = result.scalar_one_or_none()
    version_payload = {
        "workflow_id": workflow.id,
        "version": workflow.version,
        "definition": deepcopy(workflow.definition),
        "description": "Seeded Multica organization collaboration workflow",
        "published_at": workflow.published_at,
    }
    if version:
        for key, value in version_payload.items():
            setattr(version, key, value)
        version.is_deleted = False
    else:
        session.add(AIWorkflowVersion(**version_payload))

    return workflow


async def upsert_agents(session, fixture: Dict[str, Any], workflow_id: str) -> int:
    from sqlalchemy import select

    from ai_platform.models.agent import Agent
    from app.base_model import generate_nanoid

    count = 0
    default_chat_model_id = await resolve_default_chat_model_id(session)
    for agent_fixture in fixture["agents"]:
        payload = build_agent_payload(agent_fixture, workflow_id)
        if (
            payload.get("mode", "autonomous") == "autonomous"
            and not payload.get("model_id")
            and default_chat_model_id
        ):
            payload["model_id"] = default_chat_model_id
        result = await session.execute(select(Agent).where(Agent.code == payload["code"]))
        agent = result.scalar_one_or_none()
        if agent:
            for key, value in payload.items():
                setattr(agent, key, value)
            agent.is_deleted = False
        else:
            session.add(Agent(id=generate_nanoid(), **payload))
        count += 1
    return count


async def apply_seed(dry_run: bool) -> Dict[str, Any]:
    fixture = load_fixture()
    counts = validate_fixture(fixture)
    workflow_payload = build_workflow_payload(fixture)

    if dry_run:
        return {
            "action": "dry-run",
            "workflow_code": workflow_payload["code"],
            "agent_count": len(fixture["agents"]),
            **counts,
        }

    from app.database import AsyncSessionLocal

    async with AsyncSessionLocal() as session:
        workflow = await upsert_workflow(session, workflow_payload)
        agent_count = await upsert_agents(session, fixture, workflow.id)

        await session.commit()
        action = "applied"

    return {
        "action": action,
        "workflow_code": workflow_payload["code"],
        "agent_count": agent_count,
        **counts,
    }


async def rollback_seed(dry_run: bool) -> Dict[str, Any]:
    from sqlalchemy import select

    from ai_platform.models.agent import Agent
    from ai_platform.models.workflow import AIWorkflow, AIWorkflowVersion
    from app.database import AsyncSessionLocal

    fixture = load_fixture()
    workflow_code = fixture["workflow"]["code"]
    agent_codes = [agent["code"] for agent in fixture["agents"]]

    async with AsyncSessionLocal() as session:
        result = await session.execute(select(AIWorkflow).where(AIWorkflow.code == workflow_code))
        workflow = result.scalar_one_or_none()
        workflow_count = 0
        version_count = 0
        if workflow:
            workflow.is_deleted = True
            workflow_count = 1
            versions = await session.execute(
                select(AIWorkflowVersion).where(AIWorkflowVersion.workflow_id == workflow.id)
            )
            for version in versions.scalars().all():
                version.is_deleted = True
                version_count += 1

        agents = await session.execute(select(Agent).where(Agent.code.in_(agent_codes)))
        agent_count = 0
        for agent in agents.scalars().all():
            agent.is_deleted = True
            agent.workflow_id = None
            agent_count += 1

        if dry_run:
            await session.rollback()
            action = "rollback-dry-run"
        else:
            await session.commit()
            action = "rolled-back"

    return {
        "action": action,
        "workflow_code": workflow_code,
        "workflow_count": workflow_count,
        "workflow_version_count": version_count,
        "agent_count": agent_count,
    }


async def main() -> None:
    parser = argparse.ArgumentParser(description="Seed Multica organization collaboration agents")
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--dry-run", action="store_true", help="Validate and simulate the seed without committing")
    mode.add_argument("--apply", action="store_true", help="Apply the seed to the configured database")
    mode.add_argument("--rollback", action="store_true", help="Soft-delete seeded workflow, versions, and agents")
    args = parser.parse_args()

    if args.rollback:
        result = await rollback_seed(dry_run=False)
    else:
        result = await apply_seed(dry_run=args.dry_run)

    print(json.dumps(result, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    asyncio.run(main())