"""Segment merge engine — assembles shard transcripts into a final result."""

from __future__ import annotations

import json

import structlog

from hive_compute.coordinator.billing import create_billing_item
from hive_compute.coordinator.database import get_db_ctx, record_event
from hive_compute.coordinator.metrics import (
    increment_job_completed,
    increment_job_failed,
    observe_job_processing_duration,
    observe_job_turnaround_duration,
)
from hive_compute.coordinator.payouts import create_payout_item
from hive_compute.coordinator.webhooks import fire_webhook

logger = structlog.get_logger()


def merge_job_shards(job_id: str) -> dict | None:
    """Merge all completed shard transcripts into the final job result.

    Called when all shards for a job have reached a terminal state.
    Returns the merged result or None if merge isn't possible.
    """
    merged_result = None
    job_failed = False
    completed_shards = []

    with get_db_ctx() as db:
        # Exclude canary shards (canary_id IS NOT NULL) from merge — they're
        # validation-only and shouldn't affect the final transcript or shard counts.
        shards = db.execute(
            "SELECT * FROM shards WHERE job_id = ? AND canary_id IS NULL ORDER BY shard_index ASC",
            (job_id,),
        ).fetchall()

        if not shards:
            return None

        completed = [s for s in shards if s["status"] == "completed"]
        failed = [s for s in shards if s["status"] == "failed"]
        in_progress = [s for s in shards if s["status"] in ("pending", "claimed", "processing", "fallback")]

        # If shards are still in progress (including fallback), don't merge yet
        if in_progress:
            return None

        # If any shard failed permanently, the whole job fails
        if failed:
            error_details = "; ".join(
                f"shard {s['shard_index']}: {s['error_message'] or 'unknown'}"
                for s in failed
            )
            failed_job = db.execute(
                """
                UPDATE jobs SET status = 'failed', error_message = ?, completed_at = datetime('now')
                WHERE job_id = ? AND status = 'pending'
                RETURNING created_at, completed_at
                """,
                (f"Failed shards: {error_details}", job_id),
            ).fetchone()
            if not failed_job:
                return None  # Already terminal
            record_event(
                db, job_id, "merge_failed", to_status="failed",
                actor="merger", detail=f"{len(failed)}/{len(shards)} shards failed",
            )
            increment_job_failed()
            observe_job_turnaround_duration(failed_job["created_at"], failed_job["completed_at"])
            job_failed = True
        else:
            # All shards completed — merge transcripts
            segments = []
            total_duration = 0.0

            for shard in completed:
                result = json.loads(shard["result_data"]) if shard["result_data"] else {}
                transcript = result.get("transcript", "")

                segments.append({
                    "shard_index": shard["shard_index"],
                    "start_seconds": shard["start_seconds"],
                    "end_seconds": shard["end_seconds"],
                    "overlap_seconds": shard["overlap_seconds"],
                    "transcript": transcript,
                    "word_count": result.get("word_count") or (len(transcript.split()) if transcript else 0),
                })

                if result.get("duration_seconds"):
                    total_duration += result["duration_seconds"]

            merged_transcript = _merge_segments(segments)
            first_result = json.loads(completed[0]["result_data"]) if completed[0]["result_data"] else {}

            merged_result = {
                "transcript": merged_transcript,
                "model": first_result.get("model", "unknown"),
                "backend": first_result.get("backend", "unknown"),
                "language": first_result.get("language", "en"),
                "word_count": len(merged_transcript.split()) if merged_transcript else 0,
                "duration_seconds": total_duration or None,
                "shard_count": len(completed),
                "segments": segments,
            }

            result_json = json.dumps(merged_result)

            # Idempotent: only update if job is still pending (prevents duplicate merges)
            merged_job = db.execute(
                """
                UPDATE jobs SET status = 'completed', result_data = ?, completed_at = datetime('now')
                WHERE job_id = ? AND status = 'pending'
                RETURNING created_at, completed_at
                """,
                (result_json, job_id),
            ).fetchone()
            if not merged_job:
                # Already completed by another merge — skip
                return None

            record_event(
                db, job_id, "merged", from_status="pending", to_status="completed",
                actor="merger", detail=f"{len(completed)} shards merged, {len(merged_transcript.split())} words",
            )
            increment_job_completed()
            observe_job_turnaround_duration(merged_job["created_at"], merged_job["completed_at"])
            observe_job_processing_duration(merged_job["created_at"], merged_job["completed_at"])
            completed_shards = [dict(s) for s in completed]

    # Webhook calls OUTSIDE the db transaction to avoid SQLite deadlock
    if job_failed:
        logger.warning("Job merge failed", job_id=job_id)
        fire_webhook(job_id)
        return None

    if merged_result:
        logger.info("Job merged", job_id=job_id, shards=merged_result["shard_count"])

        # Create billing and payout items for the completed job
        with get_db_ctx() as db:
            job = db.execute(
                "SELECT partner_id, audio_duration FROM jobs WHERE job_id = ?",
                (job_id,),
            ).fetchone()

        if job and job["partner_id"] and job["audio_duration"]:
            create_billing_item(job_id, job["partner_id"], job["audio_duration"])

        # Pay workers — aggregate duration per worker to avoid sub-cent rounding loss.
        # One payout per worker per job, rounded once at job granularity.
        worker_durations: dict[str, float] = {}
        for shard in completed_shards:
            worker = shard["claimed_by"]
            if worker and worker != "cloud-fallback":
                duration = shard["end_seconds"] - shard["start_seconds"]
                worker_durations[worker] = worker_durations.get(worker, 0.0) + duration
        for worker, total_duration in worker_durations.items():
            create_payout_item(worker, job_id, total_duration)

        fire_webhook(job_id)

    return merged_result


def _merge_segments(segments: list[dict]) -> str:
    """Merge segment transcripts, deduplicating overlap regions.

    For overlapping segments, finds the longest common suffix/prefix
    between adjacent segments and removes the duplicate.
    """
    if not segments:
        return ""

    if len(segments) == 1:
        return segments[0]["transcript"].strip()

    parts = []
    for i, seg in enumerate(segments):
        text = seg["transcript"].strip()
        if not text:
            continue

        if i == 0:
            parts.append(text)
            continue

        # If there's overlap, try to deduplicate the boundary
        if seg["overlap_seconds"] > 0 and parts:
            prev_text = parts[-1]
            deduped = _deduplicate_boundary(prev_text, text)
            if deduped is not None:
                parts[-1] = deduped
            else:
                parts.append(text)
        else:
            parts.append(text)

    return " ".join(parts)


def _deduplicate_boundary(prev_text: str, next_text: str) -> str | None:
    """Try to find and remove overlapping text at the boundary.

    Returns combined text with overlap removed, or None if no overlap found.
    """
    prev_words = prev_text.split()
    next_words = next_text.split()

    if not prev_words or not next_words:
        return None

    # Look for matching sequences at the boundary (check up to 20 words)
    max_check = min(20, len(prev_words), len(next_words))

    best_overlap = 0
    for overlap_len in range(1, max_check + 1):
        prev_tail = " ".join(prev_words[-overlap_len:]).lower()
        next_head = " ".join(next_words[:overlap_len]).lower()

        if prev_tail == next_head:
            best_overlap = overlap_len

    if best_overlap > 0:
        # Remove the overlapping prefix from next_text
        combined = prev_text + " " + " ".join(next_words[best_overlap:])
        return combined.strip()

    return None
