import logging

from google.protobuf.duration_pb2 import Duration
from google.protobuf.timestamp_pb2 import Timestamp

from mlflow.entities import Assessment, Span, Trace, TraceData, TraceInfo
from mlflow.entities.trace_info_v2 import _truncate_request_metadata, _truncate_tags
from mlflow.entities.trace_location import (
    InferenceTableLocation,
    MlflowExperimentLocation,
    TraceLocation,
    TraceLocationType,
    UCSchemaLocation,
    UnityCatalog,
)
from mlflow.protos import assessments_pb2
from mlflow.protos import databricks_tracing_pb2 as pb
from mlflow.tracing.utils import (
    construct_trace_id_v4,
    parse_trace_id_v4,
)

_logger = logging.getLogger(__name__)


def parse_uc_location(location: str) -> tuple[str, str, str | None]:
    parts = location.split(".")
    if len(parts) == 2:
        return parts[0], parts[1], None
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    raise ValueError(
        f"Invalid UC location: {location}. Expected format: <catalog>.<schema>[.<table_prefix>]."
    )


def uc_location_to_str(catalog: str, schema: str, table_prefix: str | None = None) -> str:
    return f"{catalog}.{schema}.{table_prefix}" if table_prefix else f"{catalog}.{schema}"


def uc_schema_location_to_proto(
    uc_schema_location: UCSchemaLocation,
) -> pb.UCSchemaLocation:
    return pb.UCSchemaLocation(
        catalog_name=uc_schema_location.catalog_name,
        schema_name=uc_schema_location.schema_name,
        otel_spans_table_name=uc_schema_location._otel_spans_table_name,
        otel_logs_table_name=uc_schema_location._otel_logs_table_name,
    )


def uc_schema_location_from_proto(
    proto: pb.UCSchemaLocation,
) -> UCSchemaLocation:
    location = UCSchemaLocation(
        catalog_name=proto.catalog_name,
        schema_name=proto.schema_name,
    )

    if proto.HasField("otel_spans_table_name"):
        location._otel_spans_table_name = proto.otel_spans_table_name
    if proto.HasField("otel_logs_table_name"):
        location._otel_logs_table_name = proto.otel_logs_table_name
    return location


def uc_table_prefix_location_to_proto(
    location: UnityCatalog,
) -> pb.UcTablePrefixLocation:
    proto = pb.UcTablePrefixLocation(
        catalog_name=location.catalog_name,
        schema_name=location.schema_name,
        table_prefix=location.table_prefix,
    )
    if location._otel_spans_table_name:
        proto.spans_table_name = location._otel_spans_table_name
    if location._otel_logs_table_name:
        proto.logs_table_name = location._otel_logs_table_name
    if location._annotations_table_name:
        proto.annotations_table_name = location._annotations_table_name
    return proto


def uc_table_prefix_location_from_proto(proto: pb.UcTablePrefixLocation) -> UnityCatalog:
    location = UnityCatalog(
        catalog_name=proto.catalog_name,
        schema_name=proto.schema_name,
        table_prefix=proto.table_prefix,
    )
    if proto.HasField("spans_table_name"):
        location._otel_spans_table_name = proto.spans_table_name
    if proto.HasField("logs_table_name"):
        location._otel_logs_table_name = proto.logs_table_name
    if proto.HasField("annotations_table_name"):
        location._annotations_table_name = proto.annotations_table_name
    return location


def inference_table_location_to_proto(
    inference_table_location: InferenceTableLocation,
) -> pb.InferenceTableLocation:
    return pb.InferenceTableLocation(full_table_name=inference_table_location.full_table_name)


def mlflow_experiment_location_to_proto(
    mlflow_experiment_location: MlflowExperimentLocation,
) -> pb.MlflowExperimentLocation:
    return pb.MlflowExperimentLocation(experiment_id=mlflow_experiment_location.experiment_id)


def trace_location_to_proto(trace_location: TraceLocation) -> pb.TraceLocation:
    if trace_location.type == TraceLocationType.UC_SCHEMA:
        return pb.TraceLocation(
            type=pb.TraceLocation.TraceLocationType.UC_SCHEMA,
            uc_schema=uc_schema_location_to_proto(trace_location.uc_schema),
        )
    elif trace_location.type == TraceLocationType.UC_TABLE_PREFIX:
        return pb.TraceLocation(
            type=pb.TraceLocation.TraceLocationType.UC_TABLE_PREFIX,
            uc_table_prefix=uc_table_prefix_location_to_proto(trace_location.uc_table_prefix),
        )
    elif trace_location.type == TraceLocationType.MLFLOW_EXPERIMENT:
        return pb.TraceLocation(
            type=pb.TraceLocation.TraceLocationType.MLFLOW_EXPERIMENT,
            mlflow_experiment=mlflow_experiment_location_to_proto(trace_location.mlflow_experiment),
        )
    elif trace_location.type == TraceLocationType.INFERENCE_TABLE:
        return pb.TraceLocation(
            type=pb.TraceLocation.TraceLocationType.INFERENCE_TABLE,
            inference_table=inference_table_location_to_proto(trace_location.inference_table),
        )
    else:
        raise ValueError(f"Unsupported trace location type: {trace_location.type}")


def trace_location_type_from_proto(proto: pb.TraceLocation.TraceLocationType) -> TraceLocationType:
    return TraceLocationType(pb.TraceLocation.TraceLocationType.Name(proto))


def trace_location_from_proto(proto: pb.TraceLocation) -> TraceLocation:
    type_ = trace_location_type_from_proto(proto.type)

    if proto.WhichOneof("identifier") == "uc_schema":
        return TraceLocation(
            type=type_,
            uc_schema=uc_schema_location_from_proto(proto.uc_schema),
        )
    elif proto.WhichOneof("identifier") == "uc_table_prefix":
        return TraceLocation(
            type=type_,
            uc_table_prefix=uc_table_prefix_location_from_proto(proto.uc_table_prefix),
        )
    elif proto.WhichOneof("identifier") == "mlflow_experiment":
        return TraceLocation(
            type=type_,
            mlflow_experiment=MlflowExperimentLocation.from_proto(proto.mlflow_experiment),
        )
    elif proto.WhichOneof("identifier") == "inference_table":
        return TraceLocation(
            type=type_,
            inference_table=InferenceTableLocation.from_proto(proto.inference_table),
        )
    else:
        return TraceLocation(TraceLocationType.TRACE_LOCATION_TYPE_UNSPECIFIED)


def trace_info_to_v4_proto(trace_info: TraceInfo) -> pb.TraceInfo:
    request_time = Timestamp()
    request_time.FromMilliseconds(trace_info.request_time)
    execution_duration = None
    if trace_info.execution_duration is not None:
        execution_duration = Duration()
        execution_duration.FromMilliseconds(trace_info.execution_duration)

    if trace_info.trace_location.uc_schema or trace_info.trace_location.uc_table_prefix:
        _, trace_id = parse_trace_id_v4(trace_info.trace_id)
    else:
        trace_id = trace_info.trace_id

    return pb.TraceInfo(
        trace_id=trace_id,
        client_request_id=trace_info.client_request_id,
        trace_location=trace_location_to_proto(trace_info.trace_location),
        request_preview=trace_info.request_preview,
        response_preview=trace_info.response_preview,
        request_time=request_time,
        execution_duration=execution_duration,
        state=pb.TraceInfo.State.Value(trace_info.state),
        trace_metadata=_truncate_request_metadata(trace_info.trace_metadata),
        tags=_truncate_tags(trace_info.tags),
        assessments=[assessment_to_proto(a) for a in trace_info.assessments],
    )


def trace_to_proto(trace: Trace) -> pb.Trace:
    return pb.Trace(
        trace_info=trace.info.to_proto(),
        spans=[span.to_otel_proto() for span in trace.data.spans],
    )


def trace_from_proto(proto: pb.Trace, location_id: str) -> Trace:
    return Trace(
        info=TraceInfo.from_proto(proto.trace_info),
        data=TraceData(spans=[Span.from_otel_proto(span, location_id) for span in proto.spans]),
    )


def assessment_to_proto(assessment: Assessment) -> pb.Assessment:
    assessment_proto = pb.Assessment()
    assessment_proto.assessment_name = assessment.name
    location, trace_id = parse_trace_id_v4(assessment.trace_id)
    if location:
        catalog, schema, table_prefix = parse_uc_location(location)
        if table_prefix:
            uc_table_prefix = pb.UcTablePrefixLocation(
                catalog_name=catalog,
                schema_name=schema,
                table_prefix=table_prefix,
            )
            assessment_proto.trace_location.CopyFrom(
                pb.TraceLocation(
                    type=pb.TraceLocation.TraceLocationType.UC_TABLE_PREFIX,
                    uc_table_prefix=uc_table_prefix,
                )
            )
        else:
            uc_schema = pb.UCSchemaLocation(catalog_name=catalog, schema_name=schema)
            assessment_proto.trace_location.CopyFrom(
                pb.TraceLocation(
                    type=pb.TraceLocation.TraceLocationType.UC_SCHEMA,
                    uc_schema=uc_schema,
                )
            )
    assessment_proto.trace_id = trace_id

    assessment_proto.source.CopyFrom(assessment.source.to_proto())

    # Convert time in milliseconds to protobuf Timestamp
    assessment_proto.create_time.FromMilliseconds(assessment.create_time_ms)
    assessment_proto.last_update_time.FromMilliseconds(assessment.last_update_time_ms)

    if assessment.span_id is not None:
        assessment_proto.span_id = assessment.span_id
    if assessment.rationale is not None:
        assessment_proto.rationale = assessment.rationale
    if assessment.assessment_id is not None:
        assessment_proto.assessment_id = assessment.assessment_id

    if assessment.expectation is not None:
        assessment_proto.expectation.CopyFrom(assessment.expectation.to_proto())
    elif assessment.feedback is not None:
        assessment_proto.feedback.CopyFrom(assessment.feedback.to_proto())

    if assessment.metadata:
        for key, value in assessment.metadata.items():
            assessment_proto.metadata[key] = str(value)
    if assessment.overrides:
        assessment_proto.overrides = assessment.overrides
    if assessment.valid is not None:
        assessment_proto.valid = assessment.valid

    return assessment_proto


def get_trace_id_from_assessment_proto(proto: pb.Assessment | assessments_pb2.Assessment) -> str:
    if "trace_location" in proto.DESCRIPTOR.fields_by_name and proto.trace_location.HasField(
        "uc_schema"
    ):
        trace_location = proto.trace_location
        return construct_trace_id_v4(
            uc_location_to_str(
                trace_location.uc_schema.catalog_name,
                trace_location.uc_schema.schema_name,
            ),
            proto.trace_id,
        )
    elif "trace_location" in proto.DESCRIPTOR.fields_by_name and proto.trace_location.HasField(
        "uc_table_prefix"
    ):
        trace_location = proto.trace_location
        return construct_trace_id_v4(
            uc_location_to_str(
                trace_location.uc_table_prefix.catalog_name,
                trace_location.uc_table_prefix.schema_name,
                trace_location.uc_table_prefix.table_prefix,
            ),
            proto.trace_id,
        )
    else:
        return proto.trace_id