from __future__ import annotations

from agents.crypto_venture import CryptoVentureService
from control_plane.ventures.models import VentureCohort
from graph.native_runtime import GraphExecutionContext
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec


def crypto_venture_cohort_graph_v1() -> ExecutionGraphSpec:
    nodes = [
        "prepare_crypto_mandate",
        "portfolio_crypto_thesis_review",
        "generate_independent_protocols",
        "onchain_necessity_gate",
        "token_necessity_gate",
        "novelty_gate",
        "regenerate_rejected_slots",
        "light_market_research",
        "protocol_research",
        "protocol_security_gate",
        "token_red_team",
        "tokenomics_simulation",
        "crypto_ic_first_pass",
        "top5_deep_research",
        "rescore",
        "portfolio_crypto_ic",
        "regulatory_gate",
        "capability_analysis",
        "update_crypto_thesis_registry",
        "produce_crypto_cohort_report",
        "complete",
    ]
    spec = ExecutionGraphSpec(name="crypto_venture_cohort", version=4, graph_type="CRYPTO_VENTURE_COHORT", entry="prepare_crypto_mandate", nodes={node: GraphNodeSpec(node, node if node == "complete" else f"crypto_venture_{node}") for node in nodes}, edges=[GraphEdgeSpec(nodes[index], nodes[index + 1], "success") for index in range(len(nodes) - 1)], terminal_nodes=["complete"], metadata={"description": "Crypto / Protocol Venture Cohort V0.3.1: Sol-native problem-first ideation, repaired stage ordering, structural native-token utility judge, primary-source-first research."})
    spec.validate()
    return spec


class CryptoCohortNode:
    idempotent = True
    replay_safe = True
    destructive = False

    def __init__(self, service: CryptoVentureService, node_type: str, *, cohort_size: int = 10, concurrency: int = 2) -> None:
        self.service = service
        self.node_type = node_type
        self.cohort_size = cohort_size
        self.concurrency = concurrency

    def cohort(self, context: GraphExecutionContext) -> VentureCohort:
        return VentureCohort.objects.get(id=context.graph_run.metadata["cohort_id"])


class PrepareCryptoMandateNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.service.prepare_cohort(size=self.cohort_size, graph_run=context.graph_run, concurrency=self.concurrency)
        context.graph_run.metadata = {**context.graph_run.metadata, "cohort_id": str(cohort.id), "cohort_stable_id": cohort.cohort_id, "no_real_spend": True, "no_user_contact": True, "no_token_sale": True, "no_fundraising": True, "no_mainnet_issuance": True}
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"cohort_id": cohort.cohort_id, "size": cohort.cohort_size})


class GenerateProtocolsNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        proposals = self.service.generate_protocols(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"proposal_count": len(proposals)})


class TokenNecessityNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        assessments = self.service.token_necessity_gate(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"assessment_count": len(assessments)})


class OnchainNecessityNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        assessments = self.service.onchain_necessity_gate(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"assessment_count": len(assessments)})


class NoveltyGateNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        return NodeResult("COMPLETE", "success", self.service.novelty_gate(self.cohort(context)))


class RegenerateRejectedSlotsNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        created = self.service.regenerate_rejected_slots(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"regenerated": len(created)})


class LightResearchNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.cohort(context)
        self.service.light_market_research(cohort)
        return NodeResult("COMPLETE", "success", cohort.metrics)


class ProtocolResearchNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        self.service.protocol_research(self.cohort(context))
        return NodeResult("COMPLETE", "success")


class ProtocolSecurityGateNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        return NodeResult("COMPLETE", "success", self.service.protocol_security_gate(self.cohort(context)))


class TokenRedTeamNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        assessments = self.service.token_red_team(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"red_team_count": len(assessments)})


class TokenomicsSimulationNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        simulations = self.service.tokenomics_simulation(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"simulation_count": len(simulations)})


class CryptoICFirstPassNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        self.service.crypto_ic_first_pass(self.cohort(context))
        return NodeResult("COMPLETE", "success")


class Top5DeepResearchNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        self.service.top5_deep_research(self.cohort(context))
        return NodeResult("COMPLETE", "success")


class PortfolioCryptoICNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        review = self.service.portfolio_crypto_ic(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"ranked_count": len(review.rankings), "top_3": review.top_3})


class RegulatoryGateNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        self.service.regulatory_gate(self.cohort(context))
        return NodeResult("COMPLETE", "success")


class CapabilityAnalysisNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        gaps = self.service.capability_analysis(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"capability_gap_count": len(gaps)})


class RegistryUpdateNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        updates = self.service.update_crypto_thesis_registry(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"update_count": len(updates)})


class ReportNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        artifact = self.service.produce_crypto_cohort_report(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"artifact_id": str(artifact.id), "artifact_type": artifact.artifact_type})


class NoopNode(CryptoCohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        return NodeResult("COMPLETE", "success")


def crypto_venture_cohort_registry(service: CryptoVentureService, *, cohort_size: int = 10, concurrency: int = 2) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    handlers = [
        PrepareCryptoMandateNode(service, "crypto_venture_prepare_crypto_mandate", cohort_size=cohort_size, concurrency=concurrency),
        NoopNode(service, "crypto_venture_portfolio_crypto_thesis_review"),
        GenerateProtocolsNode(service, "crypto_venture_generate_independent_protocols"),
        OnchainNecessityNode(service, "crypto_venture_onchain_necessity_gate"),
        TokenNecessityNode(service, "crypto_venture_token_necessity_gate"),
        NoveltyGateNode(service, "crypto_venture_novelty_gate"),
        RegenerateRejectedSlotsNode(service, "crypto_venture_regenerate_rejected_slots"),
        LightResearchNode(service, "crypto_venture_light_market_research"),
        ProtocolResearchNode(service, "crypto_venture_protocol_research"),
        ProtocolSecurityGateNode(service, "crypto_venture_protocol_security_gate"),
        TokenRedTeamNode(service, "crypto_venture_token_red_team"),
        TokenomicsSimulationNode(service, "crypto_venture_tokenomics_simulation"),
        CryptoICFirstPassNode(service, "crypto_venture_crypto_ic_first_pass"),
        Top5DeepResearchNode(service, "crypto_venture_top5_deep_research"),
        CryptoICFirstPassNode(service, "crypto_venture_rescore"),
        PortfolioCryptoICNode(service, "crypto_venture_portfolio_crypto_ic"),
        RegulatoryGateNode(service, "crypto_venture_regulatory_gate"),
        CapabilityAnalysisNode(service, "crypto_venture_capability_analysis"),
        RegistryUpdateNode(service, "crypto_venture_update_crypto_thesis_registry"),
        ReportNode(service, "crypto_venture_produce_crypto_cohort_report"),
    ]
    for handler in handlers:
        registry.register(handler)
    return registry