"""Finalize structured and community analysis for docsctx."""

from __future__ import annotations

from typing import cast

from docsctx.analysis import analyze_document, build_communities
from docsctx.config import DocsSettings
from docsctx.db.repository import DocsRepository
from docsctx.embedder import validate_embedder_access, decode_vector, encode_vector, embed_texts
from docsctx.provider import analyze_with_retries, create_provider, validate_provider_access
import typer
from docsctx.extractor import extract_entities, extract_claims, summarize_community

try:
    from docsctx.vectorstore import DocsVectorStore
except ImportError:
    DocsVectorStore = None  # type: ignore[misc,assignment]


def run_finalize(
    repository: DocsRepository,
    settings: DocsSettings | None = None,
    embedder: object | None = None,
    vectorstore: object | None = None,
) -> dict[str, int]:
    """Build document structures, entities, relationships, and communities."""
    resolved_settings = settings or DocsSettings()
    requested_provider = resolved_settings.chat.provider.lower().strip() or "none"
    provider = validate_provider_access(resolved_settings, create_provider(resolved_settings))
    embedding_provider = resolved_settings.embedding.provider.lower().strip()
    if embedding_provider not in {"", "none"} or requested_provider not in {"none", "mock"}:
        validate_embedder_access(resolved_settings)

    documents = repository.list_documents_for_analysis()
    analyzed = []
    provider_successes = 0
    provider_fallbacks = 0
    provider_deterministic = 0
    for document in documents:
        document_id = cast(int, document["id"])
        title = cast(str, document["title"])
        source_uri = cast(str, document["source_uri"])
        doc_type = cast(str, document["doc_type"])
        content = cast(str, document["content"])
        
        doc_data = repository.get_document(document_id)
        chunks = doc_data.get("chunks", []) if doc_data else []
        chunk_texts = [c["content"] for c in chunks]

        all_entities = []
        all_relationships = []
        all_claims = []
        
        batch_size = 5
        total_batches = (len(chunk_texts) + batch_size - 1) // batch_size
        
        if total_batches > 0:
            typer.echo(f"  🔍 Extracting graph data from '{title}' ({len(chunk_texts)} chunks in {total_batches} batches)...")

        batch_idx = 0
        for i in range(0, len(chunk_texts), batch_size):
            batch_idx += 1
            batch = chunk_texts[i:i + batch_size]
            if not batch:
                continue
            
            typer.echo(f"    ⏳ Processing batch {batch_idx}/{total_batches} with provider '{requested_provider}'...", err=True)
            
            ext_res = extract_entities(resolved_settings, batch)
            all_entities.extend([(e.name, e.score) for e in ext_res.entities])
            all_relationships.extend([(r.source, r.target, r.weight) for r in ext_res.relationships])
            
            if resolved_settings.extract_claims:
                claims_res = extract_claims(resolved_settings, batch)
                all_claims.extend([(c.entity_name, c.claim, c.status) for c in claims_res])

        if total_batches > 0:
            typer.echo(f"    ✨ Extracted {len(all_entities)} entities and {len(all_relationships)} relationships total.")

        analyzed_doc = analyze_document(
            document_id=document_id,
            title=title,
            source_uri=source_uri,
            doc_type=doc_type,
            content=content,
        )

        provider_attempt = analyze_with_retries(
            provider,
            title,
            content,
            max_retries=resolved_settings.chat.max_retries,
            backoff_seconds=resolved_settings.chat.retry_backoff_seconds,
        )
        if provider_attempt.succeeded and provider_attempt.analysis is not None:
            provider_result = provider_attempt.analysis
            analyzed_doc.structure["summary"] = provider_result.summary or analyzed_doc.structure.get("summary", "")
            analyzed_doc.structure["headings"] = provider_result.headings or analyzed_doc.structure.get("headings", [])
            
            if all_entities:
                analyzed_doc.entities = all_entities
            else:
                analyzed_doc.entities = provider_result.entities or analyzed_doc.entities
                
            if all_relationships:
                analyzed_doc.relationships = all_relationships
            else:
                analyzed_doc.relationships = provider_result.relationships or analyzed_doc.relationships
                
            analyzed_doc.claims = all_claims
            
            if requested_provider == "mock":
                provider_deterministic += 1
                status_str = "deterministic"
            else:
                provider_successes += 1
                status_str = "succeeded"

            provider_meta = {
                "requested_provider": requested_provider,
                "used_provider": requested_provider,
                "status": status_str,
                "attempts": provider_attempt.attempts,
                "fallback_used": False,
                "last_error": "",
            }
        else:
            # Fallback: Use extracted entities/relationships without provider analysis
            provider_fallbacks += 1
            if all_entities:
                analyzed_doc.entities = all_entities
            if all_relationships:
                analyzed_doc.relationships = all_relationships
            analyzed_doc.claims = all_claims
            
            provider_meta = {
                "requested_provider": requested_provider,
                "used_provider": "fallback",
                "status": "fallback",
                "attempts": provider_attempt.attempts,
                "fallback_used": True,
                "last_error": provider_attempt.error,
            }
        analyzed_doc.structure["provider"] = provider_meta
        analyzed.append(analyzed_doc)
        
        # Phase 1: Embedding Chunks (if embedder provided)
        if embedder and chunks:
            typer.echo(f"    🧠 Generating embeddings for {len(chunks)} chunks...")
            chunk_embeddings = []
            vectors = embed_texts(embedder, chunk_texts)
            for idx, vec in enumerate(vectors):
                if idx < len(chunks):
                    chunk_embeddings.append((chunks[idx]["id"], encode_vector(vec)))
            repository.store_chunk_embeddings(chunk_embeddings)

            # Also store in ChromaDB for HNSW search
            if vectorstore is not None and DocsVectorStore is not None:
                chunk_ids = [chunks[idx]["id"] for idx in range(min(len(vectors), len(chunks)))]
                metadatas = [
                    {"document_id": document_id, "chunk_index": chunks[idx].get("chunk_index", idx)}
                    for idx in range(len(chunk_ids))
                ]
                vectorstore.upsert_chunk_embeddings(
                    chunk_ids=chunk_ids,
                    embeddings=vectors[:len(chunk_ids)],
                    texts=chunk_texts[:len(chunk_ids)],
                    metadatas=metadatas,
                )

    for document in analyzed:
        repository.store_document_analysis(
            document.document_id,
            structure=document.structure,
            entities=document.entities,
            relationships=document.relationships,
            claims=document.claims,
        )
        
        # Phase 1: Embedding Entities
        if embedder and document.entities:
            typer.echo(f"    🧠 Generating embeddings for {len(document.entities)} entities...")
            entity_names = [e[0] for e in document.entities]
            entity_vectors = embed_texts(embedder, entity_names)
            
            # We need the IDs of the newly created entities to store embeddings
            # The repository doesn't return them, so we'll fetch them by name for this doc
            db_entities = repository.list_entities(document.document_id)
            entity_embeddings = []
            for name, vec in zip(entity_names, entity_vectors):
                # Find the entity by name in the fetched list
                entity_row = next((e for e in db_entities if e["name"] == name), None)
                if entity_row:
                    entity_embeddings.append((entity_row["id"], encode_vector(vec)))
            
            if entity_embeddings:
                repository.store_entity_embeddings(entity_embeddings)

            # Also store entity embeddings in ChromaDB
            if vectorstore is not None and DocsVectorStore is not None:
                entity_ids_for_chroma = []
                entity_vecs_for_chroma = []
                entity_names_for_chroma = []
                for name, vec in zip(entity_names, entity_vectors):
                    entity_row = next((e for e in db_entities if e["name"] == name), None)
                    if entity_row:
                        entity_ids_for_chroma.append(entity_row["id"])
                        entity_vecs_for_chroma.append(vec)
                        entity_names_for_chroma.append(name)
                if entity_ids_for_chroma:
                    vectorstore.upsert_entity_embeddings(
                        entity_ids=entity_ids_for_chroma,
                        embeddings=entity_vecs_for_chroma,
                        names=entity_names_for_chroma,
                    )

    communities = build_communities(analyzed)
    
    # Phase 2: Hierarchical Community Summarization
    if communities and resolved_settings.chat.provider.lower() != "none":
        typer.echo(f"  🏘️  Summarizing {len(communities)} semantic communities...")
        for i, community in enumerate(communities):
            typer.echo(f"    ⏳ Summarizing community {i+1}/{len(communities)}: {community.label}...", err=True)
            
            # Collect entities and relationships for context
            # In a real GraphRAG, we would pull all entity descriptions in this community
            # For parity, we'll use the top terms and related docs
            entity_descriptions = []
            rel_descriptions = []
            
            for doc_id in community.document_ids:
                doc = next((d for d in analyzed if d.document_id == doc_id), None)
                if doc:
                    entity_descriptions.extend([f"{name} (Score: {score})" for name, score in doc.entities[:10]])
                    rel_descriptions.extend([f"{s} -> {t} (Weight: {w})" for s, t, w in doc.relationships[:10]])
            
            title, summary = summarize_community(
                resolved_settings, 
                entity_descriptions, 
                rel_descriptions
            )
            if title:
                community.label = title
            if summary:
                community.summary = summary

    repository.replace_communities(communities)

    # Store community embeddings in ChromaDB
    if vectorstore is not None and DocsVectorStore is not None and embedder and communities:
        typer.echo(f"    🧠 Generating community embeddings for ChromaDB...")
        community_summaries = [c.summary for c in communities if c.summary]
        if community_summaries:
            community_vecs = embed_texts(embedder, community_summaries)
            # Fetch community IDs from DB
            db_communities = repository.list_communities()
            community_ids = [c["id"] for c in db_communities[:len(community_vecs)]]
            if community_ids:
                vectorstore.upsert_community_embeddings(
                    community_ids=community_ids,
                    embeddings=community_vecs[:len(community_ids)],
                    summaries=community_summaries[:len(community_ids)],
                )

    return {
        "documents": len(analyzed),
        "communities": len(communities),
        "provider_successes": provider_successes,
        "provider_fallbacks": provider_fallbacks,
        "provider_deterministic": provider_deterministic,
    }
