from __future__ import annotations

from agents.venture_discovery import VentureDiscoveryService
from control_plane.ventures.models import CompanyMandate, CompanyProposal, ICDiligence, ICDecision, PortfolioCapabilityGap
from graph.native_runtime import GraphExecutionContext
from graph.registry import NodeHandlerRegistry, NodeResult
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec


def venture_discovery_graph_v1() -> ExecutionGraphSpec:
    nodes = ["prepare_mandate", "generate_company", "conduct_research", "board_review", "revise_pitch", "ic_first_pass", "generate_questions", "company_response", "red_team", "final_response", "score", "ic_decision", "capability_analysis", "produce_investment_memo", "complete"]
    spec = ExecutionGraphSpec(
        name="venture_discovery",
        version=1,
        graph_type="VENTURE_DISCOVERY",
        entry="prepare_mandate",
        nodes={node: GraphNodeSpec(node, node if node == "complete" else f"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": "Venture Discovery V0: exactly one startup proposal through Board, IC diligence, scoring, memo, and capability-gap analysis."},
    )
    spec.validate()
    return spec


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

    def __init__(self, service: VentureDiscoveryService, node_type: str) -> None:
        self.service = service
        self.node_type = node_type

    def mandate(self, context: GraphExecutionContext) -> CompanyMandate:
        return CompanyMandate.objects.get(id=context.graph_run.metadata["mandate_id"])

    def proposal(self, context: GraphExecutionContext) -> CompanyProposal:
        return CompanyProposal.objects.get(id=context.graph_run.metadata["proposal_id"])

    def diligence(self, context: GraphExecutionContext) -> ICDiligence:
        return ICDiligence.objects.get(id=context.graph_run.metadata["diligence_id"])

    def gap(self, context: GraphExecutionContext) -> PortfolioCapabilityGap:
        return PortfolioCapabilityGap.objects.get(id=context.graph_run.metadata["capability_gap_id"])


class PrepareMandateNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        mandate = self.service.create_v0_mandate()
        metadata = dict(context.graph_run.metadata)
        metadata["mandate_id"] = str(mandate.id)
        metadata["exactly_one_company_required"] = True
        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", {"mandate_id": str(mandate.id), "objective": mandate.objective})


class GenerateCompanyNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        proposal = self.service.generate_single_company(self.mandate(context), graph_run=context.graph_run)
        metadata = dict(context.graph_run.metadata)
        metadata["proposal_id"] = str(proposal.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"proposal_id": str(proposal.id), "company": proposal.title, "generation_source": proposal.metadata.get("generation_source")})


class BoardReviewNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        review = self.service.board_review(self.proposal(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"board_review_id": str(review.id), "recommendation": review.recommendation})


class ResearchNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        research = self.service.conduct_market_research(self.proposal(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"source_count": len(research.get("sources", [])), "unverified_categories": research.get("unverified_categories", [])})


class RevisePitchNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        proposal = self.proposal(context)
        latest = proposal.board_reviews.order_by("-created_at").first()
        if latest and latest.revised_pitch:
            proposal.pitch = latest.revised_pitch
            proposal.save(update_fields=["pitch", "updated_at"])
        return NodeResult("COMPLETE", "success", {"proposal_id": str(proposal.id), "pitch_revised": latest is not None})


class ICFirstPassNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        diligence = self.service.start_ic_diligence(self.proposal(context), graph_run=context.graph_run)
        metadata = dict(context.graph_run.metadata)
        metadata["diligence_id"] = str(diligence.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"diligence_id": str(diligence.id), "status": diligence.status})


class ICQuestionsNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        questions = self.service.generate_ic_questions(self.diligence(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"question_count": len(questions), "question_ids": [str(q.id) for q in questions]})


class CompanyResponseNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        responses = self.service.answer_questions(self.diligence(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"response_count": len(responses)})


class RedTeamNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        challenge = self.service.red_team(self.diligence(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"concern_count": len(challenge.get("concerns", []))})


class FinalResponseNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        response = self.service.final_company_response(self.diligence(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"kill_criteria_count": len(response.get("kill_criteria", []))})


class ScoreNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        decision = self.service.score_and_decide(self.diligence(context), graph_run=context.graph_run)
        metadata = dict(context.graph_run.metadata)
        metadata["decision_id"] = str(decision.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"decision_id": str(decision.id), "score": decision.composite_score})


class ICDecisionNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        decision = ICDecision.objects.get(id=context.graph_run.metadata["decision_id"])
        return NodeResult("COMPLETE", "success", {"decision": decision.decision, "initial_tranche": str(decision.initial_tranche or "")})


class CapabilityAnalysisNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        gap = self.service.capability_analysis(self.proposal(context), graph_run=context.graph_run)
        metadata = dict(context.graph_run.metadata)
        metadata["capability_gap_id"] = str(gap.id)
        context.graph_run.metadata = metadata
        context.graph_run.save(update_fields=["metadata", "updated_at"])
        return NodeResult("COMPLETE", "success", {"capability_gap_id": str(gap.id), "missing_count": len(gap.missing)})


class InvestmentMemoNode(VentureNode):
    def run(self, context: GraphExecutionContext) -> NodeResult:
        memo = self.service.produce_investment_memo(self.diligence(context), self.gap(context), graph_run=context.graph_run)
        return NodeResult("COMPLETE", "success", {"memo_artifact_id": str(memo.id), "artifact_type": memo.artifact_type})


def venture_discovery_registry(service: VentureDiscoveryService) -> NodeHandlerRegistry:
    registry = NodeHandlerRegistry()
    for handler in [PrepareMandateNode(service, "venture_prepare_mandate"), GenerateCompanyNode(service, "venture_generate_company"), ResearchNode(service, "venture_conduct_research"), BoardReviewNode(service, "venture_board_review"), RevisePitchNode(service, "venture_revise_pitch"), ICFirstPassNode(service, "venture_ic_first_pass"), ICQuestionsNode(service, "venture_generate_questions"), CompanyResponseNode(service, "venture_company_response"), RedTeamNode(service, "venture_red_team"), FinalResponseNode(service, "venture_final_response"), ScoreNode(service, "venture_score"), ICDecisionNode(service, "venture_ic_decision"), CapabilityAnalysisNode(service, "venture_capability_analysis"), InvestmentMemoNode(service, "venture_produce_investment_memo")]:
        registry.register(handler)
    return registry