from __future__ import annotations

from agents.venture_discovery import VentureDiscoveryService
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 venture_discovery_cohort_graph_v1() -> ExecutionGraphSpec:
    nodes = ["prepare_cohort", "generate_independent_proposals", "initial_research", "fingerprint_theses", "collision_analysis", "run_individual_diligence", "portfolio_compare", "portfolio_ic", "aggregate_capabilities", "produce_cohort_report", "complete"]
    spec = ExecutionGraphSpec(name="venture_discovery_cohort", version=1, graph_type="VENTURE_DISCOVERY_COHORT", entry="prepare_cohort", nodes={node: GraphNodeSpec(node, node if node == "complete" else f"venture_cohort_{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": "Venture Discovery V0.2: 10-company independent cohort, portfolio IC, collisions, capability demand, and report."})
    spec.validate()
    return spec


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

    def __init__(self, service: VentureDiscoveryService, node_type: str, *, cohort_size: int = 10, concurrency: int = 1) -> 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 PrepareCohortNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.service.prepare_cohort(size=self.cohort_size, graph_run=context.graph_run, concurrency=self.concurrency)
        metadata = dict(context.graph_run.metadata)
        metadata["cohort_id"] = str(cohort.id)
        metadata["cohort_stable_id"] = cohort.cohort_id
        metadata["no_real_spend"] = True
        metadata["no_real_customer_outreach"] = True
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"cohort_id": cohort.cohort_id, "size": cohort.cohort_size})


class GenerateIndependentNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        proposals = self.service.generate_independent_proposals(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"proposal_count": len(proposals), "independent_generation": True})


class ResearchNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.cohort(context)
        self.service.research_cohort(cohort)
        cohort.refresh_from_db()
        return NodeResult("COMPLETE", "success", cohort.metrics)


class FingerprintNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        fingerprints = self.service.fingerprint_cohort(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"fingerprint_count": len(fingerprints)})


class CollisionNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        collisions = self.service.analyze_collisions(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"pair_count": len(collisions)})


class DiligenceNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.cohort(context)
        self.service.run_individual_diligence_for_cohort(cohort)
        return NodeResult("COMPLETE", "success", {"diligence_count": cohort.members.count()})


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


class CapabilityNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.cohort(context)
        demand = self.service.aggregate_capability_demand(cohort)
        top_3 = self.service.aggregate_capability_demand(cohort, top_3_only=True)
        return NodeResult("COMPLETE", "success", {"capability_count": len(demand), "top_3_gap_count": len(top_3)})


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


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


def venture_discovery_cohort_registry(service: VentureDiscoveryService, *, cohort_size: int = 10, concurrency: int = 1) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    for handler in [PrepareCohortNode(service, "venture_cohort_prepare_cohort", cohort_size=cohort_size, concurrency=concurrency), GenerateIndependentNode(service, "venture_cohort_generate_independent_proposals"), ResearchNode(service, "venture_cohort_initial_research"), FingerprintNode(service, "venture_cohort_fingerprint_theses"), CollisionNode(service, "venture_cohort_collision_analysis"), DiligenceNode(service, "venture_cohort_run_individual_diligence"), NoopNode(service, "venture_cohort_portfolio_compare"), PortfolioNode(service, "venture_cohort_portfolio_ic"), CapabilityNode(service, "venture_cohort_aggregate_capabilities"), ReportNode(service, "venture_cohort_produce_cohort_report")]:
        registry.register(handler)
    return registry