from typing import Dict, Any, Optional, List
from dataclasses import dataclass

from .gsz3_decompressor import GSZ3Decompressor, GSZ3DecompressionError
from .segment_runtime import SegmentRuntime, Segment, SegmentRuntimeError
from .execution_tracer import ExecutionTracer
from .profiler import SegmentProfiler


@dataclass
class CompressionManifest:
    source_file: str
    codex_lineage: Dict[str, Any]
    compression_format: str = "gsz3"


class CompressedEngineError(Exception):
    pass


class CompressedEngine:
    def __init__(self, verbose: bool = False, profile: bool = False):
        self.verbose = verbose
        self.profile = profile
        self.tracer = ExecutionTracer(verbose=verbose)
        self.profiler = SegmentProfiler() if profile else None

    def simulate(self, compressed_payload: bytes, manifest: CompressionManifest) -> Dict[str, Any]:
        try:
            decompressed = GSZ3Decompressor.decompress(compressed_payload)
            segments = self._build_segments(decompressed, manifest)
            return {
                "status": "simulated",
                "segment_count": len(segments),
                "total_bytes": sum(len(s.code.encode('utf-8')) for s in segments)
            }
        except GSZ3DecompressionError as e:
            raise CompressedEngineError(f"Decompression failed: {e}")

    def execute(self, compressed_payload: bytes, manifest: CompressionManifest, debug: bool = False) -> Dict[str, Any]:
        try:
            decompressed = GSZ3Decompressor.decompress(compressed_payload)
            segments = self._build_segments(decompressed, manifest)

            for segment in segments:
                if self.profile:
                    self.profiler.start(segment.segment_id)

                with self.tracer.trace_segment(
                    segment.segment_id,
                    0,
                    len(segment.code.encode('utf-8'))
                ):
                    try:
                        SegmentRuntime.execute_segments([segment], debug=debug)
                    except SegmentRuntimeError as e:
                        if debug:
                            raise
                        if self.verbose:
                            print(f"[ERROR] {segment.segment_id}: {e}")

                if self.profile:
                    self.profiler.stop(segment.segment_id)

            result = {
                "status": "executed",
                "segment_count": len(segments),
                "traces": [
                    {
                        "segment_id": t.segment_id,
                        "duration": t.duration,
                        "exception": t.exception
                    }
                    for t in self.tracer.get_traces()
                ]
            }

            if self.profile:
                result["profiles"] = {
                    k: {
                        "elapsed": v.elapsed,
                        "call_count": v.call_count
                    }
                    for k, v in self.profiler.get_all_profiles().items()
                }

            return result

        except GSZ3DecompressionError as e:
            raise CompressedEngineError(f"Decompression failed: {e}")
        except Exception as e:
            raise CompressedEngineError(f"Execution failed: {e}")

    def _build_segments(self, decompressed: str, manifest: CompressionManifest) -> List[Segment]:
        lineage = manifest.codex_lineage
        segments_meta = lineage.get("segments", [])

        if not segments_meta:
            return [Segment(
                segment_id="seg_0",
                code=decompressed,
                namespace={}
            )]

        segments = []
        lines = decompressed.split('\n')

        for i, seg_meta in enumerate(segments_meta):
            segment_id = seg_meta.get("id", f"seg_{i}")
            start = seg_meta.get("start", 0)
            end = seg_meta.get("end", len(lines))
            code = '\n'.join(lines[start:end])

            segments.append(Segment(
                segment_id=segment_id,
                code=code,
                namespace={}
            ))

        return segments