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", "portfolio_thesis_review", "create_ideation_mandate", "allocate_search_territories", "generate_independent_proposals", "novelty_gate", "initial_research", "fingerprint_theses", "collision_analysis", "cluster_theses", "run_individual_diligence", "autonomous_operability_review", "portfolio_compare", "portfolio_ic", "saturation_analysis", "update_thesis_registry", "aggregate_capabilities", "produce_cohort_report", "complete"]
    spec = ExecutionGraphSpec(name="venture_discovery_cohort", version=4, 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.5: autonomous operability track, research smoke hardening, finalist deep research, and autonomous-only finalists."})
    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 PortfolioThesisReviewNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        review = self.service.portfolio_thesis_review(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"registry_size": review["registry_size"], "saturated_count": len(review["saturated_thesis_areas"])})


class IdeationMandateNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        mandate = self.service.create_ideation_mandate(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"mandate_id": str(mandate.id), "territory_count": len(mandate.opportunity_territories)})


class SearchTerritoryNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        territories = self.service.allocate_search_territories(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"territories": territories})


class NoveltyGateNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        cohort = self.cohort(context)
        return NodeResult("COMPLETE", "success", {"accepted_proposals": cohort.metrics.get("accepted_proposals", cohort.members.count()), "rejections": cohort.generation_rejections.count(), "failed_slots": cohort.metrics.get("failed_slots", 0)})


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 ClusterThesesNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        clusters = self.service.cluster_theses(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"cluster_count": len(clusters)})


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 AutonomousOperabilityNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        assessments = self.service.assess_autonomous_operability_for_cohort(self.cohort(context))
        eligible = len([item for item in assessments if item.gate_result == "AUTONOMOUS_ELIGIBLE"])
        return NodeResult("COMPLETE", "success", {"assessment_count": len(assessments), "autonomous_eligible_count": eligible})


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 SaturationNode(CohortNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        rows = self.service.saturation_analysis(self.cohort(context))
        return NodeResult("COMPLETE", "success", {"analysis_count": len(rows)})


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


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), PortfolioThesisReviewNode(service, "venture_cohort_portfolio_thesis_review"), IdeationMandateNode(service, "venture_cohort_create_ideation_mandate"), SearchTerritoryNode(service, "venture_cohort_allocate_search_territories"), GenerateIndependentNode(service, "venture_cohort_generate_independent_proposals"), NoveltyGateNode(service, "venture_cohort_novelty_gate"), ResearchNode(service, "venture_cohort_initial_research"), FingerprintNode(service, "venture_cohort_fingerprint_theses"), CollisionNode(service, "venture_cohort_collision_analysis"), ClusterThesesNode(service, "venture_cohort_cluster_theses"), DiligenceNode(service, "venture_cohort_run_individual_diligence"), AutonomousOperabilityNode(service, "venture_cohort_autonomous_operability_review"), NoopNode(service, "venture_cohort_portfolio_compare"), PortfolioNode(service, "venture_cohort_portfolio_ic"), SaturationNode(service, "venture_cohort_saturation_analysis"), RegistryUpdateNode(service, "venture_cohort_update_thesis_registry"), CapabilityNode(service, "venture_cohort_aggregate_capabilities"), ReportNode(service, "venture_cohort_produce_cohort_report")]:
        registry.register(handler)
    return registry