"""
Serialization and deserialization for Cyvest investigations.

Provides JSON export/import and Markdown generation for LLM consumption.
"""

from __future__ import annotations

import json
from decimal import Decimal
from pathlib import Path
from typing import TYPE_CHECKING, Any

from cyvest.levels import Level
from cyvest.model import Check, Container, Enrichment, Observable, Relationship, ThreatIntel
from cyvest.model_schema import InvestigationSchema, StatsChecksSchema
from cyvest.score import ScoreMode

if TYPE_CHECKING:
    from cyvest.cyvest import Cyvest


def serialize_investigation(cv: Cyvest) -> InvestigationSchema:
    """
    Serialize a complete investigation to an InvestigationSchema.

    Uses InvestigationSchema for validation and automatic serialization via
    Pydantic's field_serializer decorators.

    Args:
        cv: Cyvest investigation to serialize

    Returns:
        InvestigationSchema instance (use .model_dump() for dict)
    """

    # Helper to resolve proxies
    def resolve(obj: Any) -> Any:
        return obj._resolve() if hasattr(obj, "_resolve") else obj

    # Resolve all observables, checks, threat intels, enrichments, containers
    observables = {key: resolve(obs) for key, obs in cv.get_all_observables().items()}
    threat_intels = {key: resolve(ti) for key, ti in cv.get_all_threat_intels().items()}
    enrichments = {key: resolve(enr) for key, enr in cv.get_all_enrichments().items()}
    containers = {key: resolve(ctr) for key, ctr in cv.get_all_containers().items()}

    # Build checks organized by scope (resolve proxies)
    checks_by_scope: dict[str, list[Check]] = {}
    for check in cv.get_all_checks().values():
        check = resolve(check)
        if check.scope not in checks_by_scope:
            checks_by_scope[check.scope] = []
        checks_by_scope[check.scope].append(check)

    # Build checks organized by level
    checks_by_level: dict[str, list[str]] = {}
    for check in cv.get_all_checks().values():
        check = resolve(check)
        if check.level.name not in checks_by_level:
            checks_by_level[check.level.name] = []
        checks_by_level[check.level.name].append(check.key)

    # Get root type
    root = cv.observable_get_root()
    root_type_value = None
    if root:
        root = resolve(root)
        root_type_value = root.obs_type.value if hasattr(root.obs_type, "value") else str(root.obs_type)

    # Build and validate using Pydantic model
    investigation = InvestigationSchema(
        score=float(cv.get_global_score()),
        level=cv.get_global_level(),
        whitelisted=cv.investigation_is_whitelisted(),
        whitelists=list(cv.investigation_get_whitelists()),
        observables=observables,
        checks=checks_by_scope,
        checks_by_level=checks_by_level,
        threat_intels=threat_intels,
        enrichments=enrichments,
        containers=containers,
        stats=cv.get_statistics(),
        stats_checks=StatsChecksSchema(
            checks=len(cv.get_all_checks()),
            applied=sum(1 for c in cv.get_all_checks().values() if c.level != Level.NONE),
        ),
        data_extraction={
            "root_type": root_type_value,
            "score_mode": cv._investigation._score_engine._score_mode.value,
        },
    )

    return investigation


def save_investigation_json(cv: Cyvest, filepath: str | Path) -> None:
    """
    Save an investigation to a JSON file.

    Args:
        cv: Cyvest investigation to save
        filepath: Path to save the JSON file
    """
    data = serialize_investigation(cv)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(data.model_dump_json(indent=2, by_alias=True))


def generate_markdown_report(
    cv: Cyvest,
    include_containers: bool = False,
    include_enrichments: bool = False,
    include_observables: bool = True,
) -> str:
    """
    Generate a Markdown report of the investigation for LLM consumption.

    Args:
        cv: Cyvest investigation
        include_containers: Include containers section in the report (default: False)
        include_enrichments: Include enrichments section in the report (default: False)
        include_observables: Include observables section in the report (default: True)

    Returns:
        Markdown formatted report
    """
    lines = []

    # Header
    lines.append("# Cybersecurity Investigation Report")
    lines.append("")
    lines.append(f"**Global Score:** {cv.get_global_score()}")
    lines.append(f"**Global Level:** {cv.get_global_level().name}")
    whitelists = cv.investigation_get_whitelists()
    whitelist_status = "Yes" if whitelists else "No"
    lines.append(f"**Whitelisted Investigation:** {whitelist_status}")
    if whitelists:
        lines.append(f"**Whitelist Entries:** {len(whitelists)}")
    lines.append("")

    # Statistics
    lines.append("## Statistics")
    lines.append("")
    stats = cv.get_statistics()
    lines.append(f"- **Total Observables:** {stats.total_observables}")
    lines.append(f"- **Internal Observables:** {stats.internal_observables}")
    lines.append(f"- **External Observables:** {stats.external_observables}")
    lines.append(f"- **Whitelisted Observables:** {stats.whitelisted_observables}")
    lines.append(f"- **Total Checks:** {stats.total_checks}")
    lines.append(f"- **Applied Checks:** {stats.applied_checks}")
    lines.append(f"- **Total Threat Intel:** {stats.total_threat_intel}")
    lines.append("")

    # Whitelists
    if whitelists:
        lines.append("## Whitelists")
        lines.append("")
        for entry in whitelists:
            lines.append(f"- **{entry.identifier}** - {entry.name}")
            if entry.justification:
                lines.append(f"  - Justification: {entry.justification}")
        lines.append("")

    # Checks by Scope
    lines.append("## Checks by Scope")
    lines.append("")
    for scope, _count in cv.get_statistics().checks_by_scope.items():
        lines.append(f"### {scope}")
        lines.append("")
        for check in cv.get_all_checks().values():
            if check.scope == scope and check.level != Level.NONE:
                lines.append(f"- **{check.check_id}**: Score: {check.score}, Level: {check.level.name}")
                lines.append(f"  - Description: {check.description}")
                if check.comment:
                    lines.append(f"  - Comment: {check.comment}")
        lines.append("")

    # Observables
    if include_observables and cv.get_all_observables():
        lines.append("## Observables")
        lines.append("")
        for obs in cv.get_all_observables().values():
            lines.append(f"### {obs.obs_type}: {obs.value}")
            lines.append(f"- **Key:** {obs.key}")
            lines.append(f"- **Score:** {obs.score}")
            lines.append(f"- **Level:** {obs.level.name}")
            lines.append(f"- **Internal:** {obs.internal}")
            lines.append(f"- **Whitelisted:** {obs.whitelisted}")
            if obs.comment:
                lines.append(f"- **Comment:** {obs.comment}")
            if obs.relationships:
                lines.append("- **Relationships:**")
                for rel in obs.relationships:
                    direction_symbol = {
                        "outbound": "→",
                        "inbound": "←",
                        "bidirectional": "↔",
                    }.get(rel.direction if isinstance(rel.direction, str) else rel.direction.value, "→")
                    lines.append(f"  - {rel.relationship_type} {direction_symbol} {rel.target_key}")
            if obs.threat_intels:
                lines.append("- **Threat Intelligence:**")
                for ti in obs.threat_intels:
                    lines.append(f"  - {ti.source}: Score {ti.score}, Level {ti.level.name}")
                    if ti.comment:
                        lines.append(f"    - {ti.comment}")
            lines.append("")

    # Enrichments
    if include_enrichments and cv.get_all_enrichments():
        lines.append("## Enrichments")
        lines.append("")
        for enr in cv.get_all_enrichments().values():
            lines.append(f"### {enr.name}")
            if enr.context:
                lines.append(f"- **Context:** {enr.context}")
            lines.append(f"- **Data:** {json.dumps(enr.data, indent=2)}")
            lines.append("")

    # Containers
    if include_containers and cv.get_all_containers():
        lines.append("## Containers")
        lines.append("")
        for ctr in cv.get_all_containers().values():
            lines.append(f"### {ctr.path}")
            lines.append(f"- **Description:** {ctr.description}")
            lines.append(f"- **Aggregated Score:** {ctr.get_aggregated_score()}")
            lines.append(f"- **Aggregated Level:** {ctr.get_aggregated_level().name}")
            lines.append(f"- **Checks:** {len(ctr.checks)}")
            lines.append(f"- **Sub-containers:** {len(ctr.sub_containers)}")
            lines.append("")

    return "\n".join(lines)


def save_investigation_markdown(
    cv: Cyvest,
    filepath: str | Path,
    include_containers: bool = False,
    include_enrichments: bool = False,
    include_observables: bool = True,
) -> None:
    """
    Save an investigation as a Markdown report.

    Args:
        cv: Cyvest investigation to save
        filepath: Path to save the Markdown file
        include_containers: Include containers section in the report (default: False)
        include_enrichments: Include enrichments section in the report (default: False)
        include_observables: Include observables section in the report (default: True)
    """
    markdown = generate_markdown_report(cv, include_containers, include_enrichments, include_observables)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(markdown)


def load_investigation_json(filepath: str | Path) -> Cyvest:
    """
    Load an investigation from a JSON file into a Cyvest object.

    Args:
        filepath: Path to the JSON file

    Returns:
        Reconstructed Cyvest investigation
    """
    from cyvest.cyvest import Cyvest
    from cyvest.investigation import Investigation

    with open(filepath, encoding="utf-8") as handle:
        data = json.load(handle)

    data_payload = data.get("data")
    extraction = data.get("data_extraction", {})

    root_type_raw = extraction.get("root_type")
    root_type = "file"
    if root_type_raw:
        root_type = root_type_raw.value if hasattr(root_type_raw, "value") else str(root_type_raw)
    if root_type not in ("file", "artifact"):
        root_type = "file"

    score_mode_raw = extraction.get("score_mode")
    try:
        score_mode = ScoreMode(score_mode_raw) if score_mode_raw else ScoreMode.MAX
    except (TypeError, ValueError):
        score_mode = ScoreMode.MAX

    cv = Cyvest(data=data_payload, root_type=root_type, score_mode=score_mode)

    # Reset internal state to avoid default root pollution
    cv._investigation = Investigation(data_payload, root_type=root_type, score_mode=score_mode)

    # Load whitelists using Pydantic validation
    whitelists = data.get("whitelists") or []
    for whitelist_info in whitelists:
        try:
            # Use model_validate for Pydantic models, with fallback for backward compatibility
            identifier = str(whitelist_info.get("identifier", "")).strip()
            name = str(whitelist_info.get("name", "")).strip()
            if identifier and name:
                cv._investigation.add_whitelist(
                    identifier,
                    name,
                    whitelist_info.get("justification"),
                )
        except ValueError:
            continue

    # Backward compatibility for older exports
    if data.get("whitelisted") and not cv._investigation.is_whitelisted():
        cv._investigation.add_whitelist("default", "Whitelisted", data.get("whitelisted_reason"))

    # Observables - leverage Pydantic model_validate
    for obs_info in data.get("observables", {}).values():
        # Prepare data for Pydantic validation
        obs_data = {
            "obs_type": obs_info.get("type", "unknown"),
            "value": obs_info.get("value", ""),
            "internal": obs_info.get("internal", True),
            "whitelisted": obs_info.get("whitelisted", False),
            "comment": obs_info.get("comment", ""),
            "extra": obs_info.get("extra", {}),
            "score": Decimal(str(obs_info.get("score", 0))),
            "level": obs_info.get("level", "INFO"),
            "key": obs_info.get("key", ""),
            "relationships": [Relationship.model_validate(rel) for rel in obs_info.get("relationships", [])],
        }
        obs = Observable.model_validate(obs_data)
        # Restore private attribute
        obs._generated_by_checks = obs_info.get("generated_by_checks", [])
        cv._investigation.add_observable(obs)

    # Threat intel - leverage Pydantic model_validate
    for ti_info in data.get("threat_intels", {}).values():
        ti_data = {
            "source": ti_info.get("source", ""),
            "observable_key": ti_info.get("observable_key", ""),
            "comment": ti_info.get("comment", ""),
            "extra": ti_info.get("extra", {}),
            "score": Decimal(str(ti_info.get("score", 0))),
            "level": ti_info.get("level", "INFO"),
            "taxonomies": ti_info.get("taxonomies", []),
            "key": ti_info.get("key", ""),
        }
        ti = ThreatIntel.model_validate(ti_data)
        observable = cv._investigation.get_observable(ti.observable_key)
        if observable:
            cv._investigation.add_threat_intel(ti, observable)

    # Checks - leverage Pydantic model_validate
    for scope_checks in data.get("checks", {}).values():
        for check_info in scope_checks:
            check_data = {
                "check_id": check_info.get("check_id", ""),
                "scope": check_info.get("scope", ""),
                "description": check_info.get("description", ""),
                "comment": check_info.get("comment", ""),
                "extra": check_info.get("extra", {}),
                "score": Decimal(str(check_info.get("score", 0))),
                "level": check_info.get("level", "NONE"),
                "score_policy": check_info.get("score_policy", "auto"),
                "key": check_info.get("key", ""),
            }
            check = Check.model_validate(check_data)
            observable_keys = check_info.get("observables", [])
            for obs_key in observable_keys:
                observable = cv._investigation.get_observable(obs_key)
                if observable:
                    check.add_observable(observable)
            cv._investigation.add_check(check)

    # Enrichments - leverage Pydantic model_validate
    for enr_info in data.get("enrichments", {}).values():
        enr_data = {
            "name": enr_info.get("name", ""),
            "data": enr_info.get("data", {}),
            "context": enr_info.get("context", ""),
            "key": enr_info.get("key", ""),
        }
        enrichment = Enrichment.model_validate(enr_data)
        cv._investigation.add_enrichment(enrichment)

    # Containers
    def build_container(container_info: dict[str, Any]) -> Container:
        container_data = {
            "path": container_info.get("path", ""),
            "description": container_info.get("description", ""),
            "key": container_info.get("key", ""),
        }
        container = Container.model_validate(container_data)
        container = cv._investigation.add_container(container)

        for check_key in container_info.get("checks", []):
            check = cv._investigation.get_check(check_key)
            if check:
                container.add_check(check)

        for sub_info in container_info.get("sub_containers", {}).values():
            sub_container = build_container(sub_info)
            container.add_sub_container(sub_container)

        return container

    for container_info in data.get("containers", {}).values():
        build_container(container_info)

    # Note: Root observable is managed by Investigation, no need to set it here

    return cv
