from __future__ import annotations

import json
import os

from django.core.management.base import BaseCommand, CommandError

from agents.venture_discovery import VentureDiscoveryService
from control_plane.ventures.models import VentureCohort
from graph.bootstrap import champion_venture_discovery_cohort_graph_v1
from graph.langgraph_runtime import LangGraphRuntime
from graph.models import GraphRun, GraphRunStatus
from graph.venture_cohort import venture_discovery_cohort_registry
from model_router.providers import providers_from_resources
from model_router.router import ModelRouter


class Command(BaseCommand):
    help = "Run a Venture Discovery cohort through the champion graph."

    def add_arguments(self, parser):
        parser.add_argument("--size", type=int, default=10)
        parser.add_argument("--concurrency", type=int, default=2)
        parser.add_argument("--qwen-only", action="store_true", help="Route Venture ideation, research, and portfolio IC roles to Qwen.")
        parser.add_argument("--no-web-research", action="store_true", help="Disable SearXNG/page-fetch research for this run.")
        parser.add_argument("--persist-requests", action="store_true", help="Persist sanitized model request telemetry.")
        parser.add_argument("--indent", type=int, default=2, help="JSON indentation for the summary output. Use 0 for compact JSON.")

    def handle(self, *args, **options):
        size = max(1, int(options["size"]))
        concurrency = max(1, int(options["concurrency"]))
        model_env_keys = ["ARTIFEX_VENTURE_IDEATION_MODEL", "ARTIFEX_VENTURE_RESEARCH_MODEL", "ARTIFEX_VENTURE_PORTFOLIO_IC_MODEL"]
        previous_model_env = {key: os.environ.get(key) for key in model_env_keys}
        if options["qwen_only"]:
            os.environ["ARTIFEX_VENTURE_IDEATION_MODEL"] = "qwen"
            os.environ["ARTIFEX_VENTURE_RESEARCH_MODEL"] = "qwen"
            os.environ["ARTIFEX_VENTURE_PORTFOLIO_IC_MODEL"] = "qwen"
        try:
            providers = providers_from_resources()
            if not providers:
                raise CommandError("No model providers configured. Run seed_spark_resources first.")
            if options["qwen_only"] and "qwen" not in providers:
                raise CommandError("--qwen-only requested but no Qwen/local_inference provider is configured.")

            version = champion_venture_discovery_cohort_graph_v1()
            graph_run = GraphRun.objects.create(execution_graph_version=version, current_node=version.graph_spec["entry"])
            service = VentureDiscoveryService(
                ModelRouter(providers, persist_requests=bool(options["persist_requests"])),
                web_research_available=not bool(options["no_web_research"]),
            )
            LangGraphRuntime(venture_discovery_cohort_registry(service, cohort_size=size, concurrency=concurrency)).run_until_terminal_or_paused(graph_run)
            graph_run.refresh_from_db()
            summary = self.summary(graph_run)
            indent = None if int(options["indent"]) <= 0 else int(options["indent"])
            self.stdout.write(json.dumps(summary, indent=indent, default=str))
            if graph_run.status != GraphRunStatus.COMPLETE:
                raise CommandError(f"Venture cohort run ended with status {graph_run.status}")
        finally:
            for key, value in previous_model_env.items():
                if value is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = value

    def summary(self, graph_run: GraphRun) -> dict[str, object]:
        summary: dict[str, object] = {
            "graph_run_id": str(graph_run.id),
            "status": graph_run.status,
            "failure": graph_run.failure_reason,
            "current_node": graph_run.current_node,
            "metadata": graph_run.metadata,
        }
        cohort_id = graph_run.metadata.get("cohort_id") if isinstance(graph_run.metadata, dict) else None
        if not cohort_id:
            return summary
        cohort = VentureCohort.objects.get(id=cohort_id)
        members = list(cohort.members.select_related("proposal").order_by("rank", "created_at"))
        proposals = [member.proposal for member in members]
        fallback_count = sum(1 for proposal in proposals if proposal.metadata.get("generation_source") == "deterministic_fallback")
        report = cohort.mandate.artifacts.filter(artifact_type="VENTURE_DISCOVERY_COHORT_REPORT").order_by("-created_at").first()
        summary.update(
            {
                "cohort_id": cohort.cohort_id,
                "cohort_pk": str(cohort.id),
                "cohort_status": cohort.status,
                "size": cohort.cohort_size,
                "members": len(members),
                "concurrency": cohort.concurrency,
                "metrics": cohort.metrics,
                "fallback_count": fallback_count,
                "generation_sources": sorted({str(proposal.metadata.get("generation_source", "unknown")) for proposal in proposals}),
                "top_3": [
                    {
                        "rank": member.rank,
                        "title": member.proposal.title,
                        "score": member.portfolio_score,
                        "generation_source": member.proposal.metadata.get("generation_source"),
                        "research": member.proposal.metadata.get("research", {}),
                    }
                    for member in members
                    if member.is_top_3
                ],
                "report_artifact_id": str(report.id) if report else None,
            }
        )
        return summary